from .index import Index
from .option import BlackSwaption, VolatilitySurface
from db import dbengine
from warnings import warn
import pandas as pd
import numpy as np

serenitasdb = dbengine('serenitasdb')

def portf_repr(method):
    def f(*args):
        obj = args[0]
        thousands = lambda x: "{:,.2f}".format(x)
        percent = lambda x: "N/A" if np.isnan(x) else f"{100*x:.2f}%"
        header = "Portfolio {}\n\n".format(obj.trade_date)
        kwargs = {'formatters': {'Notional': thousands,
                                 'PV': thousands,
                                 'Delta': percent,
                                 'Gamma': percent,
                                 'Theta': thousands,
                                 'Vega': thousands,
                                 'Vol': percent,
                                 'Ref': thousands},
                  'index': False}
        if method == 'string':
            kwargs['line_width'] = 100
        s = getattr(obj._todf(), 'to_' + method)(**kwargs)
        return header + s
    return f

class Portfolio:
    def __init__(self, trades):
        self.trades = trades
        self.indices = [t for t in trades if isinstance(t, Index)]
        self.swaptions = [t for t in trades if isinstance(t, BlackSwaption)]
        trade_dates = set(index.trade_date for index in self.indices)
        self._keys = [(index.index_type, index.series, index.tenor) for index in self.indices]
        for swaption in self.swaptions:
            trade_dates.add(swaption.index.trade_date)
        self._trade_date = trade_dates.pop()
        if len(trade_dates) >= 1:
            warn("not all instruments have the same trade date, picking {}".
                 format(self._trade_date))
        self._vs = {}

    @property
    def pnl(self):
        return sum(t.pnl for t in self.trades)

    @property
    def pnl_list(self):
        return [t.pnl for t in self.trades]

    @property
    def pv(self):
        return sum(t.pv for t in self.trades)

    @property
    def pv_list(self):
        return [t.pv for t in self.trades]

    def reset_pv(self):
        for t in self.trades:
            t.reset_pv()

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

    @trade_date.setter
    def trade_date(self, d):
        #we try to keep everybody in sync
        for index in self.indices:
            index.trade_date = d
        if len(self.indices) == 0:
            for swaption in self.swaptions:
                self.swaption.trade_date = d
        self._trade_date = d

    def mark(self, source_list=[], option_type=None, model="black", surface_id=None):
        #add None so that we always try everything
        source_list = source_list + [None]
        for index, (index_type, series, tenor) in zip(self.indices, self._keys):
            index.mark()
            if tenor != '5yr':
                continue
            k = (index.trade_date, index_type, series, tenor)
            if self.swaptions:
                if k not in self._vs:
                    vs = VolatilitySurface(index_type, series, tenor, index.trade_date)
                    if surface_id is None:
                        for source in source_list:
                            if len(vs.list(source, option_type, model)) >=1:
                                break
                        else:
                            raise ValueError("No market data available for this day")
                        self._vs[k] = vs[vs.list(source, option_type, model)[-1]]
                    else:
                        self._vs[k] = vs[surface_id]
        for swaption in self.swaptions:
            vol_surface = self._vs[(swaption.index.trade_date, ) + \
                                   _key_from_index(swaption.index)]
            swaption.sigma = float(self._vs[(swaption.index.trade_date, ) \
                                            + _key_from_index(swaption.index)].
                                   ev(swaption.T, swaption.moneyness))

    @property
    def ref(self):
        if len(self.indices) == 1:
            return self.indices[0].ref
        else:
            return [index.ref for index in self.indices]

    @ref.setter
    def ref(self, val):
        if len(self.indices) == 1:
            self.indices[0].ref = val
        elif len(self.indices) == 0:
                ##no index, so set the individual refs
            for t in trades:
                t.index.ref = val
        elif len(self.indices) == len(val):
            for index, val in zip(self.indices, val):
                index.ref = val
        else:
            raise ValueError("The number of refs doesn't match the number of indices")

    @property
    def spread(self):
        if len(self.indices) == 1:
            return self.indices[0].spread
        else:
            return [index.spread for index in self.indices]

    @spread.setter
    def spread(self, val):
        if len(self.indices) == 1:
            self.indices[0].spread = val
        elif len(self.indices) == 0:
                ##no index, so set the individual refs
            for t in trades:
                t.index.spread = val
        elif len(self.indices) == len(val):
            for index, val in zip(self.indices, val):
                index.spread = val
        else:
            raise ValueError("The number of spreads doesn't match the number of indices")

    @property
    def delta(self):
        """returns the equivalent protection notional

        makes sense only where there is a single index."""
        return sum([getattr(t, 'delta', -t._direction) * t.notional for t in self.trades])

    @property
    def gamma(self):
        return sum([getattr(t, 'gamma', 0) * t.notional for t in self.trades])

    @property
    def dv01(self):
        return sum(t.dv01 for t in self.trades)

    @property
    def theta(self):
        return sum(t.theta for t in self.trades)

    def _todf(self):
        headers = ["Product", "Index", "Notional", "Ref", "Direction", "Expiry",
                   "Vol", "PV", "Delta", "Gamma", "Theta", "Vega"]
        rec = []
        for t in self.trades:
            if isinstance(t, Index):
                index_type, series, tenor = _key_from_index(t)
                r = ("Index", f"{index_type}{series} {tenor}",
                     t.notional, t.ref, t.direction, "N/A", None, t.pv, 1., 0., t.theta, 0.)
            elif isinstance(t, BlackSwaption):
                index_type, series, tenor = _key_from_index(t.index)
                r = ("Swaption", f"{index_type}{series} {tenor}",
                     t.notional, t.ref, t.direction, t.forward_date, t.sigma, t.pv,
                     t.delta, t.gamma, t.theta, t.vega)
            else:
                raise TypeError
            rec.append(r)
        return pd.DataFrame.from_records(rec, columns=headers)

    __repr__ = portf_repr('string')

    _repr_html_ = portf_repr('html')