from .index_data import get_index_quotes, get_singlenames_curves
from .db import _engine
from dateutil.relativedelta import relativedelta
from pyisda.credit_index import CreditIndex
from typing import List
import numpy as np
import pandas as pd
import datetime
from scipy.optimize import brentq
from pandas.tseries.offsets import BDay

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], *,
                 trade_date: pd.Timestamp=pd.Timestamp.today().normalize() - BDay()):
        self.index_type = index_type
        self.series = series
        if index_type == 'IG' or index_type == 'EU':
            self.recovery = 0.4
        else:
            self.recovery = 0.3
        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",
                                            _engine,
                                            index_col='tenor',
                                            params=(index_type, series),
                                            parse_dates=['maturity', 'issue_date'])
        r = _engine.execute("SELECT lastdate, indexfactor/100 AS factor, cumulativeloss/100, version " \
                            "FROM index_version " \
                            "WHERE index = %s AND series = %s" \
                            "ORDER BY lastdate", (index_type, series))
        self._version = tuple(tuple(t) for t in r)
        self.issue_date = self.index_desc.issue_date[0]
        maturities = self.index_desc.maturity.sort_values().dt.to_pydatetime()
        self.index_desc = self.index_desc.reset_index().set_index('maturity')
        curves, args = get_singlenames_curves(index_type, series, trade_date)
        _, jp_yc, _, step_in_date, value_date, _ = args
        self.yc = jp_yc
        self.step_in_date = step_in_date
        self.value_date = value_date
        self._trade_date = trade_date
        self.tweaks = []
        super().__init__(self.issue_date, maturities, curves)

    def _query_version(self, i):
        for lastdate, *data in self._version:
            if lastdate >= self.trade_date.date():
                return data[i]

    @property
    def factor(self):
        return self._query_version(0)

    @property
    def cumloss(self):
        return self._query_version(1)

    @property
    def version(self):
        return self._query_version(2)

    def _get_quotes(self):
        pass

    @property
    def trade_date(self):
        return self._trade_date

    @trade_date.setter
    def trade_date(self, d: pd.Timestamp):
        curves, args = get_singlenames_curves(self.index_type, self.series, d)
        _, jp_yc, _, step_in_date, value_date, _ = args
        self.yc = jp_yc
        self.step_in_date = step_in_date
        self.value_date = value_date
        self._trade_date = d
        self.curves = curves

    def pv(self, maturity: pd.Timestamp, epsilon=0.):
        return super().pv(self.step_in_date, self.value_date, maturity, self.yc,
                          self.recovery, self.coupon(maturity), epsilon)

    def duration(self, maturity):
        return super().duration(self.step_in_date, self.value_date, maturity, self.yc)

    def theta(self, maturity):
        if self.step_in_date.date() > maturity - relativedelta(years=1):
            return np.NaN
        else:
            index_quote = self.index_quotes.loc[self.trade_date, maturity]
            return super().theta(self.step_in_date, self.value_date, maturity,
                                 self.yc, self.recovery, self.coupon(maturity), index_quote)
    def coupon(self, maturity):
        return self.index_desc.loc[maturity, 'coupon']

    def tweak(self):
        """ tweak the singlename curves to match index quotes"""
        quotes = self._get_quotes()
        self.tweaks = []
        for m, index_quote in quotes.items():
            lo, hi = -0.3, 0.3
            while lo > -1:
                try:
                    eps = brentq(lambda epsilon: self.pv(m, epsilon) -
                                 index_quote, lo, hi)
                except ValueError:
                    lo *= 1.1
                    hi *= 1.1
                else:
                    break
            else:
                print("couldn't calibrate for date: {} and maturity: {}".
                      format(self.trade_date.date(), m.date()))
                self.tweaks.append(np.NaN)
                continue
            self.tweaks.append(eps)
            self.tweak_portfolio(eps, m)

class MarkitBasketIndex(BasketIndex):
    def __init__(self, index_type: str, series: int, tenors: List[str], *,
                 trade_date: pd.Timestamp=pd.Timestamp.today().normalize() - BDay()):
        super().__init__(index_type, series, tenors, trade_date=trade_date)
        self.index_quotes = (get_index_quotes(index_type, series,
                                              tenors, years=None)['closeprice'].
                             unstack().
                             groupby(level='date', as_index=False).nth(0).
                             reset_index(['index', 'series', 'version'], drop=True))
        self.index_quotes.columns = (self.index_desc.reset_index().
                                     set_index('tenor').
                                     loc[self.index_quotes.columns, "maturity"])
        self.index_quotes = 1 - self.index_quotes / 100

    def _get_quotes(self):
        return self.index_quotes.loc[self.trade_date]

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(cds_schedule.to_npdates().view('int') + 134774)