import pandas as pd
from pandas import DataFrame
from pathlib import Path
import os
import pickle
import numpy as np
import torch
from sklearn.preprocessing import MinMaxScaler
from tool_data import get_stock_filted
from B_date_tool import calc_date_pre

from A_Global import path_stock, path_relation, path_input

def get_path_file_input(args):
    market      = args['market']
    test_start  = args['test_start']
    test_end    = args['test_end']
    feat_name   = args['feat_name']
    return 'INPUT_' + feat_name + '_' + market + '_' + test_start + '_' + test_end + '.npy'

def get_static_data(date_dict, args, feat_name):
    seq_len                 = 1
    args['seq_len']         = seq_len
    args['date_dict']       = calc_date_pre(date_dict, seq_len)
    args['stocks'], _       = get_stock_filted(args['market'])
    args['feat_name']       = feat_name

    if os.path.exists(path_input / get_path_file_input(args)):
        data_dict = load_x_y(args)
    else:
        data_dict = gen_x_y(args)
    return data_dict

def get_seq_data(date_dict, args, feat_name):
    seq_len                 = args['seq_len']
    args['seq_len']         = seq_len
    args['date_dict']       = calc_date_pre(date_dict, seq_len)
    args['stocks'], _       = get_stock_filted(args['market'])
    args['feat_name']       = feat_name

    if os.path.exists(path_input / get_path_file_input(args)):
        data_dict = load_x_y(args)
    else:
        data_dict = gen_x_y(args)
    return data_dict

# def get_data_adj(args={}):
#
#     market              = args['market']
#     adj_end_date        = args['adj_end_date']
#     # test_end        = args['test_end']
#     threshold_discrete  = args['threshold_discrete']
#     nday_relation       = args['nday_relation']
#     min_support_same    = args['min_support_same']
#     min_support_diff    = args['min_support_diff']
#
#     '''1 Relation Data'''
#     file_adj = name_adj(market             =market,
#                          threshold_discrete =threshold_discrete,
#                          nday_relation      =nday_relation,
#                          min_support        =min_support_same,
#                          min_support_diff   =min_support_diff,
#                          end_date           =adj_end_date
#                          )
#     path_adj = path_relation / market
#
#     if not (path_adj / file_adj).exists():
#         print('FILE NOT EXISTED', path_adj / file_adj)
#         adj = get_adj(market            =market,
#                       threshold_discrete=threshold_discrete,
#                       end_date          =adj_end_date,
#                       nday_relation     =nday_relation,
#                       min_support       =min_support_same,
#                       min_support_diff  =min_support_diff)
#         # np.savetxt(path_adj / file_adj, adj)
#         adj = torch.Tensor(adj)
#         print('Make RELATIONs', str(args['nday_relation']), 'days before', adj_end_date)
#     else:
#         adj = torch.Tensor(np.load(path_adj / file_adj))
#         print('Load RELATIONs', str(args['nday_relation']), 'days before', adj_end_date)
#
#     return adj

def load_x_y(args):
    path_input_file = path_input / get_path_file_input(args)
    print('LOAD INPUT from', path_input_file)
    with open(path_input_file, "rb") as f:
        data_ready = pickle.load(f)
    return data_ready

def gen_x_y(args):
    '''
    date_dict = {
                    'train_a':'2015-11-01'
                    'train_a_index': 1, 'train_b_index': 2,
                    'valid_a_index': 3, 'valid_b_index': 4,
                    'test_a_index' : 5, 'test_b_index' : 6}

    market_period_seq_nf
    path_stock/stocks_[market]/
    '''
    market      = args['market']
    stocks      = args['stocks']
    feat_name   = args['feat_name']
    feat        = args['feat_label'][feat_name]
    label       = args['feat_label']['label']
    seq_len     = args['seq_len']
    date_dict   = args['date_dict']

    data_ready = {}
    ''' 2. seq_date '''
    seq_train = seq_index(stocks, date_dict['date_all'], date_dict['date_train'], seq_len)
    seq_valid = seq_index(stocks, date_dict['date_all'], date_dict['date_valid'], seq_len)
    seq_test  = seq_index(stocks, date_dict['date_all'], date_dict['date_test'],  seq_len)

    print('SEQ Train:', seq_train[0], '-', seq_train[-1])
    print('SEQ Valid:', seq_valid[0], '-', seq_valid[-1])
    print('SEQ Test: ', seq_test[0],  '-', seq_test[-1])

    ''' 3. seq_value 
    1) 归一化
    2) 给seq_date插入值
    '''
    # date_scaler     = date_all[date_dict['train_a_index'] - seq_len: date_dict['train_b_index']]
    # date_experiment = date_all[date_dict['train_a_index'] - seq_len: date_dict['test_b_index']]
    date_scaler     = date_dict['date_pre']
    date_experiment = date_dict['date_pre']

    # if not update and not os.path.exists(src / cache_remark / 'data_ready'):
    #     update = True

    # if update:
    print('Update INPUT to', get_path_file_input(args))

    dfeat, dlabel = scale_stocks_2(market, stocks, date_scaler, date_experiment, feat, [], label)
    # dfeat, dlabel = scale_stocks(stocks, date_scaler, feat, [], label, src=src, dst=path_scaler)
    print('SEQ_VALUE Start..')
    # print(dfeat)
    data_ready['x_train'], data_ready['y_train'] = _seq_value(dfeat, dlabel, seq_train)
    data_ready['x_valid'], data_ready['y_valid'] = _seq_value(dfeat, dlabel, seq_valid)
    data_ready['x_test'],  data_ready['y_test']  = _seq_value(dfeat, dlabel, seq_test)
    print('SEQ_VALUE OVER')

    with open(path_input / get_path_file_input(args), "wb") as f:
        pickle.dump(data_ready, f)

    print('shape - X Y TRAIN', data_ready['x_train'].shape, data_ready['y_train'].shape)
    return data_ready

def seq_index(stocks, date_all, date_label, seq_len):
    '''
    date_label，找前面的seq_len个日期作为序列
    :param stocks:
    :param date_all:
    :param date_seq:
    :param seq_len:
    :return:
    '''
    date_seq = []
    i_start = np.searchsorted(date_all, date_label[0])
    i_end   = np.searchsorted(date_all, date_label[-1])
    for i in range(i_start, i_end + 1):
    #     print(i)
    # exit()
    # for date in date_label:
    #     i = np.searchsorted(date_all, date)
    #     print(i, date, date_all[i])
        for stock in stocks:
            date_seq.append((stock, tuple(date_all[i - seq_len: i + 1])))

    return date_seq

def _seq_value(dfeat, dlabel, seq_list):
    ''' extract x,y from seq_list; replace x(or y) with features(or label)

    :param seq_list: [(stock, (days{seq_len + 1}), ..]
    :return: x - (nstock * ndays, seq_len, nfeat); y - (nstock * ndays)
    '''
    key_now, month_now, year_now = None, None, None
    x, y = [], []
    for seq in seq_list:

        stock = seq[0]
        x_key, y_key = seq[1][:-1], seq[1][-1]

        # if key_now != y_key:
        #     key_now = y_key
        #     if month_now != key_now[5:7]:
        #         month_now = key_now[5:7]
        #         print('SEQ_VALUE', key_now[:7], 'START..')

        x_value = np.array([dfeat.loc[stock + '_' + k].values for k in x_key])
        y_value = dlabel.loc[stock + '_' + y_key]

        x.append(x_value)
        y.append(y_value)

    return torch.Tensor(x), torch.Tensor(y).squeeze()

def scale_stocks_2(market, stocks, days_train, days_experiment, feat_s: list, feat_uns: list, label = ['return']):
    # scalers = {}
    d_feat_label = DataFrame([], columns=feat_s + feat_uns + label)

    print('SCALE Start..')
    i = 0
    for stock in stocks:
        df = pd.read_csv(path_stock / ('stocks_' + market) / stock, index_col=0)
        df_fit = df.loc[days_train]
        df = df.loc[days_experiment]

        if df.isnull().values.any():
            print(stock, 'exist nan')
            print(df)
            exit()

        df.index = stock + '_' + df.index
        df_scaled = df[feat_uns + label]

        if feat_s:
            scaler = MinMaxScaler()
            scaler.fit(df_fit[feat_s])

            df_scaled = pd.concat([df_scaled, DataFrame(scaler.transform(df[feat_s]), index=df.index, columns=feat_s)], axis=1)

        d_feat_label = pd.concat([d_feat_label, df_scaled], sort=False)
        i += 1

    print('SCALE', i, 'Stocks OVER')
    # d_feat_label.to_csv(dst / 'ALL_SCALED.csv')

    dfeat   = d_feat_label[feat_s + feat_uns]
    dlabel  = d_feat_label[label]
    return dfeat, dlabel

if __name__ == '__main__':
    # adj_end_date = '2018-03-15'
    # args = {
    #     'market': 'SH',
    #     'adj_end_date': adj_end_date,
    #     'threshold_discrete': 0.005,
    #     'nday_relation': 30,
    #     'min_support_same': 0.5,
    #     'min_support_diff': 0.4,
    #
    # }
    # adj1 = get_data_adj(args)
    #
    # args['threshold_discrete'] = 0.008
    # adj2 = get_data_adj(args)
    #
    # print(adj1 == adj2)

    from A_Global import args
    from B_date_tool import gen_date
    args['market'] = 'NYSE'

    date_dict = gen_date(args)
    get_seq_data(date_dict, args, 'feat_seq')