import index_data as data
from dateutil.relativedelta import relativedelta
from pyisda.credit_index import CreditIndex
from typing import List
import pandas as pd
import datetime
from scipy.optimize import brentq

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], *args,
                 trade_date: pd.Timestamp = pd.Timestamp.today().normalize()):
        self.index_type = index_type
        self.series = series
        if index_type == 'IG':
            self.recovery = 0.4
        else:
            self.recovery = 0.3
        self.index_quotes = (data.get_index_quotes(index_type, series,
                                                   tenors, years=None)['closeprice'].
                             unstack().
                             groupby(level='date', as_index=False).nth(0).
                             reset_index(['index', 'series'], drop=True))
        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",
                                            data._serenitas_engine,
                                            index_col='tenor',
                                            params=(index_type, series),
                                            parse_dates=['maturity', 'issue_date'])
        self.index_quotes.columns = self.index_desc.loc[self.index_quotes.columns, "maturity"]
        self.index_quotes = 1 - self.index_quotes / 100
        self.issue_date = self.index_desc.issue_date[0]
        maturities = self.index_quotes.columns.sort_values().to_pydatetime()
        self.index_desc = self.index_desc.reset_index().set_index('maturity')
        curves, args = data.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)

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

    @trade_date.setter
    def trade_date(self, d: pd.Timestamp):
        curves, args = data.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.):
        coupon = self.index_desc.loc[maturity, 'coupon']
        return super().pv(self.step_in_date, self.value_date, maturity, self.yc,
                          self.recovery, coupon, 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:
            coupon = self.index_desc.loc[maturity, 'coupon']
            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, coupon, index_quote)
    def tweak(self):
        """ tweak the singlename curves to match index quotes"""
        quotes = self.index_quotes.loc[self.trade_date]
        self.tweaks = []
        for m, index_quote in quotes.iteritems():
            coupon = self.index_desc.loc[m, 'coupon']
            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)