# encoding: UTF-8

'''
本文件中包含的是CTA模块的回测引擎，回测引擎的API和CTA引擎一致，
可以使用和实盘相同的代码进行回测。
'''

from datetime import datetime, timedelta
from collections import OrderedDict
from itertools import product
import pymongo

# import MySQLdb
import json
import os
import cPickle
import csv

from ctaBase import *
from ctaSetting import *

from eventEngine import *
from vtConstant import *
from vtGateway import VtOrderData, VtTradeData
from vtFunction import loadMongoSetting
import logging
import copy
import pandas as pd
from ctaBase import *


########################################################################
class BacktestingEngine(object):
    """
    CTA回测引擎
    函数接口和策略引擎保持一样，
    从而实现同一套代码从回测到实盘。
    # modified by IncenseLee：
    1.增加Mysql数据库的支持；
    2.修改装载数据为批量式后加载模式。
    3.增加csv 读取bar的回测模式
    4.增加csv 读取tick合并价差的回测模式
    5.增加EventEngine，并对newBar增加发送OnBar事件，供外部的回测主体显示Bar线。
    """

    TICK_MODE = 'tick'  # 数据模式，逐Tick回测
    BAR_MODE = 'bar'  # 数据模式，逐Bar回测

    REALTIME_MODE = 'RealTime'  # 逐笔交易计算资金，供策略获取资金容量，计算开仓数量
    FINAL_MODE = 'Final'  # 最后才统计交易，不适合按照百分比等开仓数量计算

    # ----------------------------------------------------------------------
    def __init__(self, eventEngine=None, initCapital = 100000):
        """Constructor"""

        self.eventEngine = eventEngine

        # 本地停止单编号计数
        self.stopOrderCount = 0
        # stopOrderID = STOPORDERPREFIX + str(stopOrderCount)

        # 本地停止单字典
        # key为stopOrderID，value为stopOrder对象
        self.stopOrderDict = {}  # 停止单撤销后不会从本字典中删除
        self.workingStopOrderDict = {}  # 停止单撤销后会从本字典中删除

        # 引擎类型为回测
        self.engineType = ENGINETYPE_BACKTESTING

        # 回测相关
        self.strategy = None  # 回测策略
        self.mode = self.BAR_MODE  # 回测模式，默认为K线

        self.slippage = 0  # 回测时假设的滑点
        self.rate = 0  # 回测时假设的佣金比例（适用于百分比佣金）
        self.size = 1  # 合约大小，默认为1

        self.dbClient = None  # 数据库客户端
        self.dbCursor = None  # 数据库指针

        self.historyData = []  # 历史数据的列表，回测用
        self.initData = []  # 初始化用的数据
        self.backtestingData = []  # 回测用的数据

        self.dbName = ''  # 回测数据库名
        self.symbol = ''  # 回测集合名

        self.dataStartDate = None  # 回测数据开始日期，datetime对象
        self.dataEndDate = None  # 回测数据结束日期，datetime对象
        self.strategyStartDate = None  # 策略启动日期（即前面的数据用于初始化），datetime对象

        self.limitOrderDict = OrderedDict()  # 限价单字典
        self.workingLimitOrderDict = OrderedDict()  # 活动限价单字典，用于进行撮合用
        self.limitOrderCount = 0  # 限价单编号

        # 持仓缓存字典
        # key为vtSymbol，value为PositionBuffer对象
        self.posBufferDict = {}

        self.tradeCount = 0  # 成交编号
        self.tradeDict = OrderedDict()  # 成交字典

        self.logList = []  # 日志记录

        # 当前最新数据，用于模拟成交用
        self.tick = None
        self.bar = None
        self.dt = None  # 最新的时间
        self.gatewayName = u'BackTest'

        # csvFile相关
        self.barTimeInterval = 60  # csv文件，属于K线类型，K线的周期（秒数）,缺省是1分钟

        # 回测计算相关
        self.calculateMode = self.FINAL_MODE
        self.usageCompounding = False  # 是否使用简单复利 （只针对FINAL_MODE有效）

        self.initCapital = 100000  # 期初资金
        self.capital = self.initCapital  # 资金  （相当于Balance）
        self.maxCapital = self.initCapital  # 资金最高净值

        # 费用情况
        self.avaliable = self.initCapital
        self.percent = EMPTY_FLOAT
        self.percentLimit = 30  # 投资仓位比例上限

        self.maxPnl = 0  # 最高盈利
        self.minPnl = 0  # 最大亏损
        self.maxVolume = 1  # 最大仓位数
        self.winningResult = 0  # 盈利次数
        self.losingResult = 0  # 亏损次数

        self.totalResult = 0  # 总成交数量
        self.totalWinning = 0  # 总盈利
        self.totalLosing = 0  # 总亏损
        self.totalTurnover = 0  # 总成交金额（合约面值）
        self.totalCommission = 0  # 总手续费
        self.totalSlippage = 0  # 总滑点

        self.timeList = []  # 时间序列
        self.pnlList = []  # 每笔盈亏序列
        self.capitalList = []  # 盈亏汇总的时间序列
        self.drawdownList = []  # 回撤的时间序列
        self.drawdownRateList = []  # 最大回撤比例的时间序列

        self.exportTradeList = []  # 导出交易记录列表

        self.fixCommission = EMPTY_FLOAT  # 固定交易费用

    def getAccountInfo(self):
        """返回账号的实时权益，可用资金，仓位比例,投资仓位比例上限"""
        if self.capital == EMPTY_FLOAT:
            self.percent = EMPTY_FLOAT

        return self.capital, self.avaliable, self.percent, self.percentLimit

    # ----------------------------------------------------------------------
    def setStartDate(self, startDate='20100416', initDays=10):
        """设置回测的启动日期"""
        self.dataStartDate = datetime.strptime(startDate, '%Y%m%d')


        # 初始化天数
        initTimeDelta = timedelta(initDays)

        self.strategyStartDate = self.dataStartDate + initTimeDelta

    # ----------------------------------------------------------------------
    def setEndDate(self, endDate=''):
        """设置回测的结束日期"""
        if endDate:
            self.dataEndDate = datetime.strptime(e
ndDate, '%Y%m%d')

        else:
            self.dataEndDate = datetime.now()

    def setMinDiff(self, minDiff):
        """设置回测品种的最小跳价，用于修正数据"""
        self.minDiff = minDiff

    # ----------------------------------------------------------------------
    def setBacktestingMode(self, mode):
        """设置回测模式"""
        self.mode = mode

    # ----------------------------------------------------------------------
    def setDatabase(self, dbName, symbol):
        """设置历史数据所用的数据库"""
        self.dbName = dbName
        self.symbol = symbol

    # ----------------------------------------------------------------------
    def loadHistoryDataFromMongo(self):
        """载入历史数据"""
        host, port, log = loadMongoSetting()

        self.dbClient = pymongo.MongoClient(host, port)
        collection = self.dbClient[self.dbName][self.symbol]

        self.output(u'开始载入数据')

        # 首先根据回测模式，确认要使用的数据类
        if self.mode == self.BAR_MODE:
            dataClass = CtaBarData
            func = self.newBar
        else:
            dataClass = CtaTickData
            func = self.newTick

        # 载入初始化需要用的数据
        flt = {'datetime': {'$gte': self.dataStartDate,
                            '$lt': self.strategyStartDate}}
        initCursor = collection.find(flt)

        # 将数据从查询指针中读取出，并生成列表
        for d in initCursor:
            data = dataClass()
            data.__dict__ = d
            self.initData.append(data)

            # 载入回测数据
        if not self.dataEndDate:
            flt = {'datetime': {'$gte': self.strategyStartDate}}  # 数据过滤条件
        else:
            flt = {'datetime': {'$gte': self.strategyStartDate,
                                '$lte': self.dataEndDate}}
        self.dbCursor = collection.find(flt)

        self.output(u'载入完成，数据量：%s' % (initCursor.count() + self.dbCursor.count()))

    # ----------------------------------------------------------------------
    def connectMysql(self):
        """连接MysqlDB"""

        # 载入json文件
        fileName = 'mysql_connect.json'
        try:
            f = file(fileName)
        except IOError:
            self.writeCtaLog(u'回测引擎读取Mysql_connect.json失败')
            return

        # 解析json文件
        setting = json.load(f)
        try:
            mysql_host = str(setting['host'])
            mysql_port = int(setting['port'])
            mysql_user = str(setting['user'])
            mysql_passwd = str(setting['passwd'])
            mysql_db = str(setting['db'])

        except IOError:
            self.writeCtaLog(u'回测引擎读取Mysql_connect.json,连接配置缺少字段，请检查')
            return

        try:
            self.__mysqlConnection = MySQLdb.connect(host=mysql_host, user=mysql_user,
                                                     passwd=mysql_passwd, db=mysql_db, port=mysql_port)
            self.__mysqlConnected = True
            self.writeCtaLog(u'回测引擎连接MysqlDB成功')
        except Exception:
            self.writeCtaLog(u'回测引擎连接MysqlDB失败')

            # ----------------------------------------------------------------------

    def loadDataHistoryFromMysql(self, symbol, startDate, endDate):
        """载入历史TICK数据
        如果加载过多数据会导致加载失败,间隔不要超过半年
        """

        if not endDate:
            endDate = datetime.today()

        # 看本地缓存是否存在
        if self.__loadDataHistoryFromLoca