import pandas as pd

from tia.util.decorator import lazy_property
from tia.analysis.model.interface import TxnColumns as TC
from tia.analysis.model.pl import TxnProfitAndLoss
from tia.analysis.model.ret import RoiiRetCalculator
from tia.analysis.util import is_decrease, is_increase, crosses_zero


__all__ = ['Intent', 'Action', 'iter_txns', 'Txns']


class Intent(object):
    Open = 1
    Close = 2
    Increase = 3
    Decrease = 4

    Labels = {
        Open: 'Open',
        Close: 'Close',
        Increase: 'Increase',
        Decrease: 'Decrease',
    }


class Action(object):
    Buy = 1
    Sell = 2
    SellShort = 3
    Cover = 4

    Labels = {
        Buy: 'Buy',
        Sell: 'Sell',
        SellShort: 'SellShort',
        Cover: 'Cover',
    }


def iter_txns(trds):
    """ iterator of trades which splits trades to ensure proper long/short accounting"""
    pos = 0
    for trd in trds:
        if pos != 0 and is_decrease(pos, trd.qty) and crosses_zero(pos, trd.qty):
            # Split to make accounting for long/short possible
            closing_trd, opening_trd = trd.split(-pos)
            # setattr(closing_trd, '_txn_id', 1)
            # setattr(opening_trd, '_txn_id', 2)
            yield closing_trd
            pos = opening_trd.qty
            yield opening_trd
        else:
            pos += trd.qty
            yield trd


class Txns(object):
    def __init__(self, trades, pricer, ret_calc=None):
        """
        #TODO - rethink if user should split trades prior to calling this method...
        :param trades: list of Trade objects
        :param pricer:
        """
        # split into l/s positions
        self.trades = tuple(iter_txns(trades))
        self.pricer = pricer
        self._ret_calc = ret_calc or RoiiRetCalculator()

    pids = property(lambda self: self.frame[TC.PID].unique())
    pl = lazy_property(lambda self: TxnProfitAndLoss(self), 'pl')
    performance = lazy_property(lambda self: self.ret_calc.compute(self), 'performance')

    @property
    def ret_calc(self):
        return self._ret_calc

    @ret_calc.setter
    def ret_calc(self, calc):
        self._ret_calc = calc
        if hasattr(self, '_performance'):
            delattr(self, '_performance')

    @lazy_property
    def frame(self):
        """Convert the trades to transaction level details necessary for long/short accouting.

        :param trades:
        :param pricer: provides the interface to get premium for a specified quanity, price, and timestamp.
        :return:
        """
        rows = []
        pricer = self.pricer
        pos = open_val = pid = 0
        for txn in self.trades:
            # These values always get copied
            qty = txn.qty
            premium = pricer.get_premium(qty, txn.px, ts=txn.ts)
            if pos == 0:  # Open position
                side = qty > 0 and Action.Buy or Action.SellShort
                open_val = premium
                pid += 1
                side = side
                intent = Intent.Open
                pos = qty
            elif pos + qty == 0:  # close position
                side = qty > 0 and Action.Cover or Action.Sell
                open_val = 0
                side = side
                intent = Intent.Close
                pos = 0
            elif is_increase(pos, qty):
                side = txn.qty > 0 and Action.Buy or Action.SellShort
                open_val += premium
                pos += qty
                intent = Intent.Increase
                side = side
            else:  # decrease - no worries about split since iterator takes care of it
                side = txn.qty > 0 and Action.Cover or Action.Sell
                open_val *= ((pos + qty) / pos)
                pos += qty
                intent = Intent.Decrease
                side = side

            # Get rid of anything but the date
            dt = txn.ts.to_period('B').to_timestamp()
            rows.append([dt, txn.ts, pid, txn.tid, txn.qty, txn.px, txn.fees, premium, open_val, pos, intent, side])

        df = pd.DataFrame.from_records(rows, columns=[TC.DT, TC.TS, TC.PID, TC.TID, TC.QTY, TC.PX, TC.FEES, TC.PREMIUM,
                                                      TC.OPEN_VAL, TC.POS, TC.INTENT, TC.ACTION])
        df.index.name = 'seq'
        return df

    def get_pid_txns(self, pid):
        pmask = self.frame[TC.PID] == pid
        assert len(pmask.index) == len(self.trades), 'assume 1-1 ratio of trade to row in frame'
        return tuple(pd.Series(self.trades)[pmask.values])

    def subset(self, pids):
        pmask = self.frame[TC.PID].isin(pids)
        if pmask.all():
            return self
        else:
            # 1 to 1 mapping of txn to row (so can figure out trades from mask)
            trds = tuple(pd.Series(self.trades)[pmask.values])
            # build the object
            result = Txns(trds, self.pricer, self.ret_calc)
            result._frame = self.frame.loc[pmask]
            if hasattr(self, '_profit_and_loss'):
                pl = self.profit_and_loss
                result._profit_and_loss = pl.subset(result)
            return result