import math
import os
import pandas as pd
import feather
from index_data import index_returns
from arch import arch_model
from math import log, exp, sqrt
import numpy as np
from scipy.optimize import minimize_scalar

returns = index_returns(index=['IG', 'HY'], tenor='5yr')
returns_hy = (returns.
              xs('HY', level=1).
              dropna().
              reset_index(level='series').
              groupby(level=['date']).
              nth(-1))
returns_hy = returns_hy.set_index('series', append=True)
returns_ig = returns.xs('IG', level=1).reset_index('tenor', drop=True)
# hy starts trading later than ig, so we line it up based on hy series
df = pd.merge(returns_hy, returns_ig, left_index=True, right_index=True,
              suffixes=('_hy','_ig'))
returns = df[['price_return_hy', 'price_return_ig']]
returns.columns = ['hy', 'ig']
returns = returns.reset_index('series', drop=True)
feather.write_dataframe(returns.reset_index(),
                        os.path.join(os.environ["DATA_DIR"], "index_returns.fth"))

#returns = returns.groupby('date').nth(-1)
# three ways of computing the volatility
# 20 days simple moving average
vol_sma = returns.hy.rolling(20).std() * math.sqrt(252)
vol_ewma = returns.hy.ewm(span=20).std() * math.sqrt(252)
# GARCH(1,1)
# we scale returns by 10 to help with the fitting
scale = 10
am = arch_model(scale * returns.hy.dropna())
res = am.fit()
vol_garch = res.conditional_volatility * math.sqrt(252)/scale
vol = pd.concat([vol_sma, vol_ewma, vol_garch], axis=1, keys=['sma', 'ewma', 'garch'])

## let's get the betas
beta_ewma = (returns.
             ewm(span=20).
             cov().
             groupby(level='date').
             apply(lambda df: df.values[0,1]/df.values[1,1]))

beta_ewma5 = (returns.
             ewm(span=5).
             cov().
             groupby(level='date').
             apply(lambda df: df.values[0,1]/df.values[1,1]))

feather.write_dataframe(beta_ewma.to_frame('beta'),
                        os.path.join(os.environ['DATA_DIR'], "beta.fth"))

def loglik(beta, returns):
    x = (returns.hy - beta*returns.ig)
    model = AR(x, missing='drop')
    fit = model.fit(maxlag=1)
    return - fit.llf

r = []
for beta in np.arange(3, 5, 0.01):
    prog = minimize(loglik, np.array([0.1, 0.1, 0.1]), args=(returns, beta),
                    bounds=[(None, None), (1e-6, None), (None, None)],
                    method='L-BFGS-B')
    r.append(prog.fun)

r = []
for beta in np.arange(3, 5, 0.01):
    r.append(test(returns, beta))