import datetime
import pandas as pd
import numpy as np
import argparse

from risk.portfolio import build_portfolio, generate_vol_surface

import serenitas.analytics as ana
from serenitas.analytics.scenarios import run_portfolio_scenarios
from serenitas.analytics.base import Trade
from serenitas.utils.db2 import dbconn
from serenitas.utils.db import dawn_engine
from serenitas.analytics.dates import prev_business_day


def parse_args():
    """Parses command line arguments"""
    parser = argparse.ArgumentParser(description="Shock data 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")
    return parser.parse_args()


def vol_surface(portf, try_days_back):
    for source in ("BAML", "GS", "MS", "JPM"):
        try:
            vol_surface = generate_vol_surface(portf, 10, source)
        except IndexError:
            pass
        else:
            return vol_surface


def gen_spreads(shock_date, fund):
    Trade.init_ontr(shock_date)
    ana._local = False
    spread_shock = np.array([-100.0, -25.0, 1.0, +25.0, 100.0, 200.0, 500, 1000])
    spread_shock /= Trade._ontr["HY"].spread
    portf, _ = build_portfolio(shock_date, shock_date, fund)
    vol_surface = generate_vol_surface(portf, 10)
    portf.reset_pv()
    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 = {}
    strategies["options"] = [
        "HYOPTDEL",
        "HYPAYER",
        "HYREC",
        "IGOPTDEL",
        "IGPAYER",
        "IGREC",
    ]
    strategies["tranches"] = [
        "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 = {}
    for i, g in scens.groupby(level="scen_type", axis=1):
        temp = g.groupby(level="strategy", axis=1).sum()
        for key, item in strategies.items():
            exist_columns = list(set(temp.columns).intersection(item))
            temp[key] = temp[exist_columns].sum(axis=1)
            temp.drop(exist_columns, axis=1, inplace=True)
        temp["total"] = temp.sum(axis=1)
        results[i] = temp
    results = pd.concat(results)
    return results


def process_dataframe(raw_df):
    """Clean and transform the input dataframe to insert into database."""
    transformed_df = raw_df.reset_index()
    transformed_df = transformed_df.rename(columns={"level_0": "unit"})
    strategy_columns = transformed_df.columns[3:]
    transformed_df = pd.melt(
        transformed_df,
        id_vars=["unit", "date", "spread_shock"],
        value_vars=strategy_columns,
    )
    transformed_df = transformed_df.rename(
        columns={"variable": "strategy", "value": "value", "unit": "risk_type"}
    )
    transformed_df.risk_type = transformed_df.risk_type.str.upper()
    return transformed_df


if __name__ == "__main__":
    args = parse_args()
    conn = dbconn("dawndb")
    for fund in ("SERCGMAST", "BOWDST", "ISOSEL", "BRINKER"):
        results = gen_spreads(args.date, fund)
        with conn.cursor() as c:
            c.execute(
                "DELETE FROM shocks WHERE fund=%s AND date=%s",
                (
                    fund,
                    args.date,
                ),
            )
        df = process_dataframe(results)
        df["fund"] = fund
        df.to_sql("shocks", dawn_engine, if_exists="append", index=False)
        conn.commit()