from serenitas.analytics.bbg_helpers import init_bbg_session, retrieve_data
from serenitas.analytics.dates 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 = ""
    elif index_type == "EUFS":
        index_name = ""
        index_type = "SNRFIN"
    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,
    session,
    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)
        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(session, index_type: str, series: int) -> pd.DataFrame:
    name = bbg_name(index_type, series)
    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.rename(columns={"column   4": "ticker"})
    df = df.set_index("ticker").join(company_ids)
    df = df.drop("column   5", axis=1)
    df.columns = ["name", "weight", "currency", "seniority", "company_id"]
    df.seniority = df.seniority.replace("LAC", "SLA")
    return df


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


def add_missing_issuers(conn, session, df):
    all_company_ids = pd.read_sql_query(
        "SELECT company_id, seniority FROM bbg_issuers",
        conn,
        index_col=["company_id", "seniority"],
    )
    df = df.reset_index().set_index(["company_id", "seniority"])
    df = df.loc[df.index.difference(all_company_ids.index)]
    if df.empty:
        return []
    tickers = [f"{t} Corp" for t in df.ticker]
    bbg_data = retrieve_data(
        session,
        tickers,
        [
            "cds_restructuring_type",
            "cds_corp_tkr",
            "restructuring_type_short_code",
            "cds_curve_info",
        ],
    )
    sql_str = (
        "INSERT INTO bbg_issuers "
        "(name, company_id, ticker, currency, seniority, doc_clause, cds_curve, short_code)"
        "VALUES(%s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING"
    )
    df = df.reset_index()
    with conn.cursor() as c:
        for r in df.itertuples(index=False):
            data = bbg_data[f"{r.ticker} Corp"]
            c.execute(
                sql_str,
                (
                    r.name,
                    r.company_id,
                    data["cds_corp_tkr"],
                    r.currency,
                    r.seniority,
                    data["cds_restructuring_type"].replace("-", " "),
                    data["cds_curve_info"].Ticker.to_list(),
                    data["restructuring_type_short_code"],
                ),
            )
    conn.commit()
    return df


def markit_mappings(conn, df):
    sql_str = "INSERT INTO bbg_markit_mapping VALUES('infinity', %s, %s, %s, %s, %s)"
    with conn.cursor() as c:
        for r in df.itertuples(index=False):
            c.execute(sql_str, r)
    conn.commit()


if __name__ == "__main__":
    from serenitas.utils.db import serenitas_pool

    conn = serenitas_pool.getconn()
    with init_bbg_session() as session:
        df = get_index_members(session, "EUFS", 36)
        missing = add_missing_issuers(conn, session, df)
        basketid = new_basket_id(conn, session, "EUFS", 36, coupon=100)
    new_index_series(conn, basketid, df)
    serenitas_pool.putconn(conn)