# -*- coding:utf-8 -*- 
"""
pro init 
Created on 2018/07/01
@author: Jimmy Liu
@group : https://tushare.pro
@contact: jimmysoa@sina.cn
"""
from __future__  import division
import datetime
from tushare.pro import client
from tushare.util import upass
from tushare.util.formula import MA

PRICE_COLS = ['open', 'close', 'high', 'low', 'pre_close']
FORMAT = lambda x: '%.4f' % x
FREQS = {'D': '1DAY',
         'W': '1WEEK',
         'Y': '1YEAR',
         }
FACT_LIST = {
           'tor': 'turnover_rate',
           'turnover_rate': 'turnover_rate',
           'vr': 'volume_ratio',
           'volume_ratio': 'volume_ratio',
           'pe': 'pe',
           'pe_ttm': 'pe_ttm',
        }

def pro_api(token=''):
    """
    初始化pro API,第一次可以通过ts.set_token('your token')来记录自己的token凭证，临时token可以通过本参数传入
    """
    if token == '' or token is None:
        token = upass.get_token()
    if token is not None and token != '':
        pro = client.DataApi(token)
        return pro
    else:
        raise Exception('api init error.') 
        

def pro_bar(ts_code='', api=None, start_date='', end_date='', freq='D', asset='E', 
           exchange='',
           adj = None,
           ma = [],
           factors = None,
           adjfactor = False,
           contract_type = '',
           retry_count = 3):
    """
    BAR数据
    Parameters:
    ------------
    ts_code:证券代码，支持股票,ETF/LOF,期货/期权,港股,数字货币
    start_date:开始日期  YYYYMMDD
    end_date:结束日期 YYYYMMDD
    freq:支持1/5/15/30/60分钟,周/月/季/年
    asset:证券类型 E:股票和交易所基金，I:沪深指数,C:数字货币,FT:期货 FD:基金/O期权/H港股/CB可转债
    exchange:市场代码，用户数字货币行情
    adj:复权类型,None不复权,qfq:前复权,hfq:后复权
    ma:均线,支持自定义均线频度，如：ma5/ma10/ma20/ma60/maN
    factors因子数据，目前支持以下两种：
        vr:量比,默认不返回，返回需指定：factor=['vr']
        tor:换手率，默认不返回，返回需指定：factor=['tor']
                    以上两种都需要：factor=['vr', 'tor']
    retry_count:网络重试次数
    
    Return
    ----------
    DataFrame
    code:代码
    open：开盘close/high/low/vol成交量/amount成交额/maN均价/vr量比/tor换手率
    
         期货(asset='FT')
    code/open/close/high/low/avg_price：均价  position：持仓量  vol：成交总量
    """
    today= datetime.datetime.today().date()
    today = str(today)[0:10]
    start_date = '' if start_date is None else start_date
    end_date = today if end_date == '' or end_date is None else end_date
    ts_code = ts_code.strip().upper() if asset != 'C' else ts_code.strip().lower()
    start_date = start_date.replace('-', '')
    end_date = end_date.replace('-', '')
    if len(freq.strip())>=3:
        freq = freq.strip().lower()
    else:
        freq = freq.strip().upper() if asset != 'C' else freq.strip().lower()
    asset = asset.strip().upper()
    api = api if api is not None else pro_api()
    for _ in range(retry_count):
        try:
            if asset == 'E':
                if freq == 'D':
                    data = api.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
                    if factors is not None and len(factors) >0 :
                        ds = api.daily_basic(ts_code=ts_code, start_date=start_date, end_date=end_date)[['trade_date', 'turnover_rate', 'volume_ratio']]
                        ds = ds.set_index('trade_date')
                        data = data.set_index('trade_date')
                        data = data.merge(ds, left_index=True, right_index=True)
                        data = data.reset_index()
                        if ('tor' in factors) and ('vr' not in factors):
                            data = data.drop('volume_ratio', axis=1)
                        if ('vr' in factors) and ('tor' not in factors):
                            data = data.drop('turnover_rate', axis=1)
                if freq == 'W':
                    data = api.weekly(ts_code=ts_code, start_date=start_date, end_date=end_date)
                if freq == 'M':
                    data = api.monthly(ts_code=ts_code, start_date=start_date, end_date=end_date)
                if 'min' in freq:
                    data = api.mins(ts_code=ts_code, start_time=start_date, end_time=end_date, freq=freq)
                    data['trade_date'] = data['trade_time'].map(lambda x: x.replace('-', '')[0:8])
                    data['pre_close'] = data['close'].shift(-1)
                if adj is not None:
                    fcts = api.adj_factor(ts_code=ts_code, start_date=start_date, end_date=end_date)[['trade_date', 'adj_factor']]
                    if fcts.shape[0] == 0:
                        return None
                    data = data.set_index('trade_date', drop=False).merge(fcts.set_index('trade_date'), left_index=True, right_index=True, how='left')
                    if 'min' in freq:
                        data = data.sort_values('trade_time', ascending=False)
                    data['adj_factor'] = data['adj_factor'].fillna(method='bfill')
                    for col in PRICE_COLS:
                        if adj == 'hfq':
                            data[col] = data[col] * data['adj_factor']
                        if adj == 'qfq':
                            data[col] = data[col] * data['adj_factor'] / float(fcts['adj_factor'][0])
                        data[col] = data[col].map(FORMAT)
                        data[col] = data[col].astype(float)
                    if adjfactor is False:
                        data = data.drop('adj_factor', axis=1)
                    if 'min' not in freq:
                        data['change'] = data['close'] - data['pre_close']
                        data['pct_chg'] = data['change'] / data['pre_close'] * 100
                        data['pct_chg'] = data['pct_chg'].map(lambda x: FORMAT(x)).astype(float)
                    else:
                        data = data.drop(['trade_date', 'pre_close'], axis=1)
            elif asset == 'I':
                if freq == 'D':
                    data = api.index_daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
                if freq == 'W':
                    data = api.index_weekly(ts_code=ts_code, start_date=start_date, end_date=end_date)
                if freq == 'M':
                    data = api.index_monthly(ts_code=ts_code, start_date=start_date, end_date=end_date)
                if 'min' in freq:
                    data = api.mins(ts_code=ts_code, start_time=start_date, end_time=end_date, freq=freq)
            elif asset == 'FT':
                if freq == 'D':
                    data = api.fut_daily(ts_code=ts_code, start_date=start_date, end_date=end_date, exchange=exchange)
                if 'min' in freq:
                    data = api.mins(ts_code=ts_code, start_time=start_date, end_time=end_date, freq=freq)
            elif asset == 'O':
                if freq == 'D':
                    data = api.opt_daily(ts_code=ts_code, start_date=start_date, end_date=end_date, exchange=exchange)
                if 'min' in freq:
                    data = api.mins(ts_code=ts_code, start_time=start_date, end_time=end_date, freq=freq)
            elif asset == 'CB':
                if freq == 'D':
                    data = api.cb_daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
            elif asset == 'FD':
                if freq == 'D':
                    data = api.fund_daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
                if 'min' in freq:
                    data = api.mins(ts_code=ts_code, start_time=start_date, end_time=end_date, freq=freq)
            if asset == 'C':
                if freq == 'd':
                    freq = 'daily'
                elif freq == 'w':
                    freq = 'week'
                data = api.coinbar(exchange=exchange, symbol=ts_code, freq=freq, start_dae=start_date, end_date=end_date,
                                   contract_type=contract_type)
            if ma is not None and len(ma) > 0:
                for a in ma:
                    if isinstance(a, int):
                        data['ma%s'%a] = MA(data['close'], a).map(FORMAT).shift(-(a-1))
                        data['ma%s'%a] = data['ma%s'%a].astype(float)
                        data['ma_v_%s'%a] = MA(data['vol'], a).map(FORMAT).shift(-(a-1))
                        data['ma_v_%s'%a] = data['ma_v_%s'%a].astype(float)
            data = data.reset_index(drop=True)
        except Exception as e:
            print(e)
        else:
            return data
    raise IOError('ERROR.')


if __name__ == '__main__':
#     upass.set_token('your token here')
    pro = pro_api()
#     print(pro_bar(ts_code='000001.SZ', api=pro, start_date='20180101', end_date='20181101', factors=['tor', 'vr']))
#     print(pro_bar(ts_code='000001.SH', api=pro, start_date='20180101', end_date='20181011', asset='I', freq='60min'))
#     print(pro_bar(ts_code='000905.SH', start_date='20181001', end_date='', asset='I'))
#     print(pro.trade_cal(exchange_id='', start_date='20131031', end_date='', fields='pretrade_date', is_open='0'))
#     print(pro_bar(ts_code='CU1811.SHF', api=pro, start_date='20180101', end_date='', asset='FT', ma=[5, 10, 15]))
#     print(pro_bar(ts_code='150023.SZ', api=pro, start_date='20180101', end_date='', asset='FD', ma=[5, 10, 15]))
#     print(pro_bar(api=pro, ts_code='000528.SZ',start_date='20180101', end_date='20181121', ma=[20]))
#     df = pro_bar(ts_code='000001.SH', freq='1min', asset='I', start_date='20190131', end_date='20190413')
#     df = pro_bar(ts_code='000001.SZ', start_date='20190415', end_date='20190416', asset='E', factors=['tor', 'vr'])
#     df = pro_bar(ts_code='CU1904.SHF', asset='FT', freq='d', start_date='20181001', end_date='20190101', adj='')
#     print(pro_bar(ts_code='000528.SZ', api=pro, freq='M', start_date='20180101', end_date='20180820', adj='qfq', ma=[5, 10, 15]))
#     print(pro_bar(ts_code='btcusdt', api=pro, exchange='huobi', freq='15MIN', start_date='20180101', end_date='20191123', asset='C', ma=[5, 10]))
#     df = pro_bar(ts_code='300510.SZ', api=pro, adj='', start_date='20180601', end_date='20181205', freq='60min')
#     df = pro_bar(ts_code='300510.SZ', api=pro, adj='qfq', start_date='20100601', end_date='20181205', freq='60min', ma=[5, 10, 20])
#     df = pro_bar(ts_code='600862.SH', api=pro, start_date='20150118', end_date='20150615', factors=['tor', 'vr'])
#     df = pro_bar(api=pro, ts_code='btcusdt', start_date='20180628', end_date='20180928', freq='15min')
#     df = pro_bar(ts_code='601888.SH', adj='qfq', freq='D')
#     df = pro_bar(api=pro, ts_code='300527.SZ', start_date='20190702', end_date='', adj='qfq')
#     df = pro_bar(ts_code='300671.SZ', factors=['tor'])
    df = pro_bar(ts_code='000001.SZ',  start_date='20190101', end_date='2019301', ma=['5, 20, 50'])
    print(df)
    