import datetime
import numpy as np
import pandas as pd
from serenitas.analytics.dates import tenor_t
from pandas.tseries.offsets import BDay
from dateutil.relativedelta import relativedelta
from serenitas.analytics.yieldcurve import get_curve
from serenitas.utils.db import serenitas_engine
from pyisda.legs import FeeLeg, ContingentLeg
from pyisda.curve import SpreadCurve
from pyisda.date import previous_twentieth

tenors = {
    "IG": ("3yr", "5yr", "7yr", "10yr"),
    "HY": ("3yr", "5yr", "7yr"),
    "EU": ("3yr", "5yr", "7yr", "10yr"),
    "XO": ("3yr", "5yr", "7yr", "10yr"),
    "HYBB": ("5yr",),
}

sql_str = "INSERT INTO index_risk VALUES(%s, %s, %s)"


def get_legs(index, series, tenors):
    fee_legs = {}
    contingent_legs = {}
    df = pd.read_sql_query(
        "SELECT tenor, maturity, coupon, issue_date "
        "FROM index_maturity "
        "WHERE index=%s AND series=%s and tenor IN %s "
        "ORDER BY maturity",
        serenitas_engine,
        params=(index, series, tenors),
        parse_dates=["maturity", "issue_date"],
    )
    df.coupon *= 1e-4
    for tenor, maturity, coupon, issue_date in df.itertuples(index=False):
        maturity_short = maturity - relativedelta(years=1)
        fee_legs[tenor] = (
            FeeLeg(issue_date, maturity, True, 1.0, 1.0),
            FeeLeg(issue_date, maturity_short, True, 1.0, coupon),
        )
        contingent_legs[tenor] = (
            ContingentLeg(issue_date, maturity, True),
            ContingentLeg(issue_date, maturity_short, True),
        )
    # number of seconds since epoch
    # number of days between 1900-1-1 and epoch
    df.maturity = df.maturity.to_numpy("M8[D]").astype(int) + 134774
    return fee_legs, contingent_legs, df


def index_pv(fl, cl, value_date, step_in_date, cash_settle_date, yc, sc, recovery):
    dl_pv = cl.pv(value_date, step_in_date, cash_settle_date, yc, sc, recovery)
    cl_pv = fl.pv(value_date, step_in_date, cash_settle_date, yc, sc, True)
    return dl_pv - cl_pv


if __name__ == "__main__":
    conn = serenitas_engine.raw_connection()
    for index in ["IG", "HY", "EU", "XO", "HYBB"]:
        if index in ["HY", "HYBB"]:
            recoveries = np.full(len(tenors[index]), 0.3)
        else:
            recoveries = np.full(len(tenors[index]), 0.4)
        for series in range(18, 41):
            if index in ["EU", "XO"] and series == 40:
                continue
            fee_legs, contingent_legs, df = get_legs(index, series, tenors[index])

            index_quotes = pd.read_sql_query(
                "SELECT distinct on (date, tenor) id, date, tenor, version, close_price "
                "FROM index_quotes_pre "
                "LEFT JOIN index_risk USING (id) "
                "WHERE index=%s AND series=%s "
                "AND source='MKIT' AND duration is NULL AND tenor IN %s "
                "ORDER BY date, tenor, version desc",
                serenitas_engine,
                params=(index, series, tenors[index]),
                parse_dates=["date"],
                index_col="id",
            )
            if index_quotes.empty:
                continue
            index_quotes.tenor = index_quotes.tenor.astype(tenor_t)
            index_quotes = index_quotes.sort_values("tenor")
            index_quotes["close_price"] = 1.0 - index_quotes["close_price"] / 100
            with conn.cursor() as c:
                for value_date, data in index_quotes.groupby("date"):
                    yc = get_curve(
                        value_date, "USD" if index in ["IG", "HY"] else "EUR"
                    )
                    data = data.reset_index().merge(df, on="tenor").set_index("id")
                    step_in_date = value_date + datetime.timedelta(days=1)
                    cash_settle_date = value_date + 3 * BDay()
                    start_date = previous_twentieth(value_date)
                    try:
                        sc = SpreadCurve(
                            value_date,
                            yc,
                            start_date,
                            step_in_date,
                            cash_settle_date,
                            data.maturity.values,
                            data.coupon.values,
                            data.close_price.values,
                            recoveries,
                        )
                    except ValueError:
                        break
                    for r in data[["coupon", "tenor", "close_price"]].itertuples():
                        fl, fl_short = fee_legs[r.tenor]
                        cl, cl_short = contingent_legs[r.tenor]

                        duration = fl.pv(
                            value_date, step_in_date, cash_settle_date, yc, sc, True
                        )
                        if cl_short.end_date <= value_date.date():
                            theta = None
                        else:
                            pv = index_pv(
                                fl_short,
                                cl_short,
                                value_date,
                                step_in_date,
                                cash_settle_date,
                                yc,
                                sc,
                                recoveries[0],
                            )
                            theta = r.close_price - pv + r.coupon

                        c.execute(
                            "INSERT INTO index_risk VALUES(%s, %s, %s)",
                            (r.Index, theta, duration),
                        )
            conn.commit()
    conn.close()