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 Iterable, Tuple, Union


def get_index_portfolio(
    d: datetime.date,
    conn: connection,
    fund: str = "SERCGMAST",
    strategies: Union[Tuple[str], None] = None,
    include_strategies: Union[str, None] = None,
    exclude_strategies: Union[str, None] = None,
    exclude_redcode: Iterable[str] = (),
    by_strat: bool = True,
    **kwargs,
):
    select_cols = [
        "security_id AS redcode",
        "security_desc",
        "sum(notional) AS notional" if not by_strat else "notional",
        "maturity",
    ]
    if by_strat:
        select_cols += ["folder"]
    sql_str = f"SELECT {','.join(select_cols)} FROM list_cds_positions_by_strat(%s, %s)"
    params = (d, fund)
    folder_filter = []
    if strategies is not None:
        if isinstance(strategies, tuple):
            folder_filter.append("folder IN %s")
        else:
            folder_filter.append("folder = %s")
        params += (strategies,)
    if include_strategies is not None:
        folder_filter.append("folder::text LIKE %s")
        params += (include_strategies,)
    if exclude_strategies is not None:
        folder_filter.append("folder::text NOT LIKE %s")
        params += (exclude_strategies,)
    if folder_filter:
        sql_str += " WHERE " + " AND ".join(folder_filter)
    if not by_strat:
        sql_str += " GROUP BY security_id, security_desc, maturity"
    with conn.cursor() as c:
        c.execute(sql_str, params)
        trades = [
            (
                CreditIndex(
                    redcode=rec.redcode,
                    maturity=rec.maturity,
                    notional=rec.notional,
                    value_date=d,
                    freeze_version=True,
                ),
                (rec.folder if by_strat else "", rec.security_desc),
            )
            for rec in c
            if rec.redcode not in exclude_redcode
        ]
    if trades:
        portf = Portfolio(*zip(*trades))
        portf.mark()
    else:
        portf = Portfolio([])
    return portf


def VaR(portf: Portfolio, quantile=0.05, years: int = 5, period="monthly"):
    index_types = tuple(set(t.index_type for t in portf))
    df = index_returns(
        index=index_types,
        years=years,
        end_date=portf.value_date,
        tenor=["3yr", "5yr", "7yr", "10yr"],
    )
    df = df.reorder_levels(["date", "index", "series", "tenor"])
    returns = df.spread_return.dropna().reset_index("series")
    returns["dist_on_the_run"] = returns.groupby(["date", "index"])["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()

    spreads = pd.DataFrame(
        {
            "spread": portf.spread,
            "tenor": [ind.tenor for ind in portf.indices],
            "index": [ind.index_type for ind in portf.indices],
            "dist_on_the_run": [
                on_the_run(ind.index_type, portf.value_date) - ind.series
                for ind in portf.indices
            ],
        }
    )
    spreads = spreads.set_index(["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).values
        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,
    fund: str = "SERCGMAST",
    strategies: Tuple[str] = ("SER_IGCURVE",),
):
    sql_str = (
        "INSERT INTO curve_risk VALUES(%s, %s, %s, %s, %s) "
        "ON CONFLICT (date, strategy) DO UPDATE SET "
        '"VaR"=excluded."VaR", currency=excluded.currency, fund=excluded.fund'
    )
    # add a portfolio with all strategies
    strategies = (*strategies, strategies)
    with conn.cursor() as c:
        for strat in strategies:
            portf = get_index_portfolio(
                d, conn, fund, strat, exclude_redcode=("2I65BYDU6",)
            )
            if portf:
                var = VaR(portf, period="daily")
                strat_name = "*" if isinstance(strat, tuple) else strat
                c.execute(sql_str, (d, strat_name, var, "USD", fund))
    conn.commit()