from .index import Index
from .option import BlackSwaption, VolatilitySurface
from db import dbengine
from warnings import warn

serenitasdb = dbengine('serenitasdb')

def _key_from_index(index):
    _, index_type, _, series, tenor = index.name.split()
    series = int(series[1:])
    tenor = tenor.lower() + 'r'
    return index_type, series, tenor

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)]
        self._keys = []
        for index in self.indices:
            self._keys.append(_key_from_index(index))
        #pick the first available trade_date
        trade_dates = set()
        for index in self.indices:
            trade_dates.add(index.trade_date)
        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 set_original_pv(self):
        for t in self.trades:
            t.set_original_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=None, option_type=None, model=None, surface_id=None):
        vs_dict = {}
        for index, (index_type, series, tenor) in zip(self.indices, self._keys):
            run = serenitasdb.execute("""SELECT * FROM index_quotes
            WHERE index=%s AND series=%s AND tenor=%s AND date=%s""",
                                      (index_type, series, tenor, self.trade_date))
            rec = run.fetchone()
            index.spread = rec.closespread
            k = (self.trade_date, index_type, series, tenor)
            if k not in self._vs:
                vs = VolatilitySurface(index_type, series, tenor, self.trade_date)
                if surface_id is None and len(vs.list(source, option_type, model)) >=1:
                    surface_id = vs.list(source, option_type, model)[-1]
                else:
                    raise ValueError("No market data available for this day")
                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 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)