from bbg_helpers import BBG_IP, init_bbg_session, retrieve_data
from analytics.utils import roll_date
from psycopg2.extras import execute_values

import datetime
import pandas as pd


def bbg_name(index_type: str, series: int, tenor="5yr", version=None) -> str:
    if index_type == "EU":
        index_type = "EUR"
        index_name = "ITRX"
    elif index_type == "XO":
        index_type = "XOVER"
        index_name = "ITRX"
    elif index_type == "HYBB":
        index_name = ""
    else:
        index_name = "CDX"
    return f"{index_name} {index_type} CDSI S{series} {f'V{version} ' if version else ''}{tenor[:2].upper()} Corp".lstrip()


def new_basket_id(
    conn,
    index_type,
    series,
    version=1,
    mat_5yr=roll_date(datetime.date.today(), 5),
    coupon=100,
    red=None,
):
    if red is None:
        name = bbg_name(index_type, series, version=version)
        with init_bbg_session(BBG_IP) as session:
            data = retrieve_data(session, [name], ["CDS_RED_PAIR_CODE"])
            red = data[name]["CDS_RED_PAIR_CODE"]
    sql_str = (
        "INSERT INTO index_version(index, series, version, redindexcode) "
        "VALUES(%s, %s, %s, %s) RETURNING basketid"
    )
    with conn.cursor() as c:
        c.execute(sql_str, (index_type, series, version, red))
        (basket_id,) = c.fetchone()
    maturities = [mat_5yr.replace(year=mat_5yr.year + off) for off in [-2, 0, 2, 5]]
    with conn.cursor() as c:
        execute_values(
            c,
            "INSERT INTO index_maturity VALUES %s",
            [
                (index_type, series, t, m, coupon)
                for t, m in zip(["3yr", "5yr", "7yr", "10yr"], maturities)
            ],
        )
    conn.commit()
    return basket_id


def get_index_members(index_type: str, series: int) -> pd.DataFrame:
    name = bbg_name(index_type, series)
    with init_bbg_session(BBG_IP) as session:
        df = retrieve_data(session, [name], ["INDX_MEMBERS"])
        df = df[name]["INDX_MEMBERS"]
        tickers = df["column   4"]
        company_ids = retrieve_data(
            session, [t + " Corp" for t in tickers], ["cds_company_id"]
        )
    company_ids = pd.DataFrame.from_dict(company_ids, orient="index")
    company_ids = company_ids.reset_index()
    company_ids["index"] = company_ids["index"].str.replace(" Corp", "")
    company_ids = company_ids.set_index("index")
    df = df.set_index("column   4").join(company_ids)[["cds_company_id", "column   3"]]
    df.columns = ["company_id", "seniority"]
    return df


def new_index_series(conn, basketid, index_type, series):
    df = get_index_members(index_type, series)
    with conn.cursor() as c:
        c.executemany(
            "INSERT INTO basket_constituents VALUES(%s, %s, %s, %s)",
            [
                (
                    t.company_id,
                    "SLA" if t.seniority == "LAC" else t.seniority,
                    basketid,
                    1.0,
                )
                for t in df.itertuples(index=False)
            ],
        )
    conn.commit()