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

import math
import backtrader as bt
import t1_short.t1_short_runner
import pandas as pd
import numpy as np
import talib
from utils import stock_format, StockFormat, ts_remote_call_func, str_to_ts_date, ts_list_to_dataframe
import utils
import config
import os
import logging


def prepare_trade_detail_for_single_stock(stock):
    ts_stock = stock_format(stock, StockFormat.TINYSOFT)
    trade_detail_dir = 't1_sell_short_v5'
    if not os.path.exists(os.path.join(config.DATA_ROOT, trade_detail_dir)):
        os.mkdir(os.path.join(config.DATA_ROOT, trade_detail_dir))
    if os.path.exists(os.path.join(config.DATA_ROOT, trade_detail_dir, ts_stock)):
        print '{} already loaded. Skip.'.format(ts_stock)
        return
    ts_stock = stock_format(stock, StockFormat.TINYSOFT)

    trade_detail_list = ts_remote_call_func('Last10mTradeDetails',
                                            [ts_stock,
                                             str_to_ts_date('2010-01-01'),
                                             str_to_ts_date('2017-06-01'),
                                             '10分钟线'.decode('utf8').encode('gbk')], {})
    row_size = len(trade_detail_list)
    if row_size == 0:
        return

    headers = ['datetime', 'open', 'high', 'low', 'close', 'sectional_open']
    ts_list_to_dataframe(
        ts_list=trade_detail_list,
        columns=zip(headers, headers),
        index_type=pd.DatetimeIndex
    ).round(2).to_csv(os.path.join(config.DATA_ROOT, trade_detail_dir, ts_stock))


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

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


def prepare_full_period_df_func(df, stock):
    ts_stock = utils.stock_format(stock, StockFormat.TINYSOFT)
    last_10m_df = pd.read_csv(os.path.join(config.DATA_ROOT, 't1_sell_short_v5', ts_stock), index_col='datetime')
    last_10m_df.index = pd.DatetimeIndex(pd.DatetimeIndex(last_10m_df.index).date)
    last_10m_raise_pct = (last_10m_df.high - last_10m_df.open) / last_10m_df.sectional_open * 100
    for dt in df.index:
        date = dt.date()
        df.loc[dt, 'sig'] = 0
        if date not in last_10m_raise_pct.index:
            continue
        if last_10m_raise_pct.get(date, 0) > 1:
            df.loc[dt, 'sig'] = 1
    return df


class Strategy(bt.Strategy):
    def __init__(self, p_chg_pct=9.8):
        super(Strategy, self).__init__()
        self.last_buy_date = None
        self.p_chg_pct = p_chg_pct

    def next_open(self):
        if not self.position:
            if len(self.data0) == self.data0.buflen():
                return
            if math.isnan(self.data0.close[-1]):
                return
            if self.data.sig[0] != 0:
                # and self.data.p_change[0] < 0:
                self.sell(data=self.data0,
                          size=int(self.broker.cash / self.data0.close[0]),
                          price=self.data0.close[0],
                          exectype=bt.Order.StopLimit)
        else:
            self.close(price=self.data0.open[0], exectype=bt.Order.Limit)
