# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import talib
import config
import os
import logging
from distutils.dir_util import mkpath


class StockIndex(object):
    name = None
    stock = None  # 指数代号
    full_period = (None, None)  # 全周期，起始日、中止日
    period = {  # dict{}
    }  # key是周期名字，value是[(s0, e0), (s1, e1), (s2, e2), ...]

    def __init__(self):
        self.full_period_df = None
        self.load()

    def _data_filename(self):
        return ('%s_OHLC_%s_%s.csv' % (self.stock, \
                                       self.full_period[0].replace('-', ''), \
                                       self.full_period[1].replace('-', ''))).replace('.', '')

    def load(self):
        path = os.path.join(config.DATA_ROOT, 'stock_index', self._data_filename())
        if not os.path.exists(path):
            self.full_period_df = self._download_full_period_df()
            mkpath(os.path.join(config.DATA_ROOT, 'stock_index'))
            self.full_period_df.to_csv(path, float_format='%.02f')
        else:
            self.full_period_df = pd.read_csv(path, index_col='date', parse_dates=True)

    def _download_full_period_df(self):
        '''
        下载全周期数据，返回 pd.DataFrame
        '''
        from WindPy import w
        w.start()
        w_data = w.wsd(self.stock, 'open,high,low,close', self.full_period[0], self.full_period[1], '')
        w.stop()
        if w_data.ErrorCode != 0:
            raise Exception('WindData Error')
        df = pd.DataFrame(
            data={
                'open': w_data.Data[0],
                'high': w_data.Data[1],
                'low': w_data.Data[2],
                'close': w_data.Data[3],
            },
            columns=['open', 'high', 'low', 'close'],
            index=pd.DatetimeIndex(w_data.Times, name='date')
        )
        return df


class SH180(StockIndex):
    name = '上证180'
    stock = '000010.SH'
    full_period = ('2004-01-01', '2017-06-12')
    period = {
        'positive': [('2014-06-01', '2015-06-11')],
        'negative': [('2015-06-18', '2016-03-02')],
    }


class SH000001(StockIndex):
    name = '上证指数'
    stock = '000001.SH'
    full_period = ('2004-01-01', '2017-06-12')
    period = {
        'positive': [('2014-06-01', '2015-06-11')],
        'negative': [('2015-06-18', '2016-03-02')],
    }
