# 导入需要的程序包
import pandas as pd
import seaborn as sns
import random as rand
from datetime import datetime
SECURITY ='000001.XSHG'    #000001.XSHG  上证综指           #000300.XSHG 沪深300
AIP_FUND ='510210.XSHG'    #510210.XSHG #富国上证综指ETF    #159919.XSHE  #嘉实沪深300ETF 开始日期 2012-05-28
START_DATE='2005-04-09'    #取数的开始日期，上证综指不能早于'2005-01-05'，沪深300不能早2005-04-08
END_DATE='2016-03-31'      #取数、测算的结束日期
START_CAL_DATE='2012-06-01'    #测算起始日期
STD_AMOUNT =100            #基准定投金额

#GA参数
PSIZE=20    #种群规模
TSIZE=5    # 锦标赛成员规模
TERMINAL=10    #迭代代数
PC=0.95    #交叉比率
PM=0.2    #变异比率
ALPH=1    #适配值函数参数
BETA=10
# 获取指数历史数据
security_data=get_price(SECURITY,START_DATE,END_DATE,fields=['price','close'])  #跟踪的指数
aip_fund_data=get_price(AIP_FUND,START_CAL_DATE,END_DATE,fields=['price'])    #定投的基金
# 种群初始化
class Population(object):
    def __init__(self,chromosome,amount=0.0,quantity=0.0,nav=0.0,fitness=0.0):
        self.chromosome = chromosome
        self.amount = amount
        self.quantity=quantity
        self.nav=nav
        self.fitness=fitness
popu=[Population(chromosome=[random.randint(0,4),random.randint(0,3),random.randint(0,3)]) for i in range(PSIZE)]
#第1段表示均线类型，第2段表示偏离度，第3段表示级差。

#获取START_CAL_DATE在security_data中的index
cal_date=pd.DatetimeIndex([datetime.strptime(START_CAL_DATE,'%Y-%m-%d')])[0]
i=0
index=-1
while (i<len(security_data)):
    if(security_data.index[i]>=cal_date):
        index=i
        break
    else:
        i+=1
if(index==-1):
    print('开始测算日期录入错误')
    sys.exit()

# 求某index的num日均线值
def mav(security_data,index,num):
    if(index-num<0):
        print('现有数据不足以计算此均值，最早日期为: ',security_data.index[0],'超出天数为: ',num-index)
        sys.exit()
    j=index-1
    sumv=0.0
    while (j>=index-num and j>=0):
        sumv=sumv+security_data.iloc[j][1]    #temp_security.iloc[j][0] 表示price；temp_security.iloc[j][1] 表示close
        j=j-1 
    mav=sumv/num
    return mav

# 解码函数
def decode(p=0,size=1):    #p：起始个体index    size：解码数量
    temp_p=p
    #检查录入的均线、偏离度、级差是否正确
    if(popu[temp_p].chromosome[2]>2):
        print('编码第3位：',popu[2],' 超出级差取值范围')
        sys.exit()
    if(popu[temp_p].chromosome[1]>2):
        print('编码第2位：',popu[1],' 超出偏离度取值范围')
        sys.exit()
    if(popu[temp_p].chromosome[0]>4):
        print('编码第1位：',popu[0],' 超出均线种类取值范围')
        sys.exit()

    while(temp_p<p+size):
        fd_index=index
        #计算START_CAL_DATE日个体编码对应的均线值
        if(popu[temp_p].chromosome[0]==0):
            mavg=mav(security_data,fd_index,120)
            #print('使用120均线')
        elif(popu[temp_p].chromosome[0]==1):
            mavg=mav(security_data,fd_index,180)
            #print('使用180均线')
        elif(popu[temp_p].chromosome[0]==2):
            mavg=mav(security_data,fd_index,250)
            #print('使用250均线')            
        elif(popu[temp_p].chromosome[0]==3):
            mavg=mav(security_data,fd_index,500)
            #print('使用180均线')
        elif(popu[temp_p].chromosome[0]==4):
            mavg=mav(security_data,fd_index,2426)           #10年均线按2426个工作日算（基于2005-1-5到2015-1-4计算得到）
            #print('使用10年均线')
        else:
            sys.exit()
        aip_index=0
        popu[temp_p].amount=0.0
        popu[temp_p].quantity=0.0
        popu[temp_p].nav=0.0
        popu[temp_p].fitness=0.0
        while(fd_index<len(security_data)):
        #0：广发 (-n,-40),[-40,-30),[-30,-20),[-20,-10),[-10,-5),[-5,0),(0,15],(15,50],(50,100],(100,n)
            if(popu[temp_p].chromosome[1]==0):
                rod=(security_data.iloc[fd_index][1]-mavg)*100/mavg         #security_data.iloc[j][1] 表示close，用close计算偏离率
                if(rod<-40):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*6))
                elif(rod<-30):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*5))
                elif(rod<-20):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*4))
                elif(rod<-10):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*3))                    
                elif(rod<-5):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*2)) 
                elif(rod<0):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*1))
                elif(rod<=15):
                    in_amount=(STD_AMOUNT*(1-(popu[temp_p].chromosome[2]+1)*0.1*1))
                elif(rod<=50):
                    in_amount=(STD_AMOUNT*(1-(popu[temp_p].chromosome[2]+1)*0.1*2))                                      
                elif(rod<=100):
                    in_amount=(STD_AMOUNT*(1-(popu[temp_p].chromosome[2]+1)*0.1*3))    
                else:
                    in_amount=(STD_AMOUNT*(1-(popu[temp_p].chromosome[2]+1)*0.1*4))

                popu[temp_p].amount += in_amount
                popu[temp_p].quantity += in_amount/aip_fund_data.iloc[aip_index][0]
                popu[temp_p].nav = popu[temp_p].amount/popu[temp_p].quantity
                            
                '''print('日期: ',aip_fund_data.index[aip_index],' 指数收盘价: ',security_data.iloc[fd_index][1],' 指数均值: ',mavg,' 偏离率:',rod,\
                    ' 基金净值: ',aip_fund_data.iloc[aip_index][0],' 购买金额: ',in_amount,' 总金额: ',\
                     pop[temp_p].amount,' 总数量: ',pop[temp_p].quantity,' 成本净值: ',pop[temp_p].nav)'''
            #1: 富国 (-n,-45),[-45,-35),[-35,-25),[-25,-15),[-15,-5),[-5,5),[5,20),[20,35),[35,50),[50,n)
            elif(popu[temp_p].chromosome[1]==1):
                rod=(security_data.iloc[fd_index][1]-mavg)*100/mavg
                if(rod<-45):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*5))
                elif(rod<-35):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*4))
                elif(rod<-25):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*3))
                elif(rod<-15):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*2))
                elif(rod<-5):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*1))                    
                elif(rod<5):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1*0))      
                elif(rod<20):
                    in_amount=(STD_AMOUNT*(1-(popu[temp_p].chromosome[2]+1)*0.05*1))             
                elif(rod<35):
                    in_amount=(STD_AMOUNT*(1-(popu[temp_p].chromosome[2]+1)*0.05*2))                       
                elif(rod<50):
                    in_amount=(STD_AMOUNT*(1-(popu[temp_p].chromosome[2]+1)*0.05*3))                               
                else:
                    in_amount=(STD_AMOUNT*(1-(popu[temp_p].chromosome[2]+1)*0.05*4))
                popu[temp_p].amount += in_amount
                popu[temp_p].quantity += in_amount/aip_fund_data.iloc[aip_index][0]
                popu[temp_p].nav = popu[temp_p].amount/popu[temp_p].quantity
                '''print('日期: ',aip_fund_data.index[aip_index],' 指数收盘价: ',security_data.iloc[fd_index][1],' 指数均值: ',mavg,' 偏离率:',rod,\
                    ' 基金净值: ',aip_fund_data.iloc[aip_index][0],' 购买金额: ',in_amount,' 总金额: ',\
                     pop[temp_p].amount,' 总数量: ',pop[temp_p].quantity,' 成本净值: ',pop[temp_p].nav) '''                   
            #2: 博时 price/avg<=0.9,0.9<price/avg<=1.1,price/avg>=1.1
            elif(popu[temp_p].chromosome[1]==2):
                rod=security_data.iloc[fd_index][1]/mavg
                if(rod<=0.9):
                    in_amount=(STD_AMOUNT*(1+(popu[temp_p].chromosome[2]+1)*0.1))
                elif(rod<1.1):
                    in_amount=STD_AMOUNT
                else:
                    in_amount=(STD_AMOUNT*(1-(popu[temp_p].chromosome[2]+1)*0.1))
                popu[temp_p].amount += in_amount
                popu[temp_p].quantity += in_amount/aip_fund_data.iloc[aip_index][0]
                popu[temp_p].nav = popu[temp_p].amount/popu[temp_p].quantity
                '''print('日期: ',aip_fund_data.index[aip_index],' 指数收盘价: ',security_data.iloc[fd_index][1],' 指数均值: ',mavg,' 偏离率:',rod,\
                    ' 基金净值: ',aip_fund_data.iloc[aip_index][0],' 购买金额: ',in_amount,' 总金额: ',\
                     pop[temp_p].amount,' 总数量: ',pop[temp_p].quantity,' 成本净值: ',pop[temp_p].nav) '''
            
            if(fd_index==(len(security_data)-1)):
                #popu[temp_p].fitness=(aip_fund_data.iloc[aip_index][0]-popu[temp_p].nav)*pop[temp_p].quantity
                popu[temp_p].fitness=ALPH*exp(BETA*(aip_fund_data.iloc[aip_index][0]*popu[temp_p].quantity)/popu[temp_p].amount)
                #popu[temp_p].fitness=(aip_fund_data.iloc[aip_index][0]*popu[temp_p].quantity)/popu[temp_p].amount
            fd_index+=1
            aip_index+=1
        temp_p=temp_p+1

# 选择算子函数（使用轮盘赌的方式从父代染色体中选择）
#import pdb
def select(size,sumfit):    #size：当前的种群规模
    sumfit=sumfit
    #pdb.set_trace()
    temp_popu=[Population(chromosome=[]) for i in range(PSIZE)]
    p=0
    while(p < PSIZE):
        sv=random.random()*sumfit[size-1]
        #print(' sv: ',sv,' total_sumfit: ',sumfit[size-1])
        sid = 0
        while (sid<size):
            if(sumfit[sid] > sv):
                #print(' sumfit[sid]: ',sumfit[sid],' sid: ',sid)
                temp_popu[p]=popu[sid]
                break
            else:
                sid+=1
        p+=1
    #pdb.set_trace()
    q=0
    while(q<size):
        if(q<PSIZE):
            popu[q]=temp_popu[q]
        else:
            popu.pop()    #去掉多余的个体
        #print(pop[q].chromosome,' ',pop[q].fitness)
        q+=1

# 交叉算子函数（使用一段式）
#import pdb
def crossover(size):
    pop_size=size
    p=0
    while(p<PSIZE):     #运行PSIZE次，总交叉次数为PSIZE*PC
        if(random.random()<PC):
            #创建锦标赛群体
            #pdb.set_trace()
           
            crs_pop=[int(random.random()*PSIZE) for i in range(TSIZE)]
            
            #pdb.set_trace()
            max_v=-99999999999999999999999999999999999999999999999999999999999
            max1=-1
            for j in crs_pop:
                if(popu[j].fitness>max_v):
                    max1=j
            crs_pop.remove(max1)
            max_v=-99999999999999999999999999999999999999999999999999999999999
            max2=-1
            for j in crs_pop:
                if(popu[j].fitness>max_v):
                    max2=j
            
            #计算交叉位置，开始交叉
            crs_point=int(random.random()*(len(popu[0].chromosome)-1))
            #pdb.set_trace()
            
            temp_chr1=[]
            temp_chr2=[]
            ci=0
            while (ci<len(popu[0].chromosome)):
                if(ci<crs_point):
                    temp_chr1.append(popu[max1].chromosome[ci])
                    temp_chr2.append(popu[max2].chromosome[ci])
                else:
                    temp_chr1.append(popu[max2].chromosome[ci])
                    temp_chr2.append(popu[max1].chromosome[ci])
                ci+=1
            
            temp_pop1=Population(chromosome=temp_chr1)
            temp_pop2=Population(chromosome=temp_chr2)
            popu.append(temp_pop1)
            popu.append(temp_pop2)
            pop_size+=2
        p+=1
    return pop_size

# 变异算子函数
#import pdb
def mutation(size):
    pop_size=size
    i=0
    while (i<PSIZE):
        #pdb.set_trace()
        temp_chr=[]

        mutated=0
        j=0
        while(j<len(popu[i].chromosome)):
            random_p = random.random()
            if(random_p<PM and j == 0):
                temp_chr.append(random.randint(0,4))
                mutated=1
                j+=1
            elif(random_p<PM and j == 1):
                temp_chr.append(random.randint(0,3))
                mutated=1
                j+=1
            elif(random_p<PM and j == 2):
                temp_chr.append(random.randint(0,3))
                mutated=1
                j+=1
            else:
                temp_chr.append(popu[i].chromosome[j])
                j+=1
        i+=1
        if(mutated==1):
            temp_pop=Population(chromosome=temp_chr)
            popu.append(temp_pop)
            pop_size+=1
    return pop_size        

def GA_MAIN():
    accumulate_income=[]
    decode(0,PSIZE)
    
    print('Initail')
    max_ini_fit=popu[0].fitness
    max_ini_id=0
    

    for p0 in range(PSIZE):
    # 显示种群全部个体
    #    print('id',p0,'chromosome: ',popu[p0].chromosome,' amount: ',popu[p0].amount,' quantity: ',popu[p0].quantity,\
    #        ' nav: ',popu[p0].nav,' fitness: ',popu[p0].fitness)
    #显示种群最优个体
        if(popu[p0].fitness>max_ini_fit):
            max_ini_fit=popu[p0].fitness
            max_ini_id=p0
    best_fit=popu[max_ini_id].fitness
    best_chr=popu[max_ini_id].chromosome
    print('id',max_ini_id,'chromosome: ',popu[max_ini_id].chromosome,' amount: ',popu[max_ini_id].amount,' quantity: ',\
          popu[max_ini_id].quantity,' nav: ',popu[max_ini_id].nav,' fitness: ',popu[max_ini_id].fitness)
    
    accumulate_income.append(100*(aip_fund_data.iloc[len(aip_fund_data)-1][0]*popu[max_ini_id].quantity)/popu[max_ini_id].amount-1)
    
    for runtime in range(TERMINAL):
        pp=crossover(PSIZE)
        
    #显示交叉算子处理后的种群
    #    print('after crossover')
    #    for p1 in range(pp):
    #        print('id',p1,'chromosome: ',popu[p1].chromosome,' amount: ',popu[p1].amount,' quantity: ',popu[p1].quantity,\
    #              ' nav: ',popu[p1].nav,' fitness: ',popu[p1].fitness)

        pp=mutation(pp)
    #显示变异算子处理后的种群
    #    print('after mutation')
    #    for p2 in range(pp):
    #        print('id',p2,'chromosome: ',popu[p2].chromosome,' amount: ',popu[p2].amount,' quantity: ',popu[p2].quantity,\
    #              ' nav: ',popu[p2].nav,' fitness: ',popu[p2].fitness)

        decode(PSIZE,pp-PSIZE)
    #解码种群所有个体
    #    print('after decode')
    #    for p3 in range(pp):
    #        print('id',p3,'chromosome: ',popu[p3].chromosome,' amount: ',popu[p3].amount,' quantity: ',popu[p3].quantity,\
    #              ' nav: ',popu[p3].nav,' fitness: ',popu[p3].fitness)

        #计算个体的sumfit
        sumfit=[]
        sumi=0.0
        si=0
        while(si<pp):
            sumi+=popu[si].fitness
            sumfit.append(sumi)
            si+=1
            
        select(pp,sumfit)
        
        print('Iteration: ',runtime)
        max_ini_fit=popu[0].fitness
        max_ini_id=0

        for p5 in range(len(popu)):
        # 显示种群全部个体
        #    print('id',p5,'chromosome: ',popu[p5].chromosome,' amount: ',popu[p5].amount,' quantity: ',popu[p5].quantity,\
        #          ' nav: ',popu[p5].nav,' fitness: ',popu[p5].fitness)
        #显示种群最优个体         
            if(popu[p5].fitness>max_ini_fit):
                max_ini_fit=popu[p5].fitness
                max_ini_id=p5
        if(best_fit<popu[max_ini_id].fitness):
            best_fit=popu[max_ini_id].fitness
            best_chr=popu[max_ini_id].chromosome
        print('id',max_ini_id,'chromosome: ',popu[max_ini_id].chromosome,' amount: ',popu[max_ini_id].amount,' quantity: ',\
              popu[max_ini_id].quantity,' nav: ',popu[max_ini_id].nav,' fitness: ',popu[max_ini_id].fitness)
        accumulate_income.append(100*(aip_fund_data.iloc[len(aip_fund_data)-1][0]*popu[max_ini_id].quantity)/popu[max_ini_id].amount-1)
    print(' best_chr: ',best_chr,' best_fit: ',best_fit,' accumulate_income: ',accumulate_income)
    
    #设置横坐标和纵坐标的名称
    plt.xlabel('Iteration')
    plt.ylabel('Accumulate_income')
    # 图的标题
    plt.title('Accumulate income of each GA iteration')
    # 定义线条
    line1 = plt.plot(accumulate_income)
    plt.setp(line1, color='r', linewidth=2.0,label='GA')
    plt.legend()

GA_MAIN()