from __future__ import print_function, division
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
%matplotlib inline

from jqdata import *
import datetime
import matplotlib.pyplot as plt
import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
#参数设置
shortcode=None
opentrade_s=[]
opentrade_l=[]
closetrade_s=[]
closetrade_l=[]

#获取每个月的第n个交易日的日期
def tradedate(start,end,n=1):
    'start,end均为str类型，如“20190102”，n为当月的第几个交易日'
    #如果回测失败，可能token如果过期，需要自己在Tushare上重新申请，不然没数据，建议注册自己注册一个。
    ts.set_token('384e878fadc6ad14225108d7a471360ffab5913e30378be6a7871c35')
    tradedate=[]
    b=[]
    pro = ts.pro_api()
    #获取
    a=pro.trade_cal(exchange='', start_date=start, end_date=end)
    x=0
    for i in a.index:
        if len(tradedate)==0:
            if a.is_open[i]==1:
                x+=1
                if x==n:
                    tradedate.append(a.cal_date[i])
                    x=0
        else:
            if a.is_open[i]==1 and tradedate[-1][4:6]!=a.cal_date[i][4:6]:
                x+=1
                if x==n:
                    tradedate.append(a.cal_date[i])
                    x=0
    if len(tradedate)>0:
        for c in tradedate:
            d=datetime.datetime(int(c[:4]),int(c[4:6]),int(c[6:]))
            b.append(d) 
    print('获取每个月的第n个交易日的日期')
    return b

#获取远期合约时间
def time_change(current):
    dt=current.month
    if dt in [2,3,4]:
        longdata=[0,9]
    elif dt in [5,6,7]:
        longdata=[0,12]
    elif dt in [8,9,10]:
        longdata=[1,3]
    elif dt in [11,12]:
        longdata=[1,6]
    elif dt ==1:
        longdata=[0,6]
    return longdata

#获取交易的合约的trading_code  
def check_seed(current):
    dt=str(current.month)
    it=time_change(current)
    x=it[0]+current.year
    x=str(x)[2:]
    y=str(it[1])
    if len(y)==1:
        y='0'+y
    security='510050.XSHG'
    current_price=get_bars(security,1, unit='1d',
            fields=['open', 'close', 'high', 'low', 'volume', 'money'],
            include_now=False, end_dt=current, fq_ref_date=None)[0][3]
    #print('上证50当前价格：%s:%s'%(current,current_price))
    a=current_price*1000
    b=a*1.08
    c=round(b/50)*50
    c=int(c)
    d=str(c)
    if len(dt)==1:
        dt='0'+dt
    shortcode='510050C'+str(current.year)[2:]+dt+'M0'+d
    longcode='510050P'+x+y+'M0'+d
    return shortcode,longcode

#选择交易合约代码，交割日
def contract(current,code):
    date=current.strftime('%Y-%m-%d')
    q=query(opt.OPT_DAILY_PREOPEN.code,
    opt.OPT_DAILY_PREOPEN.expire_date,
    opt.OPT_DAILY_PREOPEN.trading_code,
    opt.OPT_DAILY_PREOPEN.date,
    ).filter(opt.OPT_DAILY_PREOPEN.trading_code==code,
    opt.OPT_DAILY_PREOPEN.date==date,
    )
    df=opt.run_query(q)
    return df

#获取合约的日行情数据
def quotation(current,code,expire_date):
    df_price=get_bars(code,count=50,unit='1d',
        fields=['date','open','high','low','close'],
        include_now=False, end_dt=expire_date)
    #把numpy数据类型，转化为pd类型
    df_price=pd.DataFrame(df_price,columns=['date','open','high','low','close'])
    #注意datetime.datetime和datetime.date不能比较，需要用方法date()转化。
    #选取当前时间之后的数据
    df_price=df_price[df_price.date>(current-datetime.timedelta(1)).date()]
    #重置索引，丢掉原来的索引
    df_price = df_price.reset_index(drop = True)
    #返回行情数据
    return df_price

#获取交易数据结果
def trede():
    monthdate=tradedate(start='20180101',end='20190630',n=7)
    for current in monthdate:
        shortdata,longdata=check_seed(current)
        df1=contract(current,shortdata)
        df2=contract(current,longdata)
        #因为有的月份数据没有，我被聚宽坑惨了，
        if df1.shape[0]!=0 and df2.shape[0]!=0:
            #获取合约交割日期
            expire_date=df1['expire_date'][0]
            #获取合约代码
            shortcode=df1['code'][0]
            longcode=df2['code'][0]
            #获取日行情数据
            df_price1=quotation(current,shortcode,expire_date)
            df_price2=quotation(current,longcode,expire_date)
            #因为有的月份数据缺失，所有需要判断
            if df_price1.shape[0]!=0 and df_price2.shape[0]!=0:
                a=df_price1.close[0]
                b=df_price2.close[0]
                num=int(b/a)
                c=a*10000*num
                c=round(c,2)
                d=b*10000
                d=round(d,2)
                ct=current.strftime('%Y-%m-%d')#时间格式转换
                opentrade_s.append({'date':ct,'so':c})
                opentrade_l.append({'date':ct,'lo':d})
                #卖出条件，盘中最高价涨10倍，收盘涨5倍，到期前两天
                for i in range(len(df_price1)-2):
                    if df_price1.high[i]>10*a:
                        x=round(df_price1.high[i]*10000*num)
                        y=round(df_price2.high[i]*10000)
                        closetrade_s.append({'date':ct,'sc':x})
                        closetrade_l.append({'date':ct,'lc':y})
                        break
                    elif df_price1.close[i]>5*a or i==len(df_price1)-3:
                        x=round(df_price1.close[i]*10000*num)
                        y=round(df_price2.close[i]*10000)
                        closetrade_s.append({'date':ct,'sc':x})
                        closetrade_l.append({'date':ct,'lc':y})
                        break
    s_1=pd.DataFrame(opentrade_s)
    s_2=pd.DataFrame(closetrade_s)
    l_1=pd.DataFrame(opentrade_l)
    l_2=pd.DataFrame(closetrade_l)
    #合并数据表
    #fally=pd.concat([s_1,s_2,l_1,l_2],axis=1,ignore_index=True)
    fally=s_1.merge(s_2,on='date')
    fally=fally.merge(l_1,on='date')
    fally=fally.merge(l_2,on='date')
    return fally

if __name__=='__main__':
    fally=trede()
    fally['share']=fally['sc']+fally['lc']-fally['so']-fally['lo']
    fally['profit']=fally.share.cumsum()
    print(fally)
    plt.figure()
    plt.plot(fally.date,fally.profit,label='profit')
    plt.plot(fally.date,fally.share,color='red',linewidth=1.0,linestyle='--',label='share')
    plt.legend(loc='best')#打印图例
    plt.xlabel('日期')
    plt.ylabel('元')
    plt.show()