import abc
import fnmatch
import heapq
from pathlib import Path
from typing import List, Dict
import pandas as pd
import numpy as np
from torch import optim
from torch.optim import lr_scheduler
from torch.utils.data import TensorDataset, DataLoader
import torch.nn.functional as F
import torch
import os
from A_Global import path_stock, path_backtest
cur_path = Path(os.path.dirname(__file__))


class BackTest(metaclass=abc.ABCMeta):

    def __init__(self, get_data, model, hyperparameter):
        # 类似模板模式，训练流程确定，中间的train，backtest等根据模型而定
        # 提供一个环境类，维护stock，time，feat等，并根据环境提供cachename等
        self.hyperparameter = hyperparameter

        ''' Model, Loss, Optim'''
        self.model = model
        self.loss_func = self.hyperparameter['loss_func']
        self.opt = optim.Adam(self.model.parameters(), lr=self.hyperparameter['lr'])
        self.scheduler = lr_scheduler.StepLR(self.opt, step_size=self.hyperparameter['step_size'],
                                        gamma=self.hyperparameter['gamma'])

        self.result_name = hyperparameter['result_name']
        if not self.result_name: self.result_name = hyperparameter['model_type']
        self.market     = hyperparameter['market']
        self.model_type = hyperparameter['model_type']
        self.feat_label = hyperparameter['feat_label']
        self.nmax, self.principal = self.hyperparameter['nmax'], self.hyperparameter['principal']
        self.hyperparameter['stocks'] = np.genfromtxt(path_stock / ('STOCKs_' + self.market + '.txt'), dtype=str)

        self.cache_model = self.model_type + '_' + self.market + '_' \
                           + str(len(self.feat_label['feat_static'])) + 's' \
                           + str(len(self.feat_label['feat_arith'])) + 'a' \
                           + str(len(self.feat_label['feat_seq'])) + 'seq'

        if 'test_periods' not in hyperparameter.keys():
            self.dl_train, self.dl_valid, self.dl_test = self.make_dataloader(get_data)

    def gen_cache_model(self):
        return self.model_type + '_' + self.market + '_' \
                           + str(len(self.feat_label['feat_static'])) + 's' \
                           + str(len(self.feat_label['feat_arith'])) + 'a' \
                           + str(len(self.feat_label['feat_seq'])) + 'seq'

    @abc.abstractmethod
    def make_dataloader(self, get_data): pass

    @abc.abstractmethod
    def train_model(self): pass
    ''' 写在一个文件，train_model 不独立 '''
    @abc.abstractmethod
    def backtest_model(self) -> Dict: pass

    def train(self):
        ''' Train '''
        loss_arr_train = []
        epochs = self.hyperparameter['epochs']
        for epoch in range(epochs):
            self.model.train()
            loss_train = self.train_model()
            loss_arr_train.append(loss_train)

            ''' ========================= Inner Test ================================ '''
            metrics = self.backtest_model()
            financials = metrics['financials']
            alpha = round(financials[-1] / financials[0], 3)
            model_name = '%s_%depoch_%1.1f_%1.6fLTrain.pt' % (
                self.cache_model, epoch, alpha, loss_train
            )

            torch.save(self.model.state_dict(), self.path_model / model_name)
            message = 'Save Model ' + model_name
            (mdd, mdd_start, mdd_end) = metrics['MDD']

            ''' ========================= Inner Test ================================ '''
            print(
                '%3d Train %2.7f MSE_Test %2.7f Return %7d MDD %.1f%% lr %.5f %s' %
                (epoch,
                 loss_arr_train[epoch],
                 metrics['MSE'],
                 financials[-1],
                 (100 * mdd),
                 self.scheduler.get_last_lr()[-1],
                 message
                 ))
            self.financials = None
        print('TRAIN OVER..')

    def backtest(self):
        save_epoch = self.hyperparameter['save_epoch']
        model_name_re = self.cache_model + '_' + str(save_epoch) + 'epoch_*.pt'
        model_name = self.find_file_name(self.path_model, model_name_re)
        if not model_name:
            print('！！！', str(self.path_model), model_name_re, '无法匹配任何文件')
            exit()

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

        ''' ========================= Backtest ================================ '''
        metrics = self.backtest_model()

        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=self.date_test, columns=[self.result_name])
        df.to_csv(Path('./') / 'z_irrs.csv')

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

    def backtest_oneday(self, y_pred, y_true, nmax=2, principal=100000):
        '''

        :return:
            [top 5 stocks with the highest yields]
            accuracy
        '''
        ids_preds = list(map(list(y_pred).index, heapq.nlargest(nmax, y_pred)))
        y_pred_max = y_pred[ids_preds]
        y_true_max = y_true[ids_preds]

        acc_arr = [1 if (y_pred_max[i] > 1 and y_true_max[i] > 1) or (y_pred_max[i] < 1 and y_true_max[i] < 1) else 0
                   for i in range(nmax)]
        acc = sum(acc_arr) / nmax

        principals = [principal / nmax for _ in range(nmax)]
        principal = sum(np.array(principals) * np.array(y_true_max))

        return ids_preds, acc, principal

    def find_file_name(self, path, name_pattern):
        names = os.listdir(path)
        for _ in names:
            if fnmatch.fnmatch(_, name_pattern):
                return _
        return None