import bottleneck as bn
import datetime
import numpy as np
import pandas as pd
import statsmodels.api as sm
import statsmodels.formula.api as smf

from analytics.basket_index import MarkitBasketIndex
from analytics import CreditIndex
from dateutil.relativedelta import relativedelta
from utils.db import dbengine


def get_dispersion(index_type, series, end_date=datetime.date.today()):
    index = MarkitBasketIndex(index_type, series, ["5yr"])
    dr = pd.bdate_range(index.issue_date, end_date)

    dispersion = []
    cumloss = []
    for d in dr:
        print(d)
        index.value_date = d
        dispersion.append(index.dispersion())
        cumloss.append(index.cumloss)

    return pd.DataFrame(
        {"dispersion": dispersion, "cumloss": cumloss,}, index=dr, name="dispersion",
    )


def get_corr_data(index_type, series, engine):
    sql_str = (
        "SELECT quotedate::date, indexrefspread, indexrefprice, index_duration, "
        "index_expected_loss, corr_at_detach "
        "FROM tranche_risk JOIN tranche_quotes "
        "ON tranche_risk.tranche_id=tranche_quotes.id "
        "WHERE index=%s and series=%s and tenor='5yr' and detach=%s order by quotedate desc"
    )
    df = pd.read_sql_query(
        sql_str,
        engine,
        params=(index_type, series, 3 if index_type == "IG" else 15),
        index_col=["quotedate"],
        parse_dates=["quotedate"],
    )
    if index_type == "HY":
        spread_equivalent = []
        index = CreditIndex(index_type, series, "5yr")
        for k, v in df.iterrows():
            index.value_date = k
            index.ref = v["indexrefprice"]
            spread_equivalent.append(index.spread)
        df["indexrefspread"] = spread_equivalent
    df = df.assign(
        fisher=lambda x: 0.5 * np.log((1 + x.corr_at_detach) / (1 - x.corr_at_detach))
    )
    return df


def get_tranche_data(index_type, engine):
    sql_string = (
        "SELECT * FROM risk_numbers "
        "LEFT JOIN index_version USING (index, series, version) "
        "WHERE index = %s"
    )
    df = pd.read_sql_query(
        sql_string, engine, parse_dates={"date": {"utc": True}}, params=[index_type]
    )
    del df["basketid"]
    df.date = df.date.dt.normalize().dt.tz_convert(None)
    df = df.groupby(
        ["date", "index", "series", "version", "tenor", "attach"], as_index=False
    ).mean()
    df = df.assign(
        moneyness=lambda x: np.clip(
            (x.detach - x.cumulativeloss) / x.indexfactor / x.index_expected_loss,
            0.0,
            1.0,
        ),
        exp_percentage=lambda x: x.expected_loss / x.index_expected_loss,
    )
    df.set_index(["index", "series", "tenor", "attach"], append=True, inplace=True)
    return df


def create_gini_models(df):
    # Takes the output of get_tranche_data
    gini_model, gini_calc = {}, {}
    for attach in df.index.get_level_values("attach").unique():
        gini_calc[attach] = df.loc(axis=0)[:, :, :, "5yr", attach]
        gini_model[attach] = smf.ols(
            "np.log(exp_percentage) ~ "
            "np.log(gini_spread) + "
            "np.log(index_duration) + "
            "np.log(moneyness)",
            data=gini_calc[attach],
        ).fit()
        gini_calc[attach]["predict"] = np.exp(
            gini_model[attach].predict(gini_calc[attach])
        )
    gini_calc = pd.concat(gini_calc, sort=False).reset_index(level=0, drop=True)
    normalization = gini_calc.groupby(["date", "index", "series", "tenor"])[
        "predict"
    ].sum()
    gini_calc = gini_calc.merge(
        normalization, left_index=True, right_index=True, suffixes=["_preN", "_sum"]
    )
    gini_calc["predict_N"] = gini_calc["predict_preN"] / gini_calc["predict_sum"]
    gini_calc["mispricing"] = (
        (gini_calc["exp_percentage"] - gini_calc["predict_N"])
        * gini_calc["index_expected_loss"]
        / (gini_calc["detach_adj"] - gini_calc["attach_adj"])
        / gini_calc["indexfactor"]
        * 10000
    )
    return gini_model, gini_calc


def gini(array):
    """Calculate the Gini coefficient of a numpy array."""
    if np.amin(array) < 0:
        array -= np.amin(array)  # values cannot be negative
    array += 0.0000001  # values cannot be 0
    array = np.sort(array)  # values must be sorted
    index = np.arange(1, array.shape[0] + 1)  # index per array element
    n = array.shape[0]  # number of array elements
    return (np.sum((2 * index - n - 1) * array)) / (n * np.sum(array))


def get_gini_spreadstdev(index_type, series, tenor, date):
    indices = MarkitBasketIndex(index_type, series, tenor, value_date=date)
    spreads = indices.spreads()
    spreads = np.ravel(spreads)
    return (gini(spreads), np.std(spreads))


if __name__ == "__main__":
    index_type = "HY"
    series = 29
    serenitas_engine = dbengine("serenitasdb")
    dispersion = get_dispersion(index_type, series)
    df = get_corr_data(index_type, series, serenitas_engine)
    df = df.join(dispersion)

    if index_type == "HY":
        formula = "fisher ~ np.log(dispersion) + cumloss + np.log(index_duration)"
    else:
        formula = "fisher ~ np.log(dispersion) + np.log(indexrefspread) + np.log(index_duration)"
    mod = smf.ols(formula=formula, data=df)