# 从EXCEL中获取数据
import pandas as pd
data = pd.read_excel('InterestData.xlsx', sheetname='Data')
yc0 = pd.read_excel('InterestData.xlsx', sheetname='YC')
data.index=data['Tenor'].values
del data['Tenor']

# 计算远期利率的协方差矩阵，及其特征值和特征向量。
from numpy import linalg as LA
import random

data_value=data.diff().shift(-1).iloc[:-2,:].T
tenor=list(data_value.index)
tenor[0]=0    #由于这个后期要被当做short rate
covmatrix=cov(data_value)*252/10000
covmatrix_df=pd.DataFrame(covmatrix,index=tenor,columns=tenor)
# PCA
w,v=LA.eig(covmatrix)
#主成分我们选取前三大的为主成分
w=w[0:3]
v=v[:,0:3]

# 校准波动率
vol_matrix=[list(ones(shape(v[:,0]))*mean(sqrt(w[0])*v[:,0]))]
poly_matrix=[mean(sqrt(w[0])*v[:,0])]
for i in range(1,3):
    vol=sqrt(w)[i]*v[:,i]
    x_data=np.array(tenor)
    y_data=vol
    poly = np.polyfit(x_data,y_data,deg=3)
    y_fit_data=np.polyval(poly,x_data)
    vol_matrix.append(list(y_fit_data))
    poly_matrix.append(poly)
vol_matrix[2]=list(np.array(vol_matrix[2])*(-1))
vol_df=pd.DataFrame(vol_matrix,index=['Vol1','Vol2','Vol3'],columns=tenor)

# 校准drift项
def M_Calculator(Tau):
    if Tau==0:
        return 0

    dTau=0.01
    N=Tau/dTau

    Vol1=ones(N+1)*poly_matrix[0]
    M1=(0.5*Vol1[0]+sum(Vol1[1:-1])+0.5*Vol1[-1])*dTau
    M1=M1*Vol1[-1]

    x_data=arange(0,Tau+dTau,dTau)
    poly=poly_matrix[1]
    y_fit_data=np.polyval(poly,x_data)
    Vol2=y_fit_data
    M2=(0.5*Vol2[0]+sum(Vol2[1:-1])+0.5*Vol2[-1])*dTau
    M2=M2*Vol2[-1]

    poly=poly_matrix[2]
    y_fit_data=np.polyval(poly,x_data)
    Vol3=y_fit_data
    M3=(0.5*Vol3[0]+sum(Vol3[1:-1])+0.5*Vol3[-1])*dTau
    M3=M3*Vol3[-1]

    M=M1+M2+M3

    return M

# 计算不同期限的m项
m_list=[]
for t in tenor:
    m=M_Calculator(t)
    m_list.append(m)

# 与之前的vol合并为一个dataframe
m_df=pd.DataFrame(m_list,index=tenor,columns=['drift']).T
df=pd.concat([m_df,vol_df],axis=0)

#生产随机数，以及布局初始远期利率曲线。
dt=0.01
Nstep=1000
yc=pd.DataFrame(index=arange(Nstep+1),columns=tenor)
phi = np.random.normal(0,1,3*Nstep).reshape((Nstep,3))
yc.iloc[0,:]=yc0.values

for t in list(range(len(yc.index)))[:-1]:
    for tau in list(range(len(yc.columns))):
        m=df.iloc[0,tau]
        Vol1=df.iloc[1,tau]
        Vol2=df.iloc[2,tau]
        Vol3=df.iloc[3,tau]
        A=yc.iloc[t,tau]
        B=m*dt+(Vol1*phi[t,0]+Vol2*phi[t,1]+Vol3*phi[t,2])*sqrt(dt)        
        if tau==list(range(len(yc.columns)))[-1]:
            C=(yc.iloc[t,tau]-yc.iloc[t,tau-1])/(tenor[tau]-tenor[tau-1])*dt
        else:
            C=(yc.iloc[t,tau+1]-yc.iloc[t,tau])/(tenor[tau+1]-tenor[tau])*dt
        yc.iloc[t+1,tau]=A+B+C

yc

#封装为一个方程
def YieldCurve_Simulator(N,yc0,df,Nstep=1000,dt=0.01):
    yc_list=[]
    tenor=list(yc0.columns)

    for i in range(N):
    # Initialize and preallocate the yield curve matrix  
        yc=pd.DataFrame(index=arange(Nstep+1),columns=tenor)
        phi = np.random.normal(0,1,3*Nstep).reshape((Nstep,3))
        yc.iloc[0,:]=yc0.values   
        Row=list(range(len(yc.index)))[:-1]
        Column=list(range(len(yc.columns)))
    # Calculate the Yield Curve    
        for t in Row:
            for tau in Column:
                m=df.iloc[0,tau]
                Vol1=df.iloc[1,tau]
                Vol2=df.iloc[2,tau]
                Vol3=df.iloc[3,tau]
                A=yc.iloc[t,tau]
                B=m*dt+(Vol1*phi[t,0]+Vol2*phi[t,1]+Vol3*phi[t,2])*sqrt(dt)        
                if tau==list(range(len(yc.columns)))[-1]:
                    C=(yc.iloc[t,tau]-yc.iloc[t,tau-1])/(tenor[tau]-tenor[tau-1])*dt
                else:
                    C=(yc.iloc[t,tau+1]-yc.iloc[t,tau])/(tenor[tau+1]-tenor[tau])*dt
                yc.iloc[t+1,tau]=A+B+C
        yc_list.append(yc)
        
    return yc_list
N=1
YieldCurve_Simulator(N,yc0,df)