from .index import CreditIndex
from .option import BlackSwaption
from warnings import warn
import pandas as pd
import numpy as np


def portf_repr(method):
    def f(*args):
        obj = args[0]
        thousands = "{:,.2f}".format

        def percent(x):
            if np.isnan(x):
                return "N/A"
            else:
                return f"{100*x:.2f}%"
        header = f"Portfolio {obj.value_date}\n\n"
        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, trade_ids=None):
        self.trades = trades
        self.trade_ids = trade_ids
        self.indices = [t for t in trades if isinstance(t, CreditIndex)]
        self.swaptions = [t for t in trades if isinstance(t, BlackSwaption)]
        value_dates = set(t.value_date for t in self.trades)
        self._keys = set([(index.index_type, index.series, index.tenor) for index in self.indices])
        for swaption in self.swaptions:
            self._keys.add((swaption.index.index_type, swaption.index.series, swaption.index.tenor))
        self._value_date = value_dates.pop()
        if len(value_dates) >= 1:
            warn(f"not all instruments have the same trade date, picking {self._value_date}")

    def __iter__(self):
        for t in self.trades:
            yield t

    def items(self):
        for trade_id, trade in zip(self.trade_ids, self.trades):
            yield (trade_id, trade)

    @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 value_date(self):
        return self._value_date

    @value_date.setter
    def value_date(self, d):
        for t in self.trades:
            t.value_date = d
        self._value_date = d

    def mark(self, **kwargs):
        for t in self.trades:
            t.mark(**kwargs)

    def shock(self, params=["pnl"], **kwargs):
        return {trade_id: trade.shock(params, **kwargs) for trade_id, trade in self.items()}

    @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 self.swaptions:
                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 self.swaptions:
                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", "Strike", "Direction",
                   "Type", "Expiry", "Vol", "PV", "Delta", "Gamma", "Theta",
                   "Vega"]
        rec = []
        for t in self.trades:
            if isinstance(t, CreditIndex):
                name = f"{t.index_type}{t.series} {t.tenor}"
                r = ("Index", name,
                     t.notional, t.ref, "N/A", t.direction, "N/A", "N/A", None, t.pv, 1., 0., t.theta, 0.)
            elif isinstance(t, BlackSwaption):
                name = f"{t.index.index_type}{t.index.series} {t.index.tenor}"
                r = ("Swaption", name,
                     t.notional, t.ref, t.strike, t.direction, t.option_type, 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, index=self.trade_ids)

    __repr__ = portf_repr('string')

    _repr_html_ = portf_repr('html')