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

import logging
import backtrader as bt
import pandas as pd
import utils
from strategy import BaseStrategy
from openpyxl import Workbook
from openpyxl.styles import Alignment
from openpyxl.utils import get_column_letter
from multiprocessing import Pool

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


class YourData(bt.feeds.PandasData):
    """
    通常包含OCHL，根据需要追加
    """
    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(YourData, self).__init__()

class YourStrategy(BaseStrategy):
    """
    编写你的策略
    """

    def __init__(self, run_data=None):
        super(YourStrategy, self).__init__()
        self.stock_id = run_data.stock_id
        self.run_data = run_data

    def start(self):
        print '%s START' % self.stock_id

    def next(self):
        pass

    def stop(self):
        # 打印交易详情，汇总数据成 DataFrame
        super(YourStrategy, self).stop()
        print '%s END' % self.stock_id
        # 汇总结果到 run_data
        self.run_data.order_history_df = self.order_history_df
        self.run_data.trade_history_df = self.trade_history_df


def run_strategy(run_data):
    logging.basicConfig(level=logging.DEBUG)
    cerebro = bt.Cerebro()
    cerebro.adddata(YourData(dataname=None))    # 添加数据，用DataFrame替换None
    cerebro.addstrategy(YourStrategy, run_data=run_data)
    cerebro.broker.setcash(run_data.cash)
    cerebro.broker.setcommission(commission=run_data.commission)
    cerebro.run()
    return run_data


def write_summary_sheet(wb, run_data_list, params):
    """
    输出汇总的Excel
    :param wb:              WorkBook 
    :param run_data_list:   list<RunData>
    :param params:          本轮测试所用的数据，参照  __main__ 里头的声明
    :return: 
    """
    stock_size = len(run_data_list)
    total_cash = params['cash'] * stock_size
    sheet = wb.active
    sheet.column_dimensions['A'].width = 14  # 设置列宽
    sheet.title = u'回测总览'
    appender = utils.WbSheetAppender(sheet=sheet)  # 方便按行写入 Excel
    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['cash'])
    appender.write_row(u'手续费率', params['commission'])
    appender.write_row(u'每股起始资金', params['cash'])
    appender.write_row(u'股票池数目', stock_size)
    appender.write_row()
    appender.write_row()

    # 计算利润，输出
    appender.write_row(u'利润', '%.02f' % 0.0)
    appender.write_row(u'利润率', '%.02f%%' % 0.0)
    appender.write_row(u'净利润', '%.02f' % 0.0)
    appender.write_row(u'净利润率', '%.02f%%' % 0.0)

    print '利润     : ', '%.2f' % 0.0
    print '利润率   : ', '%.2f%%' % 0.0
    print '净利润   : ', '%.2f' % 0.0
    print '净利润率 : ', '%.2f%%' % 0.0


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

    # 为确保后期的扩展性，本次运行的全部参数，请在此声明
    params = dict(
        start_date_str='2016-01-01',
        end_date_str='2017-04-14',
        cash=100000.0,
        commission=0.0016,  # 手续费，买入、卖出单次一般0.0003，卖出还有印花税0.001，加起来估算0.0016
    )

    stock_list = utils.get_stock_pool()  # 读入公用股票池

    # 定义 RunData 对象构建过程，并声称参数列表
    run_data_list = map(lambda stock_id: RunData(stock_id=stock_id, **params), stock_list)

    # 并行化运行，如果 process_num = 1，则在当前进程运行，适用于调试，结果将替换run_data_list中的元素
    utils.parallel_run(run_strategy, run_data_list, process_num=6)

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

    wb.save(filename="YourStrategy.xlsx")
