

import functools
import timeit
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from typing import Union
from enum import Enum

import qis.models.linear.ewm as ewm
import qis.models.linear.auto_corr as ac
import qis.plots.utils as put
import qis.plots.time_series as pts


def get_test_data(n1: int = 1000, n2: int = 3, means: Union[float, np.ndarray] = 0.0) -> pd.DataFrame:
    np_data = np.random.normal(means, 1.0, (n1, n2))
    random_nans = np.random.binomial(size=(n1, n2), n=1, p=0.5)
    np_data = np.where(random_nans, np.nan, np_data)
    df = pd.DataFrame(data=np_data,
                      index=pd.date_range(start='1Jan2020', periods=n1, freq='D'),
                      columns=[f"x{m+1}" for m in range(n2)])
    return df


def time_comp_with_pandas_ewm(span: int = 31) -> None:

    big_df = get_test_data(n1=10000, n2=10000)
    np_data = big_df.to_numpy()
    init_value = np.zeros(np_data.shape[1])

    def pd_ewm():
        big_df.ewm(span=span, adjust=False).mean()

    def np_to_ewm():
        ewm.ewm_recursion(a=np_data, span=span, init_value=init_value)

    def pd_to_ewm():
        ewm.compute_ewm(data=big_df, span=span, init_value=init_value)

    n = 20
    print(f"pandas ewm: {timeit.Timer(functools.partial(pd_ewm)).timeit(n) / n:.3f}")
    print(f"np_to_ewm: {timeit.Timer(functools.partial(np_to_ewm)).timeit(n) / n:.3f}")
    print(f"pd_to_ewm: {timeit.Timer(functools.partial(pd_to_ewm)).timeit(n) / n:.3f}")


def vol_comp_with_pandas_ewm(span: int = 31) -> None:
    df = get_test_data(n1=10000, n2=1)
    vol_pd = df.ewm(span=span, adjust=False).std()
    vol_ewm1 = ewm.compute_ewm_vol(data=df, span=span, mean_adj_type=ewm.MeanAdjType.NONE)
    vol_ewm2 = ewm.compute_ewm_vol(data=df, span=span, mean_adj_type=ewm.MeanAdjType.EWMA)
    joint = pd.concat([vol_pd.iloc[:, 0].rename('Pandas std'),
                      vol_ewm1.iloc[:, 0].rename('Ewm no adj'),
                      vol_ewm2.iloc[:, 0].rename('Ewm with Ewm adj')],
                      axis=1)

    pts.plot_time_series(df=joint,
                         title='Vols comparision',
                         legend_loc='upper left',
                         trend_line=put.TrendLine.AVERAGE,
                         var_format='{:.2f}')


class UnitTests(Enum):
    TIME_COMP_WITH_PANDAS_EWM = 1
    VOL_COMP_WITH_PANDAS_EWM = 2
    EWMA_NP = 3
    EWMA_DF = 4
    EWMA_MEAN = 5
    EWMA_VOL = 6
    EWMA_BETA = 7
    EWMA_AUTO_CORR = 8
    EWMA_CORR_MATRIX = 9


def run_unit_test(unit_test: UnitTests):

    np.random.seed(1)

    if unit_test == UnitTests.TIME_COMP_WITH_PANDAS_EWM:
        time_comp_with_pandas_ewm()

    elif unit_test == UnitTests.VOL_COMP_WITH_PANDAS_EWM:
        vol_comp_with_pandas_ewm()

    else:

        data = get_test_data(n1=10000, n2=3, means=np.array([-2.0, -1.0, 0.0]))
        ewm_lambda = 0.94
        ewm_lambda3 = np.array([0.94, 0.50, 0.10])
        plot_data, title = None, None

        if unit_test == UnitTests.EWMA_NP:
            ewm1 = ewm.ewm_recursion(a=data.iloc[:, 0].to_numpy(dtype=np.double), ewm_lambda=ewm_lambda, init_value=0.0)
            print(ewm1)

            ewm2 = ewm.ewm_recursion(a=data.to_numpy(), ewm_lambda=ewm_lambda, init_value=np.zeros(len(data.columns)))
            print(ewm2)

            ewm3 = ewm.ewm_recursion(a=data.to_numpy(), ewm_lambda=ewm_lambda3, init_value=np.zeros(len(data.columns)))
            print(ewm3)

        elif unit_test == UnitTests.EWMA_DF:

            ewm_df1 = ewm.compute_ewm(data=data.iloc[:, 0], ewm_lambda=ewm_lambda, init_type=ewm.InitType.MEAN)
            print(ewm_df1)

            ewm_df1_np = ewm.compute_ewm(data=data.iloc[:, 0].to_numpy(), ewm_lambda=ewm_lambda, init_type=ewm.InitType.MEAN)
            print(ewm_df1_np)

            ewm_df2_np = ewm.compute_ewm(data=data.to_numpy(), ewm_lambda=ewm_lambda, init_type=ewm.InitType.MEAN)
            print(ewm_df2_np)

            ewm_df2 = ewm.compute_ewm(data=data, ewm_lambda=ewm_lambda, init_type=ewm.InitType.MEAN)
            print(ewm_df2)

            ewm_df3 = ewm.compute_ewm(data=data, ewm_lambda=ewm_lambda3, init_type=ewm.InitType.MEAN)
            print(ewm_df3)

            ewm_df3.columns = [x + ' ewm' for x in data.columns]
            plot_data = pd.concat([data, ewm_df3], axis=1)
            title = 'ewm'

        elif unit_test == UnitTests.EWMA_MEAN:

            datas = {'np1': data.iloc[:, 0].to_numpy(), 'np': data.to_numpy(), 'series': data.iloc[:, 0], 'df': data}

            for key, data_ in datas.items():
                print(key)
                ewm_mean = ewm.compute_roll_mean(data=data,
                                                 mean_adj_type=ewm.MeanAdjType.EWMA,
                                                 ewm_lambda=ewm_lambda)
                data_adjusted = data - ewm_mean
                ewm_data = ewm.compute_ewm(data=data_adjusted,
                                              ewm_lambda=ewm_lambda)

            ewm_mean.columns = [x + ' ewm_mean' for x in data.columns]
            ewm_data.columns = [x + '-ewm_mean' for x in data.columns]
            plot_data = pd.concat([data, ewm_mean, ewm_data], axis=1)
            title = 'ewm-mean'

        elif unit_test == UnitTests.EWMA_VOL:

            datas = {'np1': data.iloc[:, 0].to_numpy(), 'np': data.to_numpy(), 'series': data.iloc[:, 0], 'df': data}

            for key, data_ in datas.items():
                print(key)
                ewm_data = ewm.compute_ewm_vol(data=data_,
                                                  ewm_lambda=ewm_lambda,
                                                  mean_adj_type=ewm.MeanAdjType.EWMA,
                                                  init_type=ewm.InitType.MEAN)
            print(ewm_data)
            ewm_data.columns = [x + ' ewm vol' for x in data.columns]
            data.columns = [x + '^2' for x in data.columns]
            plot_data = pd.concat([np.power(data, 2), ewm_data], axis=1)
            title = 'ewm vol'

        elif unit_test == UnitTests.EWMA_BETA:

            cross_xy_types = [ewm.CrossXyType.COVAR, ewm.CrossXyType.BETA, ewm.CrossXyType.CORR]

            for cross_xy_type in cross_xy_types:
                print(cross_xy_type)
                ewm_data = ewm.compute_ewm_cross_xy(x_data=data[[data.columns[0]]],
                                                       y_data=data,
                                                       cross_xy_type=cross_xy_type,
                                                       ewm_lambda=ewm_lambda,
                                                       mean_adj_type=ewm.MeanAdjType.EWMA,
                                                       init_type=ewm.InitType.MEAN)
            ewm_data.columns = [x + ' ewm beta' for x in data.columns]
            plot_data = pd.concat([data, ewm_data], axis=1)
            title = 'ewm beta'

        elif unit_test == UnitTests.EWMA_AUTO_CORR:
            plot_data = ac.compute_dynamic_auto_corr(data=data,
                                                     ewm_lambda=ewm_lambda,
                                                     mean_adj_type=ewm.MeanAdjType.NONE,
                                                     aggregation_type='median',
                                                     is_normalize=True)
            title = 'ewm auto-corr'

        else:
            return

        if plot_data is not None:
            pts.plot_time_series(df=plot_data,
                                 title=title,
                                 legend_loc='upper left',
                                 trend_line=put.TrendLine.AVERAGE,
                                 var_format='{:.2f}')

    plt.show()


if __name__ == '__main__':

    unit_test = UnitTests.EWMA_NP

    is_run_all_tests = False
    if is_run_all_tests:
        for unit_test in UnitTests:
            run_unit_test(unit_test=unit_test)
    else:
        run_unit_test(unit_test=unit_test)
