# -*- coding: utf-8 -*-

import logging
import backtrader as bt
import pandas as pd
import numpy as np
import utils
from strategy import BaseStrategy
from openpyxl import Workbook
from openpyxl.styles import Alignment


def write_summary_sheet(wb, run_data_list, params):
    stock_size = len(run_data_list)
    total_cash = params['cash'] * stock_size
    sheet = wb.active
    sheet.column_dimensions['A'].width = 14
    sheet.column_dimensions['A'].alignment = Alignment(horizontal='center')
    sheet.column_dimensions['B'].width = 14
    sheet.title = u'回测总览'
    appender = utils.WbSheetAppender(sheet=sheet)
    appender.write_row(u'复权方式', u'前复权')
    appender.write_row(u'起始日期', params['start_date_str'])
    appender.write_row(u'终止日期', params['end_date_str'])
    appender.write_row(u'周期（日）', params['period'])
    appender.write_row(u'每股起始资金', params['cash'])
    appender.write_row(u'手续费率', params['commission'])
    appender.write_row(u'每股起始资金', params['cash'])
    appender.write_row(u'股票池数目', stock_size)
    appender.write_row(u'成交量倍数', params['expand_ratio'])
    appender.write_row(u'买入方式', params['buy_method'])
    appender.write_row()
    appender.write_row()

    all_trade_history_df = map(lambda run_data: run_data.trade_history_df,
                               filter(lambda run_data: run_data.trade_history_df.size != 0, run_data_list)
                               )
    final_df = pd.concat(all_trade_history_df)

    final_df_sum = final_df.sum()

    appender.write_row(u'利润', '%.02f' % final_df_sum['gross'])
    appender.write_row(u'利润率', '%.02f%%' % (final_df_sum['gross'] * 100.0 / total_cash))
    appender.write_row(u'净利润', '%.02f' % final_df_sum['net'])
    appender.write_row(u'净利润率', '%.02f%%' % (final_df_sum['net'] * 100.0 / total_cash))

    print '利润     : ', '%.02f' % final_df_sum['gross']
    print '利润率   : ', '%.02f%%' % (final_df_sum['gross'] * 100.0 / total_cash)
    print '净利润   : ', '%.02f' % final_df_sum['net']
    print '净利润率 : ', '%.02f%%' % (final_df_sum['net'] * 100.0 / total_cash)

    total_op_count = final_df.index.size
    gross_series = final_df['gross']
    success_gross_op_count = gross_series[gross_series > 0].size
    net_series = final_df['net']
    success_net_op_count = net_series[net_series > 0].size

    appender.write_row()
    appender.write_row()
    appender.write_row(u'总操作次数', '%d' % total_op_count)
    appender.write_row(u'成功次数', '%d' % success_gross_op_count)
    appender.write_row(u'成功率', '%.02f%%' % (success_gross_op_count * 100.0 / total_op_count))
    appender.write_row(u'净成功次数', '%d' % success_net_op_count)
    appender.write_row(u'净成功率', '%.02f%%' % (success_net_op_count * 100.0 / total_op_count))

    print '总操作次数 : ', '%d' % total_op_count
    print '成功次数   : ', '%d' % success_gross_op_count
    print '成功率     : ', '%.02f%%' % (success_gross_op_count * 100.0 / total_op_count)
    print '净成功次数 : ', '%d' % success_net_op_count
    print '净成功率   : ', '%.02f%%' % (success_net_op_count * 100.0 / total_op_count)


def write_trade_sheet(wb, run_data_list, params):
    stock_size = len(run_data_list)
    total_cash = params['cash'] * stock_size
    sheet = wb.create_sheet(title=u"个股收益总览")
    sheet.column_dimensions['A'].width = 14
    # sheet.column_dimensions['B'].width = 15
    appender = utils.WbSheetAppender(sheet=sheet)
    appender.write_row(u"股票代号", u"利润", u"利润率", u"净利润", u"净利润率",
                       u"成功次数", u"成功率", u"净成功次数", u"净成功率")
    for run_data in run_data_list:
        op_count = run_data.trade_history_df.index.size
        if op_count == 0:
            appender.write_row(run_data.stock_id, u"无交易")
            continue

        cash = params['cash']
        sum_df = run_data.trade_history_df.sum()
        gross_op_series = run_data.trade_history_df['gross']
        success_gross_op_count = gross_op_series[gross_op_series > 0].size
        net_op_series = run_data.trade_history_df['net']
        success_net_op_count = net_op_series[net_op_series > 0].size

        appender.write_row(run_data.stock_id,
                           sum_df['gross'], '%.02f%%' % (sum_df['gross'] * 100.0 / cash),
                           sum_df['net'], '%.02f%%' % (sum_df['net'] * 100.0 / cash),
                           success_gross_op_count, '%.02f%%' % (success_gross_op_count * 100.0 / op_count),
                           success_net_op_count, '%.02f%%' % (success_net_op_count * 100.0 / op_count))


def write_order_history_sheet(wb, run_data_list, params):
    stock_size = len(run_data_list)
    total_cash = params['cash'] * stock_size
    sheet = wb.create_sheet(title=u"交易明细")
    sheet.column_dimensions['A'].width = 14
    sheet.column_dimensions['B'].width = 14
    appender = utils.WbSheetAppender(sheet=sheet)
    appender.write_row(u"股票代号", u"交易日期", u"价格", u"股数")
    for run_data in run_data_list:
        op_count = run_data.order_history_df.index.size
        appender.write_row(run_data.stock_id)

        if op_count == 0:
            appender.write_row("", u"无交易")
            appender.write_row()
            continue

        for order in run_data.order_history_df.itertuples():
            appender.write_row('', order.Index.strftime('%Y-%m-%d %H:%M:%S'), order.price, order.size)

    appender.write_row()


class TSStockDetailData(bt.feeds.PandasData):
    lines = ('open', 'close', 'high', 'low',)
    params = (
        ('datetime', None),
        ('open', -1),
        ('high', -1),
        ('close', -1),
        ('low', -1),
    )
    datafields = [x[0] for x in params]

    def __init__(self):
        super(TSStockDetailData, self).__init__()


def get_ts_trade_detail_df(stock_id, begin_date_str, end_date_str, cycle='30分钟线'):
    begin_date = utils.str_to_ts_date(begin_date_str)
    end_date = utils.str_to_ts_date(end_date_str)
    trade_detail_list = utils.ts_remote_call_func('TradeDetails',
                                                  [begin_date, end_date, cycle.decode('utf8').encode('gbk')], {
                                                      'StockID': stock_id,
                                                      'bRate': 2,  # 前复权
                                                      'RateDay': end_date
                                                  })
    row_size = len(trade_detail_list)
    if row_size == 0:
        return None

    return utils.ts_list_to_dataframe(
        ts_list=trade_detail_list,
        columns=[('datetime', '时间'), ('open', '开盘价'), ('close', '收盘价'), ('high', '最高价'), ('low', '最低价')],
        index_type=pd.DatetimeIndex
    )


def get_ts_stock_period_hq_df(stock_id, begin_date_str, end_date_str, cycle='日'):
    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, cycle, 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 RunData(object):
    """
    回测的数据记录结构，包括参数和交易记录，用于并行化处理
    """

    def __init__(self, *args, **kwargs):
        for k, v in kwargs.iteritems():
            setattr(self, k, v)
        self.kwargs = kwargs
        self.args = args
        self.order_history_df = None
        self.trade_history_df = None


# 单只股票，条件:日成交量放大为前10个交易日平均成交量的4倍及以上，同时当天涨幅为正（涨停除外，也可以单独测试不去涨停的效果）。
# 处理方式：当天收盘前卖出，测试隔天开盘15分钟内以最低价买回来
# 需要数据：
# 历史成交量，当天涨幅，个股

class VolExpandStrategy(BaseStrategy):
    def __init__(self, run_data, detail_15m_df=None, period_hq_df=None):
        super(VolExpandStrategy, self).__init__()
        # 准备回测需要的数据
        self.run_data = run_data
        self.last_sell_order = None
        self.last_buy_order = None
        self.last_buy_period = None
        self.detail_15m_df = detail_15m_df
        self.period_hq_df = period_hq_df
        # 1个交易日有16个 15分钟，所以周期*16
        self.sma_close = bt.talib.SMA(self.data.close, timeperiod=self.run_data.period * 16)

    def previous_period_average_volume(self):
        period_idx = len(self) / 16
        start_idx = period_idx - self.run_data.period
        return self.period_hq_df['volume'][start_idx: period_idx].mean() * 1.0

    def get_current_df_idx(self):
        return len(self) - 1

    def next(self):
        period_idx = len(self) / 16  # 16 * 15 minutes = 4 hours. 全日交易时间，周期坐标，从0开始
        if period_idx < self.run_data.period:
            return
        dt = self.datas[0].datetime.datetime(0)
        if dt.hour == 15 and dt.minute == 00:
            if self.position:
                price = self.data.low[1] if self.run_data.buy_method == 'low' else self.data.open[1]  # 使用未来数据
                self.last_buy_order = self.close(price=price, exectype=bt.Order.Limit)
                self.last_buy_period = period_idx
        elif dt.hour == 14 and dt.minute == 45:
            prev_avg_volume = self.previous_period_average_volume()
            sell_signal_triggered = False
            # 1、交易量是过去平均周期平均值 expand_ratio 倍数以上
            # 2、涨幅为正
            # 3、不达到涨停价格
            # 4、T+1 限制
            if self.period_hq_df['volume'][period_idx] > (prev_avg_volume * self.run_data.expand_ratio) and \
                            self.period_hq_df['p_change'][period_idx] > 0 and \
                            self.period_hq_df['close'][period_idx] < utils.calc_limit_up_price(
                        self.period_hq_df['close'][period_idx - 1]) and \
                            period_idx != self.last_buy_period:
                # print "DETECT %.02f > %.02f * %.1f" % (
                #     self.period_hq_df['volume'][period_idx], prev_avg_volume, self.run_data.expand_ratio)
                sell_signal_triggered = True

            if sell_signal_triggered and not self.position and len(self) < self.detail_15m_df.index.size - 1:
                # 使用了未来数据，最后2个bar不允许交易
                price = self.data.close[1]  # 使用未来数据
                buy_size = int((self.broker.cash / (1 + self.run_data.commission) / price) / 100) * 100
                self.last_sell_order = self.sell(size=buy_size, price=price, exectype=bt.Order.StopLimit)

    def start(self):
        # [u'股票代码', u'交易日期', u'价格', u'手数', u'总金额', u'']
        print '%s START' % self.run_data.stock_id

    def stop(self):
        # 打印交易详情
        super(VolExpandStrategy, self).stop()
        print '%s END' % self.run_data.stock_id
        self.run_data.order_history_df = self.order_history_df
        self.run_data.trade_history_df = self.trade_history_df


def worker_func(run_data):
    logging.basicConfig(level=logging.DEBUG)
    detail_15m_df = get_ts_trade_detail_df(run_data.stock_id, run_data.start_date_str, run_data.end_date_str, '15分钟线')
    period_hq_df = get_ts_stock_period_hq_df(run_data.stock_id, run_data.start_date_str, run_data.end_date_str)
    cerebro = bt.Cerebro()
    cerebro.adddata(TSStockDetailData(dataname=detail_15m_df), name=run_data.stock_id)
    cerebro.addstrategy(VolExpandStrategy, run_data=run_data, detail_15m_df=detail_15m_df, period_hq_df=period_hq_df)
    cerebro.broker.setcash(run_data.cash)
    cerebro.broker.setcommission(commission=run_data.commission)
    cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')
    results = cerebro.run(once=True)
    strat = results[0]
    pyfoliozer = strat.analyzers.getbyname('pyfolio')
    returns, positions, transactions, gross_lev = pyfoliozer.get_pf_items()
    transactions.index = transactions.index.map(lambda dt: dt.replace(hour=0, minute=0, second=0, microsecond=0))
    def make_line_data_from_linebuffer(linebuffer):
        return linebuffer.array[15:len(linebuffer):16]

    value_list = make_line_data_from_linebuffer(cerebro.runstrats[0][0].stats.broker.value)
    cash_list = make_line_data_from_linebuffer(cerebro.runstrats[0][0].stats.broker.cash)

    cash_value_df = pd.DataFrame(
        data=zip(value_list, cash_list),
        index=returns.index.copy(),
        columns=['value', 'cash']
    )
    cash_value_df.index.rename('datetime', inplace=True)

    cerebro.plot()

    # utils.plot_html(run_data.stock_id,
    #                 datetime_ohlc_df=period_hq_df,
    #                 order_history_df=transactions,
    #                 cash_value_df=cash_value_df,
    #                 cerebro=cerebro,
    #                 index_format=lambda index: index.strftime('%Y-%m-%d'),
    #                 linebuf_sample_step=16,
    #                 show=True)
    return run_data


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)

    params = dict(
        start_date_str='2016-01-01',
        end_date_str='2017-04-14',
        period=10,
        cash=100000.0,
        commission=0.001,
        expand_ratio=2,
        buy_method='open'  # open：以开盘价买入    low：开盘15分钟内最低价买入
    )

    # stock_list = utils.get_stock_pool()
    stock_list = [utils.get_stock_pool()[10]]

    run_data_list = map(lambda stock_id: RunData(stock_id=stock_id, **params), stock_list)

    utils.parallel_run(worker_func, run_data_list, process_num=1)

    # 汇总结果，输出Excel，保存
    wb = Workbook()
    write_summary_sheet(wb, run_data_list, params)
    write_trade_sheet(wb, run_data_list, params)
    write_order_history_sheet(wb, run_data_list, params)

    wb.save(filename=('VolExpandStrategy_%s_%s_%s_%d.xlsx' % (
        params['start_date_str'], params['end_date_str'],
        params['buy_method'], params['expand_ratio'] * 10)))
