from .index_data import get_index_quotes, get_singlenames_curves
from . import serenitas_engine
from .utils import tenor_t, get_fx
from functools import partial
from pyisda.cdsone import upfront_charge, spread_from_upfront
from pyisda.credit_index import CreditIndex
from pyisda.date import previous_twentieth
from typing import List
from yieldcurve import get_curve
import datetime
import logging
import numpy as np
import pandas as pd
from math import exp
from scipy.optimize import brentq
from pandas.tseries.offsets import Day, BDay

logger = logging.getLogger(__name__)


def make_index(t, d, args):
    instance = t.__new__(t)
    CreditIndex.__init__(instance, *args)
    instance.__dict__.update(d)
    return instance


class BasketIndex(CreditIndex):
    index_type: str
    series: int
    recovery: float
    step_in_date: pd.Timestamp
    value_date: pd.Timestamp
    tweaks: List[float]

    def __init__(
        self,
        index_type: str,
        series: int,
        tenors: List[str],
        *,
        value_date: pd.Timestamp = pd.Timestamp.today().normalize() - BDay(),
    ):
        self.index_type = index_type
        self.series = series
        if index_type in ("HY", "HY.BB"):
            self.recovery = 0.3
        else:
            self.recovery = 0.4
        self.index_desc = pd.read_sql_query(
            "SELECT tenor, maturity, coupon * 1e-4 AS coupon, "
            "issue_date "
            "FROM index_maturity "
            "WHERE index=%s AND series=%s",
            serenitas_engine,
            index_col="tenor",
            params=(index_type, series),
            parse_dates=["maturity", "issue_date"],
        )
        if self.index_desc.empty:
            raise ValueError(f"Index {index_type} {series} doesn't exist")
        self._index_version = tuple(
            tuple(r.values())
            for r in serenitas_engine.execute(
                "SELECT lastdate,"
                "       indexfactor/100 AS factor,"
                "       cumulativeloss/100 AS cum_loss,"
                "       version "
                "FROM index_version "
                "WHERE index = %s AND series = %s"
                "ORDER BY lastdate",
                (index_type, series),
            )
        )
        self._update_factor(value_date)
        self.issue_date = self.index_desc.issue_date[0]
        self.index_desc = self.index_desc.loc[tenors].sort_values("maturity")
        self.tenors = {t: m.date() for t, m in self.index_desc.maturity.items()}
        maturities = self.index_desc.maturity.dt.to_pydatetime()
        self.index_desc = self.index_desc.reset_index().set_index("maturity")
        self.index_desc.tenor = self.index_desc.tenor.astype(tenor_t)
        max_tenor = int(tenors[-1][:-2])
        self._curve_tenors = tuple(
            t for t in (0.5, 1, 2, 3, 4, 5, 7, 10) if t <= max_tenor
        )
        curves = get_singlenames_curves(
            index_type, series, value_date, self._curve_tenors
        )

        self.currency = "EUR" if index_type in ["XO", "EU"] else "USD"
        self.yc = get_curve(value_date, self.currency)
        self._fx = get_fx(value_date, self.currency)
        self.step_in_date = value_date + Day()
        self.cash_settle_date = value_date + 3 * BDay()
        self.tweaks = []
        self.start_date = previous_twentieth(value_date)
        self._ignore_hash = set(
            [
                "_Z",
                "_w",
                "_skew",
                "tenors",
                "index_desc",
                "tweaks",
                "_Legs",
                "_ignore_hash",
            ]
        )
        super().__init__(self.issue_date, maturities, curves, value_date=value_date)

    def __reduce__(self):
        _, args = CreditIndex.__reduce__(self)
        d = vars(self)
        return partial(make_index, self.__class__), (d, args)

    def __hash__(self):
        def aux(v):
            if isinstance(v, list):
                return hash(tuple(v))
            elif type(v) is np.ndarray:
                return hash(v.tobytes())
            else:
                return hash(v)

        return hash(
            (CreditIndex.__hash__(self),)
            + tuple(aux(v) for k, v in vars(self).items() if k not in self._ignore_hash)
        )

    def _update_factor(self, d):
        if isinstance(d, datetime.datetime):
            d = d.date()
        for lastdate, *data in self._index_version:
            if lastdate >= d:
                self._factor, self._cumloss, self._version = data
                self._lastdate = lastdate
                break

    @property
    def factor(self):
        return self._factor

    @property
    def cumloss(self):
        return self._cumloss

    @property
    def version(self):
        return self._version

    def _get_quotes(self, *args):
        """ allow to tweak based on manually inputed quotes"""
        if self.index_type == "HY":
            return {m: (100 - p) / 100 for m, p in zip(self.maturities, args[0])}
        else:
            return {
                m: self._snacpv(s * 1e-4, self.coupon(m), self.recovery, m)
                for m, s in zip(self.maturities, args[0])
            }

    value_date = property(CreditIndex.value_date.__get__)

    @value_date.setter
    def value_date(self, d: pd.Timestamp):
        self.curves = get_singlenames_curves(
            self.index_type, self.series, d, self._curve_tenors
        )
        self.yc = get_curve(d, self.currency)
        self._fx = get_fx(d, self.currency)
        self.step_in_date = d + Day()
        self.cash_settle_date = d + 3 * BDay()
        self.start_date = previous_twentieth(d)  # or d + 1?
        self._update_factor(d)
        CreditIndex.value_date.__set__(self, d)

    @property
    def recovery_rates(self):
        # we don't always have the 6 months data point
        # so pick arbitrarily the 1 year point
        return np.array([c.recovery_rates[0] for _, c in self.curves])

    def spreads(self):
        return super().spreads(self.yc)

    def dispersion(self, use_gini: bool = False, use_log: bool = True):
        if use_gini:
            w = self.weights
            spreads = self.spreads().ravel()
            if use_log:
                spreads = np.log(spreads)
            index = np.argsort(spreads)
            spreads = spreads[index]
            w = w[index]
            S = np.cumsum(w * spreads)
            return 1 - (np.inner(w[1:], (S[:-1] + S[1:])) + w[0] * S[0]) / S[-1]
        else:
            return super().dispersion(self.yc, use_log=use_log)

    def pv(self, maturity=None, epsilon=0.0, coupon=None):
        if maturity is None:
            r = []
            for m in self.maturities:
                coupon = self.index_desc.coupon[m]
                r.append(
                    super().pv(
                        self.step_in_date,
                        self.cash_settle_date,
                        m,
                        self.yc,
                        coupon,
                        epsilon,
                    )
                )
            return pd.Series(r, index=self.index_desc.tenor, name="pv")
        else:
            return super().pv(
                self.step_in_date,
                self.cash_settle_date,
                maturity,
                self.yc,
                coupon or self.coupon(maturity),
                epsilon,
            )

    def pv_vec(self):
        return (
            super().pv_vec(self.step_in_date, self.cash_settle_date, self.yc).unstack(0)
        )

    def coupon_leg(self, maturity=None):
        return self.index_desc.coupon.values * self.duration()

    def spread(self, maturity=None):
        return self.protection_leg(maturity) / self.duration(maturity) * 1e4

    def protection_leg(self, maturity=None):
        if maturity is None:
            r = []
            for m in self.maturities:
                r.append(
                    super().protection_leg(
                        self.step_in_date, self.cash_settle_date, m, self.yc
                    )
                )
            return pd.Series(r, index=self.index_desc.tenor, name="protection_leg")
        else:
            return super().protection_leg(
                self.step_in_date, self.cash_settle_date, maturity, self.yc
            )

    def duration(self, maturity=None):
        if maturity is None:
            r = []
            for m in self.maturities:
                r.append(
                    super().duration(
                        self.step_in_date, self.cash_settle_date, m, self.yc
                    )
                )
            return pd.Series(r, index=self.index_desc.tenor, name="duration")
        else:
            return super().duration(
                self.step_in_date, self.cash_settle_date, maturity, self.yc
            )

    def theta(self, maturity=None, coupon=None, theta_date=None):
        """ index thetas

        if maturity is None, returns a series of theta for all tenors.
        Otherwise computes the theta for that specific maturity (which needs
        not be an existing tenor)

        if theta_date is provided, computes the theta to that specific date
        instead of one-year theta"""
        try:
            index_quotes = self._get_quotes()
        except (ValueError, IndexError):
            index_quotes = {}
        if maturity is None:
            r = []
            for m in self.maturities:
                coupon = self.index_desc.coupon[m]
                index_quote = index_quotes.get(m, np.nan)
                r.append(
                    super().theta(
                        self.step_in_date,
                        self.cash_settle_date,
                        m,
                        self.yc,
                        coupon,
                        index_quote,
                        theta_date,
                    )
                )
            return pd.Series(r, index=self.index_desc.tenor, name="theta")
        else:
            return super().theta(
                self.step_in_date,
                self.cash_settle_date,
                maturity,
                self.yc,
                coupon or self.coupon(maturity),
                np.nan,
                theta_date,
            )

    def coupon(self, maturity=None, assume_flat=True):
        if maturity is None:
            return self.index_desc.set_index("tenor").coupon
        else:
            try:
                return self.index_desc.coupon[maturity]
            except KeyError:
                if assume_flat:
                    return self.index_desc.coupon.iat[0]
                else:
                    raise ValueError("Non standard maturity: coupon must be provided")

    def tweak(self, *args):
        """ tweak the singlename curves to match index quotes"""
        quotes = self._get_quotes(*args)
        self.tweaks = []

        for m in self.maturities:
            if m not in quotes:
                self.tweaks.append(np.nan)
                continue
            else:
                index_quote = quotes[m]
                if abs(self.pv(m) - index_quote) < 1e-12:  # early exit
                    self.tweaks.append(0.0)
                    continue
            lo, hi = -0.3, 0.3
            hi_tilde = exp(hi) - 1
            while hi_tilde < 5:
                # map range to (-1, +inf)
                lo_tilde = exp(lo) - 1
                hi_tilde = exp(hi) - 1
                try:
                    eps = brentq(
                        lambda epsilon: self.pv(m, epsilon) - index_quote,
                        lo_tilde,
                        hi_tilde,
                    )
                except ValueError:
                    lo *= 1.1
                    hi *= 1.1
                else:
                    break
            else:
                logger.warning(
                    f"couldn't calibrate for date: {self.value_date} and maturity: {m}"
                )
                self.tweaks.append(np.NaN)
                continue
            self.tweaks.append(eps)
            self.tweak_portfolio(eps, m)

    def _snacpv(self, spread, coupon, recov, maturity):
        return upfront_charge(
            self.value_date,
            self.cash_settle_date,
            self.start_date,
            self.step_in_date,
            self.start_date,
            maturity,
            coupon,
            self.yc,
            spread,
            recov,
        )

    def _snacspread(self, coupon, recov, maturity):
        return spread_from_upfront(
            self.value_date,
            self.cash_settle_date,
            self.start_date,
            self.step_in_date,
            self.start_date,
            maturity,
            coupon,
            self.yc,
            self.pv(maturity),
            recov,
        )


class MarkitBasketIndex(BasketIndex):
    def __init__(
        self,
        index_type: str,
        series: int,
        tenors: List[str],
        *,
        value_date: pd.Timestamp = pd.Timestamp.today().normalize() - BDay(),
    ):
        super().__init__(index_type, series, tenors, value_date=value_date)
        self.index_quotes = (
            get_index_quotes(
                index_type, series, tenors, years=None, remove_holidays=False
            )[["close_price", "id"]]
            .groupby(level=["date", "tenor"], as_index=True)
            .nth(0)
        )
        self.index_quotes.close_price = 1 - self.index_quotes.close_price / 100

    def _get_quotes(self):
        quotes = self.index_quotes.loc[self.value_date, "close_price"]
        return {self.tenors[t]: q for t, q in quotes.items()}


if __name__ == "__main__":
    ig28 = BasketIndex("IG", 28, ["3yr", "5yr", "7yr", "10yr"])
    from quantlib.time.api import Schedule, Rule, Date, Period, WeekendsOnly
    from quantlib.settings import Settings

    settings = Settings()

    cds_schedule = Schedule.from_rule(
        settings.evaluation_date,
        Date.from_datetime(ig28.maturities[-1]),
        Period("3M"),
        WeekendsOnly(),
        date_generation_rule=Rule.CDS2015,
    )

    sp = ig28.survival_matrix()