import jqdata
import pandas as pd
import numpy as np
import datetime
import math
import scipy.stats as stats
import matplotlib.pyplot as plt
matplotlib.rcParams['axes.unicode_minus']=False # 解决负号显示异常的问题
from jqlib.technical_analysis import *
from jqdata import *
from sklearn import linear_model
from sklearn.metrics import mean_squared_error, r2_score

# 使用聚宽api读取指数数据，参数ben为指数代码，end_date为数据截止日
# 返回一个dataframe，index为交易日，columns为开、收盘价等基本数据，数据起始为指数上市日（最早到聚宽数据初始的2005年1月1日）
def get_data(ben,end_date):
    ben_sd = get_security_info(ben).start_date
    sdate = (ben_sd if ben_sd>datetime.date(2005,1,1) else '2005-01-01')   
    return get_price(ben,sdate,end_date)

# 读取数据文件，参数为文件名（带完整路径），返回dataframe包含基本数据，index为日期
def load_data(filename):
    df = pd.read_csv(filename,encoding='gbk')
    df['date'] = df.ix[:,0]
    del df[df.columns[0]]
    df = df.set_index(['date'],drop=True)
    # 将dataframe的数据全部转换成数字型，方便后续操作
    df = df.convert_objects(convert_numeric=True)
    
    return df

# 计算最大回撤，参数为策略净值序列（list or series）
# 思路：i从序列的第二个值开始，依次检查序列从头到i这一段的max和i值，计算是否有回撤（与0比较）并记录；
# 取所有回撤中的最大值即为最大回撤（负值取最小值），找到最小回撤点ie，对应回撤终点，然后找到从头开始到ie这段之间的最大值即为回撤起点
# 返回最大回撤值（正），回撤区间（序列index，如果传入的是series，返回series的index，list返回下标）
def calc_max_drawdown(rets):
    dd=[0]*len(rets)
    tmp_mdd=0
    for j in range(1,len(rets)):
        dd[j] = min(rets[j]/max(rets[:j+1])-1, 0)
        tmp_mdd = min(min(dd[:j+1]),tmp_mdd)
#     print (tmp_mdd)
    ie=np.argmin(dd)
    ist=np.argmax(rets[:ie])
#     print (ie,ist)
    return -tmp_mdd*100.0,ist,rets.index[ie]

# 计算N日均线，参数为指数数据ind_df，时间段窗口N，返回dataframe，包含指数基本数据和新加列MA_N
# 注意，为防止未来函数，第i天存储的是i-1天收盘后求出的N日均线值
def calc_MA(ind_df,N):
    h=ind_df.copy()
    for i in range(N+1,len(h)):
        h.loc[h.index[i],'MA_N'] = h['close'][(i-N-1):i].mean()
        
    return np.round(h,6)

# 计算N日均线择时策略，参数为指数数据ben，回测起止日s_d和e_d，以及手续费cpr（单边）
# 返回一个净值记录dataframe，index为交易日，columns包括净值序列（含and不含手续费）、开平仓信号、持仓状态
def calc_strategy_result(ben,s_d,e_d,cpr):
    # 均线参数
    N=5
    th=0.01
    
    # 获取数据
    df=calc_MA(ben,N)
    df=df.dropna()
    df_t=df[s_d:e_d] # 截取回测区间数据
    
    is_holding=False # 持仓标志位
    nav=pd.DataFrame(columns={'etf','etf_cost','benchmark','flag','position','excess_return'},index=df_t.index)
    nav.loc[:,:]=-36.9 #初始化

    # 执行策略,从回测区间第一个交易日开始逐日循环，遇到开仓信号开仓、平仓信号平仓，同时记录净值及相关操作
    # nav记录的是每日收盘后的数据及状态，策略默认开盘时执行
    for i in range(0,len(df_t)):
        navv = (1.0000 if i==0 else nav['etf'][i-1])       # 开盘前临时净值，初始为1，后续为昨日收盘后净值记录
        navc = (1.0000 if i==0 else nav['etf_cost'][i-1])  # 对应含手续费净值临时值
        if not is_holding: # 当前空仓
            if df_t['open'][i]>((1+th)*df_t['MA_N'][i]): # 遇到开仓信号，买入
                # 开盘价买入，故今日收盘后净值记录应对应增长close/open
                nav['etf'][i] = df_t['close'][i] / df_t['open'][i] * navv
                nav['etf_cost'][i] = df_t['close'][i] / df_t['open'][i] * navc * (1-cpr)
                # 记录今日信号flag=1，今日收盘后持仓状态position=1
                nav['flag'][i]=1
                nav['position'][i]=1
                is_holding=True
            else:  # 没有开仓信号，等待
                # 今日未开仓，故今日净值较昨日无变化
                nav['etf'][i] = navv
                nav['etf_cost'][i] = navc
                # 记录今日信号flag=0，今日收盘后持仓状态position=0
                nav['flag'][i]=0
                nav['position'][i]=0
        else:  # 当前持仓
            if df_t['open'][i]<((1-th)*df_t['MA_N'][i]): # 遇到平仓信号，卖出
                # 开盘价买卖出，故今日收盘后净值记录应对应增长open/last-close
                nav['etf'][i] = df_t['open'][i] / df_t['close'][i-1] * navv
                nav['etf_cost'][i] = df_t['open'][i] / df_t['close'][i-1] * navc *(1-cpr)
                # 记录今日信号flag=-1，今日收盘后持仓状态position=0
                nav['flag'][i]=-1
                nav['position'][i]=0
                is_holding=False
            else: # 没有平仓信号，继续持有
                # 今日继续持有，故净值对应增长close/last-close
                nav['etf'][i] = df_t['close'][i] / df_t['close'][i-1] * navv
                nav['etf_cost'][i] = df_t['close'][i] / df_t['close'][i-1] * navc
                # 记录今日信号flag=0，今日收盘后持仓状态position=1
                nav['flag'][i]=0
                nav['position'][i]=1
    
    # 计算基准（即指数本身）净值序列
    tmp=df_t.copy()
    nav.loc[:,'benchmark'] = tmp.loc[:,'close'] / tmp.loc[tmp.index[0],'open']
    
    nav.loc[:,'excess_return']=nav.loc[:,'etf'] / nav.loc[:,'benchmark']-1.0 # 计算超额收益
    
    
    return np.round(nav,6)

def calc_risk_metrics(nav):
    risks={}
    Rf=0.04 #无风险收益，年化
    
    # 计算收益
    risks['algorithm_return'] = nav['etf'][len(nav)-1]-1
    risks['annual_algo_return']=math.pow(nav['etf'][len(nav)-1],250/len(nav))-1
    risks['benchmark_return']=nav['benchmark'][len(nav)-1]-1
    risks['annual_bm_return']=math.pow(nav['benchmark'][len(nav)-1],250/len(nav))-1
    
    # 计算alpha，beta，sharpe，volatility及IR，参考聚宽
    rets=nav[['etf','benchmark']].pct_change()[1:]
    risks['beta']=rets.cov().iloc[0,1]/rets.cov().iloc[1,1]
    risks['alpha']=risks['annual_algo_return'] - risks['beta']*(risks['annual_bm_return']-Rf) - Rf
    risks['algorithm_volatility']=rets['etf'].std() * math.sqrt(250)
    risks['benchmark_volatility']=rets['benchmark'].std() * math.sqrt(250)
    risks['sharpe']=(risks['annual_algo_return']-Rf)/risks['algorithm_volatility']
    diffvol=(rets['etf']-rets['benchmark']).std() * math.sqrt(250)
    risks['information_ratio']=(risks['annual_algo_return']-risks['annual_bm_return'])/diffvol
    
    # 计算交易次数，盈利、亏损数，胜率，盈亏比以及持仓天数
    risks['test_days'] = len(nav)
    risks['holding_days']=nav['position'].sum()
    buy_count=len(nav[nav['flag']==1])
    sell_count=len(nav[nav['flag']==-1])
    risks['trade_count']=[buy_count,sell_count]
    # 遍历净值记录，统计盈利交易数及亏损数
    wc=0
    lc=0
    cp=0.0
    cl=0.0
    for i in range(0,len(nav)):
        if nav['flag'][i]==1:
            bbnav=(1.0000 if i==0 else nav['etf'][i-1]) # 盈亏基准约为买入日前一日净值
            for j in range(i+1,len(nav)):
                if nav['flag'][j]==-1: # 找到每一次买入后最近的一次卖出
                    if nav['etf'][j]>=bbnav:  # 忽略买入日开盘价相对前一日收盘价变化
                        wc += 1
                        cp += nav['etf'][j]-bbnav
                    else:
                        lc += 1
                        cl += abs(nav['etf'][j]-bbnav)
                    break
    risks['win_count']=wc
    risks['lose_count']=lc
    risks['win_ratio']=2.0*wc/sum(risks['trade_count'])
    risks['profit_to_loss_ratio']=(cp/cl if cl!=0.0 else 99999999.999)
    
    # 计算最大回撤及对应区间
    risks['max_drawdown'],t1,t2=calc_max_drawdown(nav['etf'])
    risks['max_drawdown_period']=[t1.date(),t2.date()]
    
    return risks

# 测试指数沪深300，首先获取数据
test_index='000300.XSHG'
end_date='2018-08-26'
h=get_data(test_index,end_date)

#回测时间段
bt_start='2010-01-01'
bt_end='2018-08-26'

cost_per_trade=0.001 / 2 # 手续费双边千分之一，单边万五

s_time = datetime.datetime.now()

ind_nav_record=calc_strategy_result(h,bt_start,bt_end,cost_per_trade)

e_time = datetime.datetime.now()
interval = (e_time - s_time).seconds
print ('carry out strategy cost : ' + str(interval) + ' seconds.')

test_results=calc_risk_metrics(ind_nav_record)
print ('risk metrics:\n%s' %(test_results))

ind_nav_record[['etf','etf_cost','benchmark','excess_return']].plot()