import numpy as np
import logging
import pandas as pd

from serenitas.analytics.base import Trade
from serenitas.analytics.index_data import on_the_run
from serenitas.analytics.api import CreditIndex, BlackSwaptionVolSurface
from serenitas.analytics.exceptions import MissingDataError
from copy import deepcopy
from .tranches import get_tranche_portfolio
from .swaptions import get_swaption_portfolio
from .bonds import subprime_risk, clo_risk, crt_risk
from .indices import get_index_portfolio
from serenitas.utils.db import dbengine
from serenitas.utils.db2 import dbconn
from serenitas.utils.pool import dawn_pool
from pandas.tseries.offsets import BDay

logger = logging.getLogger(__name__)


def hy_equiv_trade(value_date, notional):
    return CreditIndex(
        "HY",
        on_the_run("HY", value_date),
        "5yr",
        value_date=value_date,
        notional=notional,
    )


def gen_bond_proxies(position_date, fund, conn):
    rmbs_pos = subprime_risk(position_date, conn, dbengine("rmbs_model"), fund=fund)
    crt_pos = crt_risk(position_date, conn, fund=fund)
    with dbconn("etdb") as etconn:
        clo_pos = clo_risk(position_date, conn, etconn, fund=fund)
    return {"rmbs_bonds": rmbs_pos, "crt_bonds": crt_pos, "clo_bonds": clo_pos}


def build_portfolio(position_date, value_date, fund="SERCGMAST"):
    """
    Output two portfolios:
    1) All synthetic + curve with just delta-proxy + dummy index as cash bonds proxy (portf)
    2) Tranches +Swaption + Index for JTD
    """

    Trade.init_ontr(value_date)
    with dawn_pool.connection() as conn:
        portf = get_tranche_portfolio(position_date, conn, False, [fund])

        swaption_portf = get_swaption_portfolio(position_date, conn, fund=fund)
        swaption_portf.trade_ids = [tid[:-1] for tid in swaption_portf.trade_ids]
        portf += swaption_portf

        syn_portf = deepcopy(portf)
        syn_portf += get_index_portfolio(position_date, conn, fund)

        nocurve_portf = get_index_portfolio(
            position_date, conn, fund, exclude_strategies="%CURVE"
        )
        portf += nocurve_portf
        basis_portf = get_index_portfolio(
            position_date, conn, fund, include_strategies="%BASIS"
        )
        portf.add_trade(
            hy_equiv_trade(value_date, -basis_portf.hy_equiv),
            ("CASH_BASIS", "negate_basis_trades"),
        )

        curve_portf = get_index_portfolio(
            position_date, conn, fund, include_strategies="%CURVE"
        )
        curve_portf.value_date = value_date
        curve_portf.mark()
        portf.add_trade(
            hy_equiv_trade(value_date, curve_portf.hy_equiv),
            ("curve_trades", "curve_trades"),
        )
        # Add the bond proxies to the portfolio only
        for name, pos in gen_bond_proxies(position_date, fund, conn).items():
            portf.add_trade(
                hy_equiv_trade(value_date, -pos.hy_equiv.sum()), (name, name)
            )

    for p in [portf, syn_portf]:
        p.value_date = value_date
        p.mark(interp_method="bivariate_linear")
        p.reset_pv()

    return portf, syn_portf


def generate_vol_surface(portf, lookback=5, source="MS"):
    vol_surface = {}

    for trade in portf.swaptions:
        index_type, series, option_type = (
            trade.index.index_type,
            trade.index.series,
            trade.option_type,
        )
        if (k := (index_type, series, option_type)) in vol_surface:
            continue
        i = 0
        while i < lookback:
            try:
                vs = BlackSwaptionVolSurface(
                    index_type,
                    series,
                    value_date=portf.value_date - BDay(i),
                    interp_method="bivariate_linear",
                )
            except MissingDataError as e:
                i += 1
                logger.info(f"Trying {portf.value_date - BDay(i)}")
            else:
                break
        vol_surface[k] = vs[vs.list(source=source, option_type=option_type)[-1]]

    return vol_surface