# -*- coding: utf-8 -*-
import config
import os
import pandas as pd
import utils
import talib
from functional import seq
from sys import exit
import backtrader as bt
import qrisk


def get_ts_stock_period_df(stock_id, begin_date_str, end_date_str):
    begin_date = utils.str_to_ts_date(begin_date_str)
    end_date = utils.str_to_ts_date(end_date_str)
    # 交易所算法，前复权
    stock_daily_data = utils.ts_remote_call_func('Stock_PeriodHq', [begin_date, end_date, '日', 2, 0], {
        'StockID': stock_id
    })
    row_size = len(stock_daily_data)
    if row_size == 0:
        return None
    stock_daily_data.reverse()  # 天软查询接口是倒序，转换成按照时间先后排序，日期小的排前面
    return utils.ts_list_to_dataframe(
        ts_list=stock_daily_data,
        columns=[('datetime', '日期'), ('open', '开盘'), ('close', '收盘'), ('high', '最高'), ('low', '最低'),
                 ('p_change', '涨幅(%)'), ('volume', '成交量(手)'), ('turnover', '换手率(%)')],
        index_type=pd.DatetimeIndex
    )


class T1ShortRunner(object):
    def __init__(self,
                 factor=None,
                 begin_date='2012-01-01',
                 end_date='2017-06-01',
                 prepare_full_period_df_func=None,
                 datafeed_cls=None,
                 strategy_cls=None,
                 strategy_args=[],
                 strategy_kwargs={},
                 full_period_progress_func=None,
                 period_progress_func=None):
        assert begin_date >= '2012-01-01'
        assert begin_date <= end_date

        self.factor = factor
        self.begin_date = begin_date
        self.end_date = end_date
        self.balance_date_list = []
        self.position_weight_df_list = []
        self.stock_pool = []
        self.stock_full_period_df_dict = {}
        self.datafeed_cls = datafeed_cls
        self.prepare_full_period_df_func = prepare_full_period_df_func
        self.full_period_progress_func = full_period_progress_func
        self.period_progress_func = period_progress_func
        self.strategy_cls = strategy_cls
        self.strategy_args = strategy_args
        self.strategy_kwargs = strategy_kwargs
        self.transactions_df = None
        self.returns_df = None
        self.cumulative_return_series = None
        self.non_cumulative_returns_series = None

    def _load_factor_position_plan(self):
        dir_path = os.path.join(config.DATA_ROOT, 'factor_position', self.factor)
        self.balance_date_list = os.listdir(dir_path)
        self.balance_date_list = map(lambda x: x[:-4], self.balance_date_list)
        self.balance_date_list = filter(lambda x: self.begin_date <= x <= self.end_date, self.balance_date_list)

        def pre_process_df(df):
            df.index = pd.Index(map(lambda x: utils.stock_format(x, fmt=utils.StockFormat.TINYSOFT),
                                    df['stock'].tolist()))
            return df.drop('stock', axis=1)

        self.position_weight_df_list = seq(self.balance_date_list) \
            .map(lambda x: pd.read_csv(os.path.join(config.DATA_ROOT, 'factor_position', self.factor, x + '.csv'),
                                       index_col=0)) \
            .map(pre_process_df).to_list()

        self.stock_pool = seq(self.position_weight_df_list) \
            .map(lambda df: df.index.tolist()) \
            .reduce(lambda a, b: a + b) \
            .distinct().list()

        for idx, stock in enumerate(self.stock_pool):
            if self.full_period_progress_func:
                self.full_period_progress_func(idx, len(self.stock_pool), stock)
            # print "loading {} / {}".format(idx + 1, len(self.stock_pool))
            df = utils.load_period_hq_df(stock, self.begin_date, self.end_date)
            if self.prepare_full_period_df_func is not None and df is not None:
                self.stock_full_period_df_dict[stock] = self.prepare_full_period_df_func(df, stock)
            else:
                self.stock_full_period_df_dict[stock] = df

    def run(self):
        self._load_factor_position_plan()
        init_cash = 10 ** 9
        period_cash = init_cash
        transactions_df = None
        cumulative_returns_series = pd.Series(index=pd.DatetimeIndex([]))

        for period_idx, begin_date in enumerate(self.balance_date_list[:-1]):
            end_date = self.balance_date_list[period_idx + 1]
            position_weight_df = self.position_weight_df_list[period_idx]
            position_weight_df['cash'] = position_weight_df['weight'] * period_cash
            # 按个股执行回测

            trade_date_list = utils.get_trade_date_str_list(begin_date, end_date)[:-1]
            value_df = pd.DataFrame(index=pd.DatetimeIndex(trade_date_list), columns=position_weight_df.index.tolist())

            for stock_idx, stock in enumerate(position_weight_df.index):
                if self.period_progress_func:
                    self.period_progress_func(begin_date, end_date, stock_idx, position_weight_df)
                period_df = self.stock_full_period_df_dict[stock]
                empty_flag = True
                if period_df is not None:  # 忽略没有数据的个股
                    # print stock, period_df.index.size
                    period_df = period_df.loc[pd.to_datetime(begin_date) <= period_df.index]
                    period_df = period_df.loc[period_df.index < pd.to_datetime(end_date)]
                    if not period_df.empty:
                        cerebro = bt.Cerebro(cheat_on_open=True, preload=True)
                        cerebro.adddata(self.datafeed_cls(dataname=period_df), name=stock)
                        cerebro.addstrategy(self.strategy_cls, *self.strategy_args, **self.strategy_kwargs)
                        cerebro.broker.setcash(position_weight_df.loc[stock]['cash'])
                        cerebro.broker.setcommission(0.0008)
                        cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')
                        # cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')
                        strategy = cerebro.run()[0]
                        pyfolio = strategy.analyzers.getbyname('pyfolio')
                        returns, positions, transactions, gross_lev = pyfolio.get_pf_items()
                        if not transactions.empty:
                            empty_flag = False
                            value_series = positions.sum(axis=1)
                            if transactions_df is None:
                                transactions_df = transactions.copy()
                            else:
                                transactions_df = transactions_df.append(transactions)
                            for dt in value_series.index:
                                value_df.at[dt.date(), stock] = value_series.get(dt)
                if empty_flag:
                    value_df[stock] = position_weight_df.at[stock, 'cash']

            # end for single period

            value_df = value_df.fillna(method='ffill').fillna(method='bfill')
            orig_period_cash = period_cash
            period_cash = value_df.iloc[-1, :].sum()
            cumulative_returns_series = cumulative_returns_series.append(value_df.sum(axis=1))

            # print begin_date + ' - ' + end_date
            # print '       stock poll size    %d' % position_weight_df.index.size
            # print '   value df index size    %d' % value_df.columns.size
            # print '         initial value    %.2f' % orig_period_cash
            # print '             new value    %.2f' % period_cash
            # print 'cumulative return rate    %.2f%%' % (period_cash / init_cash * 100 - 100)

        self.transactions_df = transactions_df
        if not transactions_df.empty:
            writer = pd.ExcelWriter(
                'transactions_{}_{}_{}.xlsx'.format(self.factor, self.begin_date, self.end_date))
            transactions_df.index.tz = None
            transactions_df.drop('sid', axis=1).to_excel(writer, u'交易记录')
            writer.save()
        cumulative_returns_series /= init_cash
        self.cumulative_return_series = cumulative_returns_series
        self.non_cumulative_returns_series = cumulative_returns_series.diff().fillna(0)


if __name__ == '__main__':
    import t1_short.short_v3

    runner = T1ShortRunner(
        factor='test_size',
        begin_date='2012-01-01',
        end_date='2016-01-01',
        datafeed_cls=t1_short.short_v3.PeriodHQData,
        strategy_cls=t1_short.short_v3.Strategy,
        prepare_full_period_df_func=t1_short.short_v3.prepare_full_period_df_func
    )

    runner.run()
3