import sys

sys.path.append("..")
import math
import os
import pandas as pd

try:
    import feather
except ImportError:
    pass
from analytics.index_data import index_returns, get_index_quotes
from arch import arch_model
from math import log, exp, sqrt
import numpy as np
from scipy.optimize import minimize_scalar
from statsmodels.tsa.ar_model import AR

import matplotlib.pyplot as plt


def calc_returns(index_list=["HY", "IG"], years=None):
    returns = (
        index_returns(index=index_list, tenor="5yr", years=years)
        .reset_index(["tenor"], drop=True)
        .swaplevel(1, 0)
    )
    returns = returns.groupby(level=["date", "index"]).nth(-1)["price_return"]
    returns = returns.unstack().dropna()
    return returns


def calc_betas(returns=None, spans=[5, 20], index_list=["HY", "IG"]):
    if returns is None:
        returns = calc_returns(index_list=index_list)
    return pd.concat(
        [
            (
                returns.ewm(span=span)
                .cov()
                .groupby(level="date")
                .apply(lambda df: df / np.diag(df))
            )
            for span in spans
        ],
        axis=1,
        keys=spans,
    )


def plot_betas(betas=None, index_list=["HY", "IG"]):
    spans = [5, 20]
    if betas is None:
        betas = calc_betas(spans=spans, index_list=index_list)
    for beta, span in zip(betas, spans):
        plt.plot(beta, label="EWMA" + str(span))
    plt.xlabel("date")
    plt.ylabel("beta")
    plt.legend()


def calc_realized_vol(returns=None):

    # three ways of computing the volatility
    # 1) 20 days simple moving average
    # 2) exponentially weighted moving average
    # 3) GARCH(1,1), we scale returns by 10 to help with the fitting
    if returns is None:
        returns = calc_returns()

    vol_sma = returns.rolling(20).std() * math.sqrt(252)
    vol_ewma = returns.ewm(span=20).std() * math.sqrt(252)
    scale = 10
    vol_garch = pd.DataFrame()
    for index in returns:
        am = arch_model(scale * returns[index].dropna())
        res = am.fit()
        vol_garch[index] = res.conditional_volatility * math.sqrt(252) / scale

    vol = pd.concat(
        [vol_sma, vol_ewma, vol_garch], axis=1, keys=["sma", "ewma", "garch"]
    )
    return vol


def spreads_ratio(series=list(range(22, 29)), index=["IG", "HY"], tenor="5yr"):
    df = get_index_quotes(series=series, index=index, tenor=tenor)
    df = df["model_spread"].groupby(["date", "index"]).last().unstack()
    df["ratio"] = df.HY / df.IG
    return df


def loglik(beta, returns):
    x = returns.HY - beta * returns.IG
    model = AR(x, missing="drop")
    fit = model.fit(maxlag=2)
    return -fit.llf


if __name__ == "__main__":
    returns = calc_returns()
    betas = calc_betas(returns)
    plot_betas(betas)
    vol = calc_realized_vol(returns)
    ratios = spreads_ratio()
    prog = minimize_scalar(loglik, bracket=(3, 5), args=(returns,))