# -*- coding: utf-8 -*-
import dateutil
import datetime
import logging
import pickle
import hashlib
import os
import config
import TSLPy2
import pandas as pd
import numpy as np
import atexit
import zmq
from multiprocessing import Pool
from html_utils import TemplateRender
import re
from multiprocessing import Pool
from jinja2 import Environment, FileSystemLoader, select_autoescape

# JINJA2_ENV = Environment(
#     loader=FileSystemLoader('./html'),
#     autoescape=select_autoescape(['html', 'xml'])
# )
#
# PLOT_TEMPLATE = JINJA2_ENV.get_template('template.html')

# 天软 Broker 相关，切勿删除
PROXY_CLIENT_SOCKET = None
PROXY_CLIENT_CONTEXT = None

PERIOD_HQ_DIR_NAME = 'period_hq_2010-01-01_2017-06-01'


def str_to_ts_date(date_str):
    d = dateutil.parser.parse(date_str)
    return TSLPy2.EncodeDate(d.year, d.month, d.day)


def str_to_ts_datetime(date_str):
    d = dateutil.parser.parse(date_str)
    return TSLPy2.EncodeDateTime(d.year, d.month, d.day, d.hour, d.minute, d.second, 0)


def str_gbk_to_utf8(s):
    return s.decode('gbk').encode('utf8')


def str_utf8_to_gbk(s):
    return s.decode('utf8').encode('gbk')


def convert_ts_dict_with_to_utf8(d):
    rs = {}
    for k, v in d.iteritems():
        rs[str_gbk_to_utf8(k)] = v
    return rs


def ts_login_if_not_yet(ts_username, ts_password):
    if TSLPy2.Logined():
        logging.info('已经登录天软')
    elif not TSLPy2.Logined():
        logging.info('开始登录天软')
        TSLPy2.ConnectServer(config.TSSERVER, config.TSPORT)
        re = TSLPy2.LoginServer(ts_username, ts_password)
        if re[0] == 0:
            logging.info('登录天软成功')
        else:
            logging.info('登录天软失败: ' + re[1])
            raise Exception("Login Tinysoft failed.")


@atexit.register
def close_proxy_client_context():
    global PROXY_CLIENT_SOCKET
    global PROXY_CLIENT_CONTEXT
    if PROXY_CLIENT_SOCKET is not None:
        logging.debug("断开数据服务连接")
        PROXY_CLIENT_SOCKET.close()
        PROXY_CLIENT_CONTEXT.term()


def ts_proxy_remote_call_func(*params):
    global PROXY_CLIENT_CONTEXT
    global PROXY_CLIENT_SOCKET
    if PROXY_CLIENT_CONTEXT is None:
        logging.debug("连接数据服务")
        PROXY_CLIENT_CONTEXT = zmq.Context()
        PROXY_CLIENT_SOCKET = PROXY_CLIENT_CONTEXT.socket(zmq.DEALER)
        # PROXY_CLIENT_SOCKET.setsockopt(zmq.RCVTIMEO, 5000)  # 5 seconds timeout
        PROXY_CLIENT_SOCKET.connect('tcp://%s:%d' % (config.TS_BROKER_SERVER, config.TS_BROKER_PORT))

    PROXY_CLIENT_SOCKET.send_string(pickle.dumps(params))
    data_text = PROXY_CLIENT_SOCKET.recv()
    return pickle.loads(data_text)


def ts_remote_call_func(*params, **kwargs):
    check = str(params)
    hashcode = hashlib.sha256(check).hexdigest()
    if not os.path.exists(config.TS_FILE_PATH):  # 检测文件夹并按需创建
        os.makedirs(config.TS_FILE_PATH)
    file_name = config.TS_FILE_PATH + str(hashcode)

    if os.path.exists(file_name):
        logging.debug('文件数据 ' + file_name)
        f = open(file_name, 'rb')
        check_str, data = pickle.load(f)
        f.close()
        # 文件读取的验证字段必须与入参的验证字段一致
        assert check_str == check
    else:
        logging.debug('网络数据 ' + file_name)

        if config.TS_USER_BROKER:
            data = ts_proxy_remote_call_func(*params)  # 原有的天软调用
        else:
            ts_login_if_not_yet(config.TSUSER, config.TSPWD)
            data = TSLPy2.RemoteCallFunc(*params)  # 原有的天软调用

        # 只保存正确的返回，错误的返回（例如网络异常、没有登录）则不保存
        if data[0] == 0:
            to_save = (check, data)
            f = open(file_name, 'wb')
            pickle.dump(to_save, f, 2)
            f.close()

    # 错误检查
    if data[0] != 0:
        raise Exception(data[2].decode('gbk').encode('utf8'))

    return data[1]


def get_ts_index_stock_code_list(index_str, date_str=None):
    """
    返回指数的成份股列表
    :param index_str:           天软指数/板块代码，默认 
    :param date_str:            查询日期，默认 datetime.date.today().isoformat()
    :return:  list<string>
     
    """
    ts_date = str_to_ts_date(date_str if date_str is not None else datetime.date.today().isoformat())
    return ts_remote_call_func("GetBKByDate", [str_utf8_to_gbk(index_str), ts_date], {})


def ts_list_to_dataframe(ts_list, columns, index_type=pd.Index):
    """
    把天软的查询返回值，转换成 pandas.DataFrame
    :param ts_list:       list
        天软的查询list返回
    :param columns:     list((column_name_str, ts_column_name_str))
        定义 DataFrame 中的 Index, Column 顺序，及其对应的天软名称翻译。
        特别地，第一个元素，会作为DataFrame的Index
    :param index_type:      pd.Index 及其派生类
        如： pd.Index, pd.DateTimeIndex
        
    调用例子    
    ts_list_to_dataframe(ts_list, 
                         [('datetime', '日期"), ('open', '开盘价'), ('close', '收盘价')], 
                         index_type=pd.DateTimeIndex])
        
    :return:                pd.DataFrame
    """
    assert type(ts_list) is list
    converted_list = map(convert_ts_dict_with_to_utf8, ts_list)

    return pd.DataFrame(
        # data=[[row[col[1]] for col in columns[1:]] for row in converted_list],
        data=[[row.get(col[1], np.nan) for col in columns[1:]] for row in converted_list],
        columns=[col[0] for col in columns[1:]],
        index=index_type(map(lambda d: d[columns[0][1]], converted_list), name=columns[0][0])
    )


def get_stock_pool(format=None):
    """
    获取  ./data/stock_pool.csv 中的股票代码列表，用天软格式
    
    :param format:  str
        
    :return: list<str> 
    """
    stock_list = pd.read_csv('./data/stock_pool.csv')
    stock_list = stock_list['stock']
    stock_list = [str(s).zfill(6) for s in stock_list]
    stock_list = [('SZ' if s[0] in ['0', '3'] else 'SH') + s for s in stock_list]
    return map(lambda x: stock_format(x, format), stock_list)


def calc_limit_up_price(price):
    """
    计算涨停价
    :param price: float
    :return: float
    """
    return round(price * 1.1 + 0.0000001, 2)  # 补精度损失


def parallel_run(target, run_data_list, process_num=1):
    if process_num == 1:
        return map(target, run_data_list)
    # 并行处理，结果乱序返回，需要缓存
    pool = Pool(processes=process_num)
    rs_run_data_dict = {}
    for finished_cnt, rs_run_data in enumerate(pool.imap_unordered(target, run_data_list)):
        rs_run_data_dict[rs_run_data.stock_id] = rs_run_data
        print "Finish %d / %d" % (finished_cnt + 1, len(run_data_list))

    # 用缓存替换原有的数据
    for idx in range(len(run_data_list)):
        run_data_list[idx] = rs_run_data_dict[run_data_list[idx].stock_id]


class WbSheetAppender(object):
    """
    openpyxl Workbook.Sheet 辅助类
    """

    def __init__(self, sheet, start_row=0):
        self.start_row = start_row
        self.sheet = sheet

    def write_row(self, *args):
        """
        按行写一个单行
        :param args:    unicode str
        :return: 
        """
        for col_idx, text in enumerate(args):
            self.sheet.cell(row=self.start_row + 1, column=col_idx + 1, value=text)
        self.start_row += 1


def plot_html(stock_id,
              datetime_ohlc_df,
              order_history_df,
              cash_value_df,
              cerebro,
              index_format=lambda index: index.strftime('%Y-%m-%d'),
              linebuf_sample_step=1,
              show=True):
    """
    :param stock_id:            str
        股票代码
    :param datetime_ohlc_df:    Open High Low Close DataFrame
        index为datetime，存储 开盘价、最高价、最低价、收盘价 的 DataFrame
    :param order_history_df: 
        RunData 中存储的交易记录
    :param cerebro:             backtrader.Cerebro
        BackTrader Cerebro对象，用于提取cash/value曲线
    :param linebuf_sample_step: int
        cash/value 曲线的取值间隔，默认是1。
        如回测使用15分钟线，但绘制使用日线，则可调整step为16.
    :param index_format:
        对 datetime_ohlc_df 的 index 的日期处理函数
    :return:                    None
    """
    output_html_path = './html/output.html'
    tr = TemplateRender(title=stock_id,
                        ohlc_df=datetime_ohlc_df,
                        order_history_df=order_history_df)
    tr.add_indicator('cash/value', cash_value_df)

    tr.save(output_html_path)

    if show:
        os.system("\"%s\" %s" % (config.CHROME_PATH, output_html_path))


def update_trade_date_list_from_wind():
    from WindPy import w
    w.start()
    with open(os.path.join(config.DATA_ROOT, 'trade_date_list'), 'w+') as f:
        f.writelines([dt.strftime('%Y-%m-%d\n') for dt in w.tdays("2007-01-1", "", "").Times])


trade_date_str_list = []


def get_trade_date_str_list(start_date=None, end_date=None, month=None):
    """
    :param start_date:  开始日期
    :param end_date:    结束日期（包括当天）
    :param month:
    :return:
    """
    global trade_date_str_list
    if not trade_date_str_list:
        with open(os.path.join(config.DATA_ROOT, 'trade_date_list'), 'r') as f:
            trade_date_str_list = f.readlines()
            trade_date_str_list = map(str.strip, trade_date_str_list)
    ret_list = trade_date_str_list[:]
    if month:
        return filter(lambda x: x.startswith(month), ret_list)
    if start_date:
        ret_list = filter(lambda x: x >= start_date, ret_list)
    if end_date:
        ret_list = filter(lambda x: x <= end_date, ret_list)
    return ret_list


class StockFormat(object):
    WIND = 'wind'
    TINYSOFT = 'tinysoft'
    TRADESTATION = 'tradestation'
    RAW = 'raw'


def stock_format(stock, fmt='raw'):
    raw_stock = stock
    if isinstance(raw_stock, int) or isinstance(raw_stock, long):
        raw_stock = str(raw_stock)
    raw_stock = re.search('[0-9]+', raw_stock)
    if not raw_stock:
        raise Exception("bad stock format")
    raw_stock = raw_stock.group(0)
    raw_stock = raw_stock.zfill(6)
    if fmt == StockFormat.RAW:
        return raw_stock
    if raw_stock.startswith('6'):
        location_code = 'SH'
    elif raw_stock.startswith('0') or raw_stock.startswith('30'):
        location_code = 'SZ'
    else:
        raise Exception('bad stock format \'%s\'', raw_stock)
    if fmt.lower() == StockFormat.TINYSOFT.lower():
        return location_code + raw_stock
    if fmt.lower() == StockFormat.TRADESTATION.lower():
        return raw_stock + '.' + location_code
    if fmt.lower() == StockFormat.WIND.lower():
        return raw_stock + '.' + location_code
    raise 'bad stock format'


def get_all_stock():
    return os.listdir(os.path.join(config.DATA_ROOT, PERIOD_HQ_DIR_NAME))


def load_period_hq_df(stock, begin_date=None, end_date=None,
                      columns=[('date', '日期'), ('open', '开盘'), ('close', '收盘'), ('high', '最高'),
                               ('low', '最低'), ('p_chg_pct', '涨幅(%)'), ('volume', '成交量(手)'),
                               ('turnover_chg_pct', '换手率(%)')]):
    """
    Inclusive
    """
    # check begin_date, end_date are in
    cache_dir_path = os.path.join(config.DATA_ROOT, 'period_hq_2010-01-01_2017-06-01')
    ts_stock = stock_format(stock, StockFormat.TINYSOFT)
    df = pd.read_csv(os.path.join(cache_dir_path, ts_stock), index_col='日期')
    df.index = pd.DatetimeIndex(df.index.tolist())
    df.index.name = columns[0][0]
    col_dict = {}
    for col in columns[1:]:
        col_dict[col[1]] = col[0]
    if begin_date is not None:
        assert begin_date >= '2010-01-01'
        df = df[df.index >= begin_date]
    if end_date is not None:
        assert end_date <= '2017-06-01'
        df = df[df.index <= end_date]
    return df.rename(columns=col_dict)


def update_local_period_hq_csv_cache_for_all_stock(begin_date='2010-01-01', end_date='2017-06-01'):
    from WindPy import w
    w.start()
    w_data = w.wset("sectorconstituent", "date=2017-06-29;sectorid=a001010100000000")
    if w_data.ErrorCode != 0:
        print w_data.ErrorCode
        raise Exception('Wind Error')
    for idx, wind_stock in enumerate(w_data.Data[1]):
        print 'fetching {} {}/{}'.format(wind_stock, idx + 1, len(w_data.Data[1]))
        update_local_period_hq_csv_cache_for_single_stock(wind_stock, begin_date, end_date)


def update_local_period_hq_csv_cache_for_single_stock(stock, begin_date='2010-01-01', end_date='2017-06-01'):
    cache_dir_path = os.path.join(config.DATA_ROOT, 'period_hq_{}_{}'.format(begin_date, end_date))
    if not os.path.exists(cache_dir_path):
        os.mkdir(cache_dir_path)
    ts_stock = stock_format(stock, StockFormat.TINYSOFT)

    if os.path.exists(os.path.join(cache_dir_path, ts_stock)):
        print '%s exists' % ts_stock
        return

    # 交易所算法，前复权
    stock_daily_data = ts_remote_call_func('Stock_PeriodHq',
                                           [str_to_ts_date(begin_date), str_to_ts_date(end_date), '日', 2, 0], {
                                               'StockID': ts_stock
                                           })
    row_size = len(stock_daily_data)
    if row_size == 0:
        return None
    stock_daily_data.reverse()  # 天软查询接口是倒序，转换成按照时间先后排序，日期小的排前面
    ts_header_columns = ['日期', '开盘', '收盘', '最高', '最低', '成交量(手)', '成交金额(万)', '涨幅(%)',
                         '超额收益率(%)', '振幅(%)', '换手率(%)', '流通市值(万)', '市价总值(万)',
                         '市盈率(动态)', '市净率(动态)', '市销率(动态)', '市现率(动态)']
    df = ts_list_to_dataframe(
        ts_list=stock_daily_data,
        columns=zip(ts_header_columns, ts_header_columns),
        index_type=pd.DatetimeIndex
    )
    df.to_csv(os.path.join(cache_dir_path, ts_stock))
    return df



if __name__ == '__main__':
    update_local_period_hq_csv_cache_for_all_stock(begin_date='2010-01-01', end_date='2017-07-16')
    # update_local_period_hq_csv_cache_for_single_stock('SH000905', begin_date='2010-01-01', end_date='2017-07-16')
    # update_local_period_hq_csv_cache_for_single_stock('SH000016', begin_date='2010-01-01', end_date='2017-07-16')