import heapq
import os
from pathlib import Path
import numpy as np
from torch.utils.data import TensorDataset, DataLoader
import torch.nn.functional as F
import torch
from torch.optim import lr_scheduler
from torch import optim
from A_Model import get_model
from B_data_tool import get_static_data, get_seq_data
from B_date_tool import gen_date
from D_relation import gen_adj
from A_Global import path_stock, args
from A_Data import get_data_gcn_mf
from tool_data import get_stock_filted
from metric_tool import IRR, SR, MDD, IRR_adj_day, IRR_acc_day, gen_metrics
import pandas as pd
cur_path = Path(os.path.dirname(__file__))
path_model = cur_path / 'Model'
from tool_backtest import backtest_oneday

def backtest_model(model, dl_test, adj):
    nmax        = 1
    financials  = [10000]
    metrics = {}
    mses = []

    model.eval()
    for arith_b, seq_b, y_b in dl_test:
        static_b = seq_b[:, -1, :]
        arith_b = arith_b.squeeze()

        preds = model(static_b, arith_b, seq_b, adj)

        # ids_stocks_buy, acc, money = backtest_oneday(preds, y_b, nmax, financials[-1])  #
        ids_stocks_buy, acc, money_add = backtest_oneday(preds, y_b, nmax, financials[0])

        financials.append(financials[-1] + money_add)
        mses.append(F.mse_loss(preds, y_b).item())

    # metrics['financials'] = Financial(financials)
    metrics['financials'] = financials
    metrics['MSEs'] = mses
    metrics['MSE'] = sum(mses) / len(mses)
    metrics['IRR'] = IRR(financials)
    metrics['irr_list'] = IRR_acc_day(financials)
    metrics['SR'] = SR(financials)
    metrics['MDD'] = MDD(financials)
    return metrics

def train_model(model, dl_train, dl_valid, dl_test, adj, epochs, loss_func, opt):
    if not os.path.exists(path_model): os.mkdir(path_model)

    loss_arr_train = []
    for epoch in range(epochs):
        model.train()

        n_train, loss_sum_train = 0, 0
        for arith_b, seq_b, y_b in dl_train:
            static_b = seq_b[:, -1, :]
            arith_b = arith_b.squeeze()
            preds = model(static_b, arith_b, seq_b, adj)
            loss = loss_func(preds, y_b)

            loss_sum_train += loss.item()
            n_train += 1

            opt.zero_grad()
            loss.backward()
            opt.step()

        # if args['step_lr'] != None:
        #     scheduler.step()
        loss_train = loss_sum_train / n_train
        loss_arr_train.append(loss_train)
        print(loss_train)

    torch.save(model.state_dict(), path_model / 'MODEL_TRMF_' + str(epochs - 1) + 'epoch_0.pt')
    print('OVER TRAIN')

def train(market):

    ''' raw data '''
    args['market']      = market
    stocks, stock2idx   = get_stock_filted(market)
    args['stocks']      = stocks

    '''data'''
    arith_data_dict, seq_data_dict, adj = get_data_gcn_mf(args)

    bs          = len(args['stocks'])
    dl_train    = DataLoader(TensorDataset(arith_data_dict['x_train'], seq_data_dict['x_train'], seq_data_dict['y_train']),
                             batch_size=bs)
    dl_valid    = DataLoader(TensorDataset(arith_data_dict['x_valid'], seq_data_dict['x_valid'], seq_data_dict['y_valid']),
                             batch_size=bs)
    dl_test     = DataLoader(TensorDataset(arith_data_dict['x_test'], seq_data_dict['x_test'], seq_data_dict['y_test']),
                             batch_size=bs)

    ''' model '''
    model       = get_model('GCN_MF')

    loss_func   = args['loss_func']
    opt         = optim.Adam(model.parameters(), lr=args['lr'])

    '''train'''
    train_model(model, dl_train, dl_valid, dl_test, adj, args['epochs'], loss_func, opt)

def backtest(market, dl_set_str, save_epoch):
    args['market']      = market
    stocks, stock2idx   = get_stock_filted(market)
    args['stocks']      = stocks

    ''' date '''
    date_dict   = gen_date(args)

    '''data'''
    arith_data_dict, seq_data_dict, adj = get_data_gcn_mf(args)

    bs          = len(args['stocks'])
    dl_train    = DataLoader(TensorDataset(arith_data_dict['x_train'], seq_data_dict['x_train'], seq_data_dict['y_train']),
                             batch_size=bs)
    dl_valid    = DataLoader(TensorDataset(arith_data_dict['x_valid'], seq_data_dict['x_valid'], seq_data_dict['y_valid']),
                             batch_size=bs)
    dl_test     = DataLoader(TensorDataset(arith_data_dict['x_test'], seq_data_dict['x_test'], seq_data_dict['y_test']),
                             batch_size=bs)
    dl_set = {'train':dl_train, 'valid':dl_valid, 'test':dl_test}[dl_set_str]

    ''' model '''
    model       = get_model('GCN_MF')

    model_name_re = 'cache_model_' + str(save_epoch) + 'epoch_*.pt'
    from tool_backtest import find_file_name
    model_name = find_file_name(path_model, model_name_re)
    if not model_name:
        print('！！！', str(path_model), model_name_re, '无法匹配任何文件')
        exit()

    loss_record = model_name[-17:-9]
    model.load_state_dict(torch.load(path_model / model_name))
    print('LOAD MODEL OVER..')

    ''' ========================= Backtest ================================ '''
    metrics = backtest_model(model, dl_set, adj)

    print('BACKTEST: Finantials - (%.2f ~ %.2f) IRR - %.2f MSE - %s SR - %.2f MDD - %.1f%%' % (
        metrics['financials'][0], metrics['financials'][-1],
        metrics['IRR'],
        loss_record,  # float(np.mean(mses)),
        metrics['SR'],
        (100 * metrics['MDD'][0])
    ))

    ''' ========================= SAVE ================================ '''
    df = pd.DataFrame(metrics['irr_list'], index=date_dict['date_' + dl_set_str], columns=[market])

    df.to_csv(Path('.') / ('z_irrs_' + dl_set_str + '.csv'))

    df_metrics = pd.DataFrame(
        [[round(metrics['IRR'], 2), loss_record, round(metrics['SR'], 2), '%.1f%%' % (100 * metrics['MDD'][0])]],
        index=[market],
        columns=['IRR', 'MSE', 'SR', 'MDD'])
    df_metrics.to_csv(Path('./') / ('z_result_' + dl_set_str + '.csv'))

if __name__ == '__main__':
    market = 'NYSE'

    ''' train '''
    train(market)

    save_epoch = 999
    ''' test '''
    backtest(market, 'test', save_epoch)

    ''' valid '''
    backtest(market, 'valid', save_epoch)