import pandas as pd
import numpy as np
from risk.portfolio import build_portfolio, generate_vol_surface


from serenitas.analytics.scenarios import run_portfolio_scenarios
from serenitas.analytics.base import Trade
from serenitas.analytics.index_data import load_all_curves


def gen_shocks(portf, shock_date, fund):
    Trade.init_ontr(shock_date)
    ontr_spread = Trade._ontr["HY"].spread
    spread_shock = np.array([-25.0, 1.0, +25.0, 100.0, 200.0, 500, 1000])
    spread_shock /= ontr_spread
    # Add in 2020 HY Wides, 2021 HY Tights, 2022 HY Wides scenarios
    historic_spreads = np.array([872, 269, 626])
    spread_shock = np.append(spread_shock, historic_spreads / ontr_spread - 1.0)
    vol_surface = generate_vol_surface(portf, lookback=10, source="BAML")
    scens = run_portfolio_scenarios(
        portf,
        date_range=[pd.Timestamp(shock_date)],
        params=["pnl", "hy_equiv"],
        spread_shock=spread_shock,
        vol_shock=[0.0],
        corr_shock=[0.0],
        vol_surface=vol_surface,
    )
    strategies = {
        s: "options"
        for s in ["HYOPTDEL", "HYPAYER", "HYREC", "IGOPTDEL", "IGPAYER", "IGREC"]
    } | {
        s: "tranches"
        for s in [
            "HYSNR",
            "HYMEZ",
            "HYINX",
            "HYEQY",
            "IGSNR",
            "IGMEZ",
            "IGINX",
            "IGEQY",
            "EUSNR",
            "EUMEZ",
            "EUINX",
            "EUEQY",
            "XOSNR",
            "XOMEZ",
            "XOINX",
            "XOEQY",
            "BSPK",
        ]
    }
    if fund == "BRINKER":
        scens = scens.xs(0, level="corr_shock")
    else:
        scens = scens.xs((0.0, 0.0), level=["vol_shock", "corr_shock"])

    scens.columns.names = ["strategy", "trade_id", "scen_type"]
    results = scens.stack(level="scen_type").reorder_levels([2, 0, 1]).sort_index()
    results = results.groupby(["strategy"], axis=1).sum()
    results = results.groupby(lambda s: strategies.get(s, s), axis=1).sum()
    # map shocks back to absolute spread diff
    results.index = results.index.set_levels(
        results.index.levels[2] * ontr_spread, level="spread_shock"
    )
    results["total"] = results.sum(axis=1)
    results = results.stack().reset_index()
    results.scen_type = results.scen_type.str.upper()
    results.insert(0, "date", results.pop("date"))
    return results


def save_shocks(conn, date, df, fund):
    with conn.cursor() as c:
        c.execute(
            "DELETE FROM shocks WHERE fund=%s AND date=%s",
            (
                fund,
                args.date,
            ),
        )
        conn.commit()
        c.executemany(
            "INSERT INTO shocks VALUES (%s, %s, %s, %s, %s, %s)",
            [(*t, fund) for t in df.itertuples(index=False)],
        )
        conn.commit()


def get_survival_curves(conn, date):
    surv_curves = load_all_curves(conn, date)
    surv_curves["spread"] = surv_curves["curve"].apply(
        lambda sc: [h for d, h in sc][5] * (1 - sc.recovery_rates[5])
    )
    return surv_curves.groupby(level=0).first()[["name", "company_id", "spread"]]


def gen_jtd(portf, survival_curves):
    jtd = portf.jtd_single_names()
    jtd = survival_curves.join(jtd.iloc[:, 0], how="right")
    jtd.columns = ["name", "company_id", "5yr_spread", "jtd"]
    return jtd.groupby(["company_id", "name"], as_index=False).sum()


def save_jtd(conn, date, df, fund):
    with conn.cursor() as c:
        c.execute(
            "DELETE FROM jtd_risks WHERE fund=%s AND date=%s",
            (
                fund,
                date,
            ),
        )
        conn.commit()
        c.executemany(
            "INSERT INTO jtd_risks(date, fund, company_id, name, five_year_spread, jtd) "
            "VALUES (%s, %s, %s, %s, %s, %s)",
            [(date, fund, *t) for t in df.itertuples(index=False)],
        )
        conn.commit()


if __name__ == "__main__":
    import datetime
    import argparse
    import warnings
    from serenitas.analytics.dates import prev_business_day
    from serenitas.utils.db2 import dbconn
    from serenitas.analytics.config import C

    parser = argparse.ArgumentParser(
        description="Shock data/ calculate JTD and insert into DB"
    )
    parser.add_argument(
        "date",
        nargs="?",
        type=datetime.date.fromisoformat,
        default=prev_business_day(datetime.date.today()),
    )
    parser.add_argument("-n", "--no-upload", action="store_true", help="do not upload")
    args = parser.parse_args()

    C.local = False
    C.dual_corr_tranche_cache_size = 16384

    survival_curves = get_survival_curves(Trade._conn, args.date)
    conn = dbconn("dawndb")
    for fund in (
        "SERCGMAST",
        "BOWDST",
        "ISOSEL",
        "BRINKER",
    ):
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore", message="pandas only supports SQLAlchemy connectable"
            )
            warnings.filterwarnings("ignore", message="skipped 1 empty curves")
            portf, _ = build_portfolio(args.date, args.date, fund)
            shocks = gen_shocks(portf, args.date, fund)
            save_shocks(conn, args.date, shocks, fund)
            print(f"{args.date}: {fund} Shocks Done")
            jtd = gen_jtd(portf, survival_curves)
            save_jtd(conn, args.date, jtd, fund)
            print(f"{args.date}: {fund} JTD Done")