from analytics import on_the_run
from analytics.basket_index import BasketIndex
from analytics.utils import bus_day
from analytics.index_data import index_returns
from utils.db import serenitas_pool
import datetime
import numpy as np
import pandas as pd
from typing import Tuple
from cds_rebook import default_adjustment

def get_defaulted_info(conn, auction_date: datetime.date) -> Tuple[int, str]:
    with conn.cursor() as c:
        c.execute("SELECT id, seniority FROM defaulted WHERE auction_date=%s",
                  (auction_date,))
        cid, seniority = c.fetchone()
    conn.commit()
    return cid, seniority

dr = pd.date_range("2015-09-27", "2019-11-25", freq=bus_day)
prices = {}
s = -1
index_list = []
conn = serenitas_pool.getconn()
for d in dr:
    print(d)
    otr = on_the_run("HY", d)
    if otr > s:
        s = otr
        index_list.append(BasketIndex("HYBB", s, ["5yr"], value_date=d))
    for index in index_list:
        index.value_date = d
        prices[(index.series, index.version, d)] = 100 * (1 - index.pv()['5yr'])
        if index._lastdate == d.date():
            cid, seniority = get_defaulted_info(conn, d.date())
            _, fee = default_adjustment(cid, seniority, index.maturities[0])
            weight = next(w for _, w, c in index.items() if c.defaulted)
            old_price = 1 - index.pv()['5yr']
            new_price = (index.factor * old_price - weight * (1 - fee))/(index.factor - weight)
            prices[(index.series, index.version + 1, d)] = new_price * 100
serenitas_pool.putconn(conn)

hybb = pd.DataFrame.from_dict(prices, orient='index')
hybb.index = pd.MultiIndex.from_tuples(hybb.index)
hybb.index.names = ['series', 'version', 'date']
hybb.columns = ['HYBB']
hybb = hybb.sort_index()
hybb = hybb.groupby(['series', 'version']).pct_change()
hybb['day_frac'] = hybb.groupby(["series", "version"]).transform(
    lambda s: s.index.get_level_values('date').to_series().diff().astype("timedelta64[D]") / 360
)
hybb["HYBB"] += hybb.day_frac * 0.05
hybb = hybb.dropna().groupby('date').nth(-1)
df = index_returns(index="HY", tenor="5yr", series=[25, 26, 27, 28, 29, 30, 31, 32, 33],
                   from_date=datetime.date(2015, 9, 27),
                   years=None)
df = df["price_return"].reset_index(level=["index", "tenor"], drop=True)

hy_returns = df.dropna().groupby('date').nth(-1)
hy_returns.name = 'HY'

returns = pd.concat([hybb['HYBB'], hy_returns], axis=1)
# weekly
beta_weekly = (returns
               .ewm(span=5)
               .cov()
               .groupby(level=0)
               .apply(lambda df: df / np.diag(df)))
beta_weekly = beta_weekly.xs("HY", level=1)["HYBB"]

## monthly
beta_monthly = (returns
               .ewm(span=20)
               .cov()
               .groupby(level=0)
               .apply(lambda df: df / np.diag(df)))
beta_monthly = beta_monthly.xs("HY", level=1)["HYBB"]