import numpy as np
from pandas import Series
from A_Global import path_date, markets, args

def gen_date(args):
    if('proportion' in args):
        return gen_date_proportion(args)
    if all(_ in args for _ in ['train_start', 'train_end', 'test_start', 'test_end']):
        return gen_date_range(args)

def gen_date_range(args):
    train_start     = args['train_start']
    train_end       = args['train_end']
    valid_start     = args['valid_start']
    valid_end       = args['valid_end']
    test_start      = args['test_start']
    test_end        = args['test_end']
    market          = args['market']

    date_all = np.genfromtxt(path_date / ("TRADE_DAYs_" + market.upper() + ".txt"), dtype=str)
    date_all = Series(index=date_all, dtype=object)

    date_train  = date_all[train_start: train_end]
    date_valid  = date_all[valid_start: valid_end]
    date_test   = date_all[test_start : test_end]

    date_dict = {
        'market'        : market,
        'date_all'      : series_index_list(date_all),
        'date_train'    : series_index_list(date_train),
        'date_valid'    : series_index_list(date_valid),
        'date_test'     : series_index_list(date_test),
        'date_experiment':series_index_list(date_all[train_start: test_end])
    }
    print_date(date_dict)

    return date_dict

def calc_date_pre(date_dict, seq_len = 5):
    date_all                = date_dict['date_all']
    date_experiment         = date_dict['date_experiment']
    train_start_index       = np.searchsorted(date_all, date_experiment[0])
    train_end_index         = np.searchsorted(date_all, date_experiment[-1])
    date_pre                = date_all[train_start_index - seq_len: train_end_index + 1]
    # print(date_experiment[0], date_experiment[-1])
    # print(date_pre[0], date_pre[-1])
    # print(date_pre[:5])

    date_dict['date_pre']   = date_pre
    return date_dict

def series_index_list(s):
    return list(s.index)

def gen_date_proportion(args):
    '''
    确定日期范围和划分 train,valid,test
    '''
    test_start      = args['test_start']
    test_end        = args['test_end']
    proportion      = args['proportion']
    market          = args['market']

    date_all = np.genfromtxt(path_date / ("TRADE_DAYs_"+market.upper()+".txt"), dtype=str)

    ''' 1. gen date '''
    test_start_index = np.searchsorted(date_all, test_start)
    test_end_index   = np.searchsorted(date_all, test_end)
    unit = (test_end_index - test_start_index) // proportion[-1]
    valid_end_index  = test_start_index
    valid_start_index= valid_end_index - unit * proportion[1]
    train_end_index  = valid_start_index
    train_start_index= train_end_index - unit * proportion[0]

    date_dict = {
        'market'        : market,
        'date_all'      : date_all,
        'train_a_index' :train_start_index, 'train_b_index' :train_end_index,
        'valid_a_index' :valid_start_index, 'valid_b_index' :valid_end_index,
        'test_a_index'  :test_start_index,  'test_b_index'  :test_end_index,
        'train_a'       :date_all[train_start_index]
    }

    date_train, date_valid, date_test = \
        date_all[date_dict['train_a_index']: date_dict['train_b_index']],\
        date_all[date_dict['valid_a_index']: date_dict['valid_b_index']],\
        date_all[date_dict['test_a_index'] : date_dict['test_b_index']]


    return date_dict

def print_date(date_dict):
    market          = date_dict['market']
    date_train      = date_dict['date_train']
    date_valid      = date_dict['date_valid']
    date_test       = date_dict['date_test']

    print(market, '- DATE EXPERIMENT: From', convert_time_str(date_train[0]), 'To', convert_time_str(date_test[-1]), len(date_train) + len(date_valid) + len(date_test))
    print('%3d TRAIN DAYs From %s To %s' % (len(date_train), convert_time_str(date_train[0]), convert_time_str(date_train[-1])))
    print('%3d VALID DAYs From %s To %s' % (len(date_valid), convert_time_str(date_valid[0]), convert_time_str(date_valid[-1])))
    print('%3d  TEST DAYs From %s To %s' % (len(date_test),  convert_time_str(date_test[0]),  convert_time_str(date_test[-1])))
    print()

def convert_time_str(timestr):
    timearr = timestr.split('-')
    return timearr[1] + '/' + timearr[2] + '/' + timearr[0]


if __name__ == '__main__':
    # use for paper 3.1 数据收集
    # note: need adjustment mannually
    market          = 'NYSE'
    args['market']  = market

    date_dict       = gen_date(args)
    date_experiment = date_dict['date_experiment']
    print(len(date_experiment))

    date_dict       = calc_date_pre(date_dict, seq_len=5)

    ''' SAVE '''
    # np.savetxt(path_date / ('EXPERIMENT_DAYs_' + market.upper() + '.txt'), date_experiment, fmt='%s')
    # np.savetxt(path_date / ('EXPERIMENT_DAYs_TRAIN_' + market.upper() + '.txt'), date_dict['date_train'], fmt='%s')


