import pandas as pd
import numpy as np

from enum import Enum, auto
from utils.db import dbengine
from yieldcurve import YC
from quantlib.termstructures.yield_term_structure import YieldTermStructure


class AssetClass(Enum):
    Subprime = auto()
    CLO = auto()
    CSO = auto()
    CRT = auto()


def latest_sim(date, engine):
    sql_string = (
        "SELECT model_id_sub FROM model_versions "
        "JOIN model_versions_nonagency USING (model_id_sub) "
        "JOIN simulations_nonagency USING (simulation_id) "
        "WHERE (date(start_time) <= %s) AND (description = 'normal') "
        "ORDER BY start_time DESC"
    )
    conn = engine.raw_connection()
    c = conn.cursor()
    c.execute(sql_string, (date,))
    model_id_sub, = next(c)
    c.close()
    return model_id_sub


def get_df(date, engine):
    model_id_sub = latest_sim(date, engine)
    df_prices = pd.read_sql_query(
        "SELECT cusip, model_version, pv, modDur, delta_yield, "
        "wal, pv_io, pv_po, pv_RnW, delta_ir_io, delta_ir_po, "
        "delta_hpi, delta_RnW, delta_mult, delta_ir, pv_FB "
        "FROM priced WHERE "
        "timestamp BETWEEN %s  AND date_add(%s, INTERVAL 1 DAY) "
        "AND model_id_sub=%s "
        "AND normalization='current_notional'",
        engine,
        ["cusip", "model_version"],
        params=(date, date, model_id_sub),
    )
    df_percentiles = pd.read_sql_query(
        "SELECT cusip, PV, percentile "
        "FROM priced_percentiles WHERE "
        "timestamp BETWEEN %s AND date_add(%s, INTERVAL 1 DAY) "
        "AND model_version=3 "
        "AND model_id_sub=%s "
        "AND percentile IN (5, 25, 50, 75, 95) "
        "AND normalization='current_notional'",
        engine,
        ["cusip", "percentile"],
        params=(date, date, model_id_sub),
    )
    df_prices = df_prices.unstack("model_version")
    df_percentiles = df_percentiles.unstack("percentile")
    return df_prices.join(df_percentiles, how="left")


def subprime_risk(date, conn, engine):
    df = get_df(date, engine)
    df_pos = get_portfolio(date, conn, AssetClass.Subprime)
    df_pv = df.xs("pv", axis=1, level=0)
    df_pv.columns = ["pv1", "pv2", "pv3"]
    df_pv_perct = df.xs("PV", axis=1, level=0)
    df_pv_perct.columns = ["pv5", "pv25", "pv50", "pv75", "pv95"]
    df_modDur = df[("modDur", 1)]
    df_modDur.name = "modDur"
    df_v1 = df.xs(1, axis=1, level="model_version")[
        ["pv_RnW", "delta_mult", "delta_hpi", "delta_ir"]
    ]
    df_v1.columns = ["v1pv_RnW", "v1_lsdel", "v1_hpidel", "v1_irdel"]
    df_pv_FB = df[("pv_FB", 3)]
    df_pv_FB.name = "pv_FB"
    df_risk = pd.concat(
        [
            df_pv,
            df_modDur,
            df_pv_perct,
            df.xs(3, axis=1, level="model_version")[
                [
                    "delta_yield",
                    "wal",
                    "pv_io",
                    "pv_po",
                    "pv_RnW",
                    "delta_ir_io",
                    "delta_ir_po",
                    "delta_hpi",
                    "delta_RnW",
                    "delta_mult",
                ]
            ],
            df_v1,
            df_pv_FB,
        ],
        axis=1,
    )

    df_calc = df_pos.join(df_risk)
    yc = YC(evaluation_date=date)

    df_calc = df_calc.assign(
        b_yield=df_calc.modDur.apply(lambda x: float(yc.zero_rate(x))),
        delta_ir=df_calc.delta_ir_io + df_calc.delta_ir_po,
        curr_ntl=df_calc.notional * df_calc.factor,
    )

    df_calc.b_yield += (
        (df_calc.pv1 * df_calc.curr_ntl / df_calc.local_market_value).log()
        / df_calc.modDur
    ).clip_upper(1.0)
    # delta scaled by ratio of market_value to model value
    df_calc.delta_yield *= df_calc.local_market_value / df_calc.pv3
    df_calc.delta_ir *= df_calc.local_market_value / df_calc.pv3
    return df_calc


def get_portfolio(date, conn, asset_class: AssetClass, fund="SERCGMAST"):
    df = pd.read_sql_query(
        "SELECT * FROM risk_positions(%s, %s, %s)",
        conn,
        params=(date, asset_class.name, fund),
    )
    if asset_class is AssetClass.CLO:
        with conn.cursor() as c:
            c.execute(
                "SELECT cusip, identifier FROM securities " "WHERE asset_class = 'CLO'"
            )
            cusip_map = {identifier: cusip for cusip, identifier in c.fetchall()}
        df["cusip"] = df["identifier"].replace(cusip_map)
    else:  # only CLOs used ISIN for now
        df["cusip"] = df.identifier.str.slice(0, 9)
    return df.set_index("cusip")


def crt_risk(date, dawn_conn, engine):
    df = get_portfolio(date, dawn_conn, AssetClass.CRT)
    df_model = pd.read_sql_query(
        "SELECT * from priced_at_market WHERE "
        "timestamp BETWEEN %s AND date_add(%s, INTERVAL 1 DAY) "
        "AND model_des = 'hpi3_ir3'",
        engine,
        "cusip",
        params=(date, date),
    )
    df = df.join(df_model)
    df["curr_ntl"] = df["notional"] * df["factor"]
    df["delta_yield"] = df["curr_ntl"] * df["duration_FW"]
    return df


def clo_risk(date, dawn_conn, et_conn):
    df = get_portfolio(date, dawn_conn, AssetClass.CLO)
    placeholders = ",".join(["%s"] * df.shape[0])
    sql_string = f"SELECT * FROM historical_cusip_risk(%s, {placeholders})"
    model = pd.read_sql_query(
        sql_string, et_conn, parse_dates=["pricingdate"], params=(date, *df.index)
    )
    model.index = df.index
    df = df.join(model, lsuffix="mark")
    df["curr_ntl"] = df["notional"] * df["factor"]
    df["hy_equiv"] = df["curr_ntl"] * df["delta"]
    return df