#!/usr/bin/env python
# coding: utf-8

from abc import ABC, abstractmethod

import pandas as pd

import matplotlib.pyplot as plt
import numpy as np
import torch.nn.functional as F

from napoleontoolbox.neural_net import roll_multi_layer_lstm
from napoleontoolbox.file_saver import dropbox_file_saver
from napoleontoolbox.neural_net import roll_multi_layer_perceptron
from napoleontoolbox.boosted_trees import roll_lightgbm
from napoleontoolbox.features import features_type
from napoleontoolbox.utility import weights
from napoleontoolbox.utility import utility
from napoleontoolbox.rebalancing import allocation

import torch
import torch.nn as nn



class AbstractRunner(ABC):
    def __init__(self, starting_date = None, running_date = None, drop_token=None, dropbox_backup = True, save_model=False, supervision_npy_file_suffix='_supervision.npy', macro_supervision_npy_file_suffix='_macro_supervision.npy', features_saving_suffix='_features.npy', features_names_saving_suffix='_features_names.npy', returns_pkl_file_name='returns.pkl', local_root_directory='../data/', user = 'napoleon', n_start = 252):
        super().__init__()
        self.supervision_npy_file_suffix=supervision_npy_file_suffix
        self.macro_supervision_npy_file_suffix=macro_supervision_npy_file_suffix
        self.features_saving_suffix=features_saving_suffix
        self.features_names_saving_suffix=features_names_saving_suffix
        self.returns_pkl_file_name=returns_pkl_file_name
        self.local_root_directory=local_root_directory
        self.user=user
        self.n_start = n_start
        self.dropbox_backup = dropbox_backup
        self.dbx = dropbox_file_saver.NaPoleonDropboxConnector(drop_token=drop_token,dropbox_backup=dropbox_backup)
        self.save_model = save_model
        self.running_date = running_date
        self.starting_date = starting_date
        self.minimize_drawdown = True
        self.equaly_weighted = True
        self.erc_weighted = True
        self.ivp_allocation = True
        self.mvp_allocation = True
        self.hrp_allocation = True
        self.mdp_allocation = True

    @abstractmethod
    def runTrial(self,saver, seed, sup, layers, epochs, n_past_features, n, s, whole_history, advance_feature, advance_signal,
                 normalize, activation_string, convolution, lr_start, lr_type, low_bound, up_bound):
        pass


class SimpleEnsemblingRunner(AbstractRunner):
    def runTrial(self, saver, seed,  n, s, low_bound, up_bound):
        meArg = (seed, n, s, low_bound, up_bound)
        meArgList = list(meArg)
        meArgList = [str(it) for it in meArgList]
        savingKey = ''.join(meArgList)
        savingKey = savingKey.replace('[', '')
        savingKey = savingKey.replace(']', '')
        savingKey = savingKey.replace(',', '')
        savingKey = savingKey.replace(' ', '')
        savingKey = savingKey.replace('.', '')
        savingKey = 'T_' + savingKey

        print('Launching computation with parameters : '+savingKey)

        df = pd.read_pickle(self.local_root_directory + self.returns_pkl_file_name)
        print(df.columns)
        df['Date'] = pd.to_datetime(df['Date'])
        df = df.set_index('Date')
        df = df.fillna(method='ffill')
        df = df[df.index >= self.starting_date]
        print('max before filtering ' + str(max(df.index)))
        df = df[df.index <= self.running_date]

        print('max after filtering ' + str(max(df.index)))

        check_drawdown_existence = saver.checkRunExistence('dd_'+savingKey)
        if self.minimize_drawdown and not check_drawdown_existence:
            print('minimizing drawdown')
            df_res = df.copy()
            portfolio, weight_mat = allocation.rolling_allocation(
                allocation.minDrawdown,
                df_res,
                n=n,
                s=s,
                normalize=True,
                low_bound=low_bound,
                up_bound=up_bound,
                filtering_threshold=0.9
            )
            df_ret = df_res.fillna(method='bfill').pct_change().fillna(0).copy()
            weight_normalized, last_predicted_weights = weights.weights_shift(weight_mat=weight_mat)
            portfolio = np.cumprod(np.prod(df_ret * weight_mat.values + 1, axis=1))
            saver.saveResults('dd_'+savingKey, portfolio, weight_normalized,last_predicted_weights)

        check_ew_existence = saver.checkRunExistence('ew_'+savingKey)
        if self.equaly_weighted and not check_ew_existence:
            print('equally weighted')
            df_ret = df_res.fillna(method='bfill').pct_change().fillna(0).copy()
            weight_normalized, last_predicted_weights = weights.weights_shift(weight_mat=weight_mat)
            portfolio = np.cumprod(np.prod(df_ret * weight_mat.values + 1, axis=1))
            saver.saveResults('ew_'+savingKey, portfolio, weight_normalized,last_predicted_weights)


        check_erc_existence = saver.checkRunExistence('erc_'+savingKey)
        if self.erc_weighted and not check_erc_existence:
            print('computing erc weights')
            df_res = df.copy()
            # Compute rolling weights
            portfolio, weight_mat = allocation.rolling_allocation(
                allocation.ERC,
                df_res,
                n=n,
                s=s,
                ret=False,
                low_bound=low_bound,
                up_bound=up_bound,
                filtering_threshold=0.9
            )
            df_ret = df_res.fillna(method='bfill').pct_change().fillna(0).copy()
            weight_normalized, last_predicted_weights = weights.weights_shift(weight_mat=weight_mat)
            portfolio = np.cumprod(np.prod(df_ret * weight_mat.values + 1, axis=1))
            saver.saveResults('erc_'+savingKey, portfolio, weight_normalized,last_predicted_weights)


        check_ivp_existence = saver.checkRunExistence('ivp_'+savingKey)
        if self.ivp_allocation and not check_ivp_existence:
            print('computing ivp weights')
            # Compute rolling weights
            df_res = df.copy()
            portfolio, weight_mat = allocation.rolling_allocation(
                allocation.IVP,
                df_res,
                n=n,
                s=s,
                normalize=True,
                low_bound=low_bound,
                up_bound=up_bound,
                filtering_threshold=0.9
            )
            df_ret = df_res.fillna(method='bfill').pct_change().fillna(0).copy()
            weight_normalized, last_predicted_weights = weights.weights_shift(weight_mat=weight_mat)
            portfolio = np.cumprod(np.prod(df_ret * weight_mat.values + 1, axis=1))
            saver.saveResults('ivp_'+savingKey, portfolio, weight_normalized,last_predicted_weights)



        check_mpv_existence = saver.checkRunExistence('mvp_'+savingKey)
        if self.mvp_allocation and not check_mpv_existence:
            print('computing mvp weights')
            df_res = df.copy()
            portfolio, weight_mat = allocation.rolling_allocation(
                allocation.MVP_uc,
                df_res,
                n=n,
                s=s,
                low_bound=low_bound,
                up_bound=up_bound,
                ret=False,
                filtering_threshold=0.9
            )
            df_ret = df_res.fillna(method='bfill').pct_change().fillna(0).copy()
            weight_normalized, last_predicted_weights = weights.weights_shift(weight_mat=weight_mat)
            portfolio = np.cumprod(np.prod(df_ret * weight_mat.values + 1, axis=1))
            saver.saveResults('mvp_'+savingKey, portfolio, weight_normalized,last_predicted_weights)


        check_hrp_existence = saver.checkRunExistence('hrp_'+savingKey)
        if self.hrp_allocation and not check_hrp_existence:
            print('computing hrp weights')
            df_res = df.copy()
            portfolio, weight_mat = allocation.rolling_allocation(
                allocation.HRP,
                df_res,
                n=n,
                s=s,
                ret=True,
                drift=True,
                method='centroid',
                metric='mse',
                up_bound=up_bound,
                low_bound=low_bound,
                filtering_threshold=0.9
            )
            df_ret = df_res.fillna(method='bfill').pct_change().fillna(0).copy()
            weight_normalized, last_predicted_weights = weights.weights_shift(weight_mat=weight_mat)
            portfolio = np.cumprod(np.prod(df_ret * weight_mat.values + 1, axis=1))
            saver.saveResults('hrp_'+savingKey, portfolio, weight_normalized,last_predicted_weights)


        check_mdp_existence = saver.checkRunExistence('mdp_'+savingKey)
        if self.mdp_allocation and not check_mdp_existence:
            print('computing mdp weights')
            # Compute rolling weights
            df_res = df.copy()
            portfolio, weight_mat = allocation.rolling_allocation(
                allocation.MDP,
                df_res,
                n=n,
                s=s,
                ret=False,
                drift=True,
                up_bound=up_bound,
                low_bound=low_bound,
                filtering_threshold=0.9
            )
            df_ret = df_res.fillna(method='bfill').pct_change().fillna(0).copy()
            weight_normalized, last_predicted_weights = weights.weights_shift(weight_mat=weight_mat)
            portfolio = np.cumprod(np.prod(df_ret * weight_mat.values + 1, axis=1))
            saver.saveResults('mdp_'+savingKey, portfolio, weight_normalized,last_predicted_weights)
