import datetime
import pandas as pd
from analytics.index import CreditIndex
from db import serenitas_pool
from statistics import median


def get_refids(
    conn,
    index,
    series,
    expiry,
    value_date=datetime.date.today(),
    sources=["GS", "MS", "CITI"],
):
    sql_str = (
        "SELECT ref_id, ref, quotedate FROM swaption_ref_quotes "
        "WHERE quotedate::date=%s "
        "  AND quote_source=%s "
        "  AND index=%s AND series=%s"
        "  AND expiry=%s "
        "ORDER BY quotedate DESC LIMIT 1"
    )
    d = {}
    with conn.cursor() as c:
        for s in sources:
            c.execute(sql_str, (value_date, s, index, series, expiry))
            d[s] = c.fetchone()
    return d


def adjust_stacks(
    index_type,
    series,
    expiry,
    value_date=datetime.date.today(),
    sources=["GS", "MS", "CITI"],
    common_ref=None,
):
    conn = serenitas_pool.getconn()
    d = get_refids(conn, index_type, series, expiry, value_date, sources)
    if all(v is None for v in d.values()):
        raise ValueError("no quotes")
    if common_ref is None:
        common_ref = median(v[1] for v in d.values())
    index = CreditIndex(
        index_type, series, "5yr", value_date=value_date, notional=10000.0
    )
    index.ref = common_ref
    old_pv = index.pv
    quotes = {}
    for s, (ref_id, ref, _) in d.items():
        index.ref = ref
        dindex_pv = index.pv - old_pv
        df = pd.read_sql_query(
            "SELECT strike, pay_bid, pay_offer, delta_pay, "
            "rec_bid, rec_offer, delta_rec FROM swaption_quotes "
            "WHERE ref_id=%s ORDER BY strike",
            conn,
            params=(ref_id,),
            index_col=["strike"],
        )
        if s == "GS":
            df["delta_rec"] = 1 - df["delta_pay"]
            if index_type == "HY":
                df[["pay_bid", "pay_offer", "rec_bid", "rec_offer"]] *= 100
        if s == "CITI":
            df["delta_rec"] *= -1
        if dindex_pv != 0.0:
            df[["pay_bid", "pay_offer"]] = df[["pay_bid", "pay_offer"]].sub(
                df.delta_pay * dindex_pv, axis=0
            )
            df[["rec_bid", "rec_offer"]] = df[["rec_bid", "rec_offer"]].add(
                df.delta_rec * dindex_pv, axis=0
            )
        quotes[s] = df
    quotes = pd.concat(quotes, names=["source"])
    quotes = quotes.swaplevel("source", "strike").sort_index()
    inside_quotes = pd.concat(
        [
            quotes[["pay_bid", "rec_bid"]].groupby(level="strike").max(),
            quotes[["pay_offer", "rec_offer"]].groupby(level="strike").min(),
        ],
        axis=1,
    ).sort_index(axis=1)
    quotes = quotes.unstack("source")
    d = {}
    for k in ["pay_bid", "rec_bid"]:
        # quotes[k].style.apply(highlight_max, axis=1)
        df = pd.concat([quotes[k], inside_quotes[k]], axis=1)
        d[k] = df.rename(columns={k: "Best"})
    for k in ["pay_offer", "rec_offer"]:
        # quotes[k].style.apply(highlight_min, axis=1)
        df = pd.concat([inside_quotes[k], quotes[k]], axis=1)
        d[k] = df.rename(columns={k: "Best"})
    serenitas_pool.putconn(conn)
    return common_ref, pd.concat(d, axis=1)