from pathlib import Path
import numpy as np
import torch.nn.functional as F
import os
cur_path = Path(os.path.dirname(__file__))


# path
path_dat        = Path(cur_path / './Data/dat')
path_stock      = Path(cur_path / './Data/dat/stocks')
path_index      = path_stock / 'stocks_index'
path_stock_tmp  = path_dat / 'tmp'
path_date       = Path(cur_path / './Data/dat/date')
path_relation   = Path(cur_path / './Data/dat/relation')
path_relation_tgc = path_dat / ('relation_tgc')
path_basket     = path_stock
path_backtest   = Path(cur_path / 'ZBackTest')
path_input      = path_dat / 'input'

path_src_data   = Path('C:/Users/liuxh/Nutstore/.nutstore_bGl1eGguemouY25AZ21haWwuY29t/我的坚果云/Paper/AData')

CLOSE       = 'CLOSE'
OPEN        = 'OPEN'
raw_cols    = ['OPEN', 'CLOSE', 'HIGH', 'LOW', 'VOLUME', 'TURN', 'RSI']

# model_name_list = ['LSTM_MF', 'LSTM_SF', '']
markets     = ['NASDAQ', 'NYSE', 'HK']
market      = 'NYSE'

# threshold - discrete
threshold_discrete_dict = {
    'r1d'   : 0.000,
    'r5d'   : 0.000,
    'r20d'  : 0.000,
    'r60d'  : 0.000
}

# threshold -


# relation
min_support_same = 0.6
min_support_diff = 0.6


# market = 'SH'
# start_date  = '2019-06-01'
# end_date    = '2020-06-01'
feat_label = {
    'feat_static':  ['OPEN', 'CLOSE', 'LOW', 'HIGH', 'VOLUME'],
    'feat_arith':   ['RSI', 'avg_5', 'avg_20', 'avg_60'],
    'feat_seq':     ['OPEN', 'CLOSE', 'LOW', 'HIGH', 'VOLUME'],
    'label':        ['label']
}
args = {
    'train_start'   : '2018-01-01',
    'train_end'     : '2018-12-31',
    'valid_start'   : '2019-01-01',
    'valid_end'     : '2019-03-29',
    'test_start'    : '2019-04-01',
    'test_end'      : '2019-06-28',

    'seq_len':      20,
    'feat_label':   feat_label,

    # MODEL
    'loss_func':    F.mse_loss,
    # 'loss_func':            Loss_Rank(len(stocks), alpha=0.5), #0.2
    'epochs':       1000,
    'dropout':      0.8,

    # Train
    'lr':           0.001,
    'step_lr':      True,
    'step_size':    20,
    'gamma':        0.8,

    # TEST
    'nmax':         1,
    'principal':    100000,

    # 'save_epoch':   100,

    # ADJ
    # 'adj_end_date':         '2018-03-15',
    # 'threshold_discrete':   0.008,
    'nday_relation'             : 30,
    'min_support_same'          : min_support_same,
    'min_support_diff'          : min_support_diff,
    'threshold_discrete_dict'   : threshold_discrete_dict
}


args_pic = {
    'title': 'Q1',
    'path_save': './Z_Result',
    'figsize': (4, 4),
    'linewidth': 1.5,
}

models_q1_1 = ['LSTM_MF', 'GRU_MF', 'RNN_MF', 'LSTM_SF', 'GRU_SF', 'RNN_SF', 'Base_Line']
models_q1_2 = ['LSTM_MF', 'Simple_MF', 'Primary_SF', 'Base_Line']
models_q2   = ['LSTM_MF', 'LSTM_SF', 'WGCN_MF', 'WGCN_SF', 'Base_Line'] # GCN
models_q3   = ['NASDAQ', 'SH', 'NYSE', 'SZ', 'HK']

colors = {
# ['blue', 'green', 'red', 'cyan', 'magenta', 'brown', 'black']
    'Base_Line': 'red',
    'Simple_MF': 'violet',
    'Primary_SF': 'yellowgreen',
    'LSTM_MF': 'chocolate',
    'LSTM_SF': 'wheat',
    'GRU_MF': 'green',
    'GRU_SF': 'lightgreen',
    'RNN_MF': 'blue',
    'RNN_SF': 'lightblue',
    'WGCN_MF': 'black',
    'WGCN_SF': 'grey',

    'NASDAQ_WGCN_MF': 'black',
    'NASDAQ_A_Index': 'red',
    'NYSE_WGCN_MF': 'black',
    'NYSE_A_Index': 'red',
    'SH_WGCN_MF': 'black',
    'SH_A_Index': 'red',
    'SZ_WGCN_MF': 'black',
    'SZ_A_Index': 'red',
    'HK_WGCN_MF': 'black',
    'HK_A_Index': 'red',
}




# file name
# suffix_basket = ''
# for (k, v) in threshold_discrete_dict.items(): suffix_basket += (str(v) + '_')

