import datetime
import pandas as pd
from analytics import Portfolio, CreditIndex
from analytics.curve_trades import on_the_run
from analytics.index_data import index_returns
from math import sqrt
from psycopg2.extensions import connection
from typing import Tuple


def get_index_portfolio(
    d: datetime.date,
    conn: connection,
    strategies: Tuple[str] = None,
    exclude_redcode=[],
    **kwargs
):
    sql_str = (
        "SELECT security_id AS redcode, notional, maturity "
        "FROM list_cds_positions_by_strat(%s) "
    )
    params = (d,)
    if strategies is not None:
        sql_str += "WHERE folder in %s"
        params += (strategies,)

    with conn.cursor() as c:
        c.execute(sql_str, params)
        trades = [
            CreditIndex(
                redcode=rec.redcode,
                maturity=rec.maturity,
                notional=rec.notional,
                value_date=d,
            )
            for rec in c
            if rec.redcode not in exclude_redcode
        ]
        portf = Portfolio(trades)
        portf.mark()
    return portf


def VaR(
    portf: Portfolio,
    index_type: str = "IG",
    quantile=0.05,
    years: int = 5,
    period="monthly",
):
    df = index_returns(
        index=index_type, years=years, tenor=["3yr", "5yr", "7yr", "10yr"]
    )
    df = df.reset_index(["index"], drop=True).reorder_levels(
        ["date", "series", "tenor"]
    )
    returns = df.spread_return.dropna().reset_index("series")
    returns["dist_on_the_run"] = returns.groupby("date")["series"].transform(
        lambda x: x.max() - x
    )
    del returns["series"]
    returns = returns.set_index("dist_on_the_run", append=True).unstack("tenor")
    returns.columns = returns.columns.droplevel(0)
    portf.reset_pv()

    otr = on_the_run(index_type)
    spreads = pd.DataFrame(
        {
            "spread": portf.spread,
            "tenor": [ind.tenor for ind in portf.indices],
            "dist_on_the_run": [otr - ind.series for ind in portf.indices],
        }
    )
    spreads = spreads.set_index(["dist_on_the_run", "tenor"])
    r = []
    for k, g in returns.groupby(level="date", as_index=False):
        shocks = g.reset_index("date", drop=True).stack("tenor")
        shocks.name = "shocks"
        portf.spread = spreads.spread * (1 + spreads.join(shocks).shocks)
        r.append((k, portf.pnl))
    pnl = pd.DataFrame.from_records(r, columns=["date", "pnl"], index=["date"])
    if period == "daily":
        return float(pnl.quantile(quantile))
    elif period == "monthly":
        return float(pnl.quantile(quantile)) * sqrt(20)
    else:
        raise ValueError("period needs to be either 'daily' or 'monthly'")


def insert_curve_risk(d: datetime.date, conn: connection, strategies=("SER_IGCURVE",)):
    sql_str = "INSERT INTO curve_risk VALUES(%s, %s, %s, %s)"
    with conn.cursor() as c:
        for strat in strategies:
            portf = get_index_portfolio(
                d, conn, (strat,), exclude_redcode=["2I65BYDU6"]
            )
            index_type = "IG" if "IG" in strat else "EU"
            var = VaR(portf, index_type, period="daily")
            c.execute(sql_str, (d, strat, var, "USD" if index_type == "IG" else "EUR"))
    conn.commit()