"""
Copyright 2016 Stephen Boyd, Enzo Busseti, Steven Diamond, BlackRock Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from abc import ABCMeta, abstractmethod
import pandas as pd
import numpy as np
import logging
import cvxpy as cvx

from .costs import BaseCost
from .returns import BaseReturnsModel
from .constraints import BaseConstraint


__all__ = ['Hold', 'FixedTrade', 'PeriodicRebalance', 'AdaptiveRebalance',
           'SinglePeriodOpt', 'MultiPeriodOpt', 'ProportionalTrade']


class BasePolicy(object):
    """ Base class for a trading policy. """
    __metaclass__ = ABCMeta

    def __init__(self):
        self.costs = []
        self.constraints = []

    @abstractmethod
    def get_trades(self, portfolio, t):
        """Trades list given current portfolio and time t.
        """
        return NotImplemented

    def _nulltrade(self, portfolio):
        return pd.Series(index=portfolio.index, data=0.)


class Hold(BasePolicy):
    """Hold initial portfolio.
    """
    def get_trades(self, portfolio, t):
        return self._nulltrade(portfolio)


class ProportionalTrade(BasePolicy):
    """Gets to target in given time steps."""
    def __init__(self, targetweight, time_steps):
        self.targetweight = targetweight
        self.time_steps = time_steps
        super(ProportionalTrade, self).__init__()

    def get_trades(self, portfolio, t):
        try:
            missing_time_steps = len(
                self.time_steps)-next(i for (i, x)
                                      in enumerate(self.time_steps)
                                      if x == t)
        except StopIteration:
            raise Exception(
                "ProportionalTrade can only trade on the given time steps")
        deviation = self.targetweight-portfolio/sum(portfolio)
        return sum(portfolio)*deviation/missing_time_steps


class SellAll(BasePolicy):
    """Sell all non-cash assets."""
    def get_trades(self, portfolio, t):
        trade = -pd.Series(portfolio, copy=True)
        trade.ix[-1] = 0.
        return trade


class FixedTrade(BasePolicy):
    """Trade a fixed trade vector.
    """
    def __init__(self, tradevec=None, tradeweight=None):
        """Trade the tradevec vector (dollars) or tradeweight weights."""
        if tradevec is not None and tradeweight is not None:
            raise(Exception(
                'only one of tradevec and tradeweight can be passed'))
        if tradevec is None and tradeweight is None:
            raise(Exception('one of tradevec and tradeweight must be passed'))
        self.tradevec = tradevec
        self.tradeweight = tradeweight
        assert(self.tradevec is None or sum(self.tradevec) == 0.)
        assert(self.tradeweight is None or sum(self.tradeweight) == 0.)
        super(FixedTrade, self).__init__()

    def get_trades(self, portfolio, t):
        if self.tradevec is not None:
            return self.tradevec
        return sum(portfolio)*self.tradeweight


class BaseRebalance(BasePolicy):

    def _rebalance(self, portfolio):
        return sum(portfolio) * self.target - portfolio


class PeriodicRebalance(BaseRebalance):
    """Track a target portfolio, rebalancing at given times.
    """
    def __init__(self, target, period, **kwargs):
        """
        Args:
            target: target weights, n+1 vector
            period: supported options are "day", "week", "month", "quarter",
                "year".
                rebalance on the first day of each new period
        """
        self.target = target
        self.period = period
        super(PeriodicRebalance, self).__init__()

    def is_start_period(self, t):
        result = not getattr(t, self.period) == getattr(self.last_t,
                                                        self.period) if \
                                                        hasattr(self,
                                                                'last_t')\
                                                        else True
        self.last_t = t
        return result

    def get_trades(self, portfolio, t):
        return self._rebalance(portfolio) if self.is_start_period(t) else \
            self._nulltrade(portfolio)


class AdaptiveRebalance(BaseRebalance):
    """ Rebalance portfolio when deviates too far from target.
    """
    def __init__(self, target, tracking_error):
        self.target = target
        self.tracking_error = tracking_error
        super(AdaptiveRebalance, self).__init__()

    def get_trades(self, portfolio, t):
        weights = portfolio/sum(portfolio)
        diff = (weights - self.target).values

        if np.linalg.norm(diff, 2) > self.tracking_error:
            return self._rebalance(portfolio)
        else:
            return self._nulltrade(portfolio)


class SinglePeriodOpt(BasePolicy):

    def __init__(self, alpha_model, costs, constraints, solver=None,
                 solver_opts={}):

        self.alpha_model = alpha_model
        solver_opts = solver_opts
        assert isinstance(self.alpha_model, BaseReturnsModel)

        super(SinglePeriodOpt, self).__init__()

        for cost in costs:
            assert isinstance(cost, BaseCost)
            self.costs.append(cost)

        for constraint in constraints:
            assert isinstance(constraint, BaseConstraint)
            self.constraints.append(constraint)

        self.solver = solver
        self.solver_opts = solver_opts

    def get_trades(self, portfolio, t):

        value = sum(portfolio)
        w = portfolio/value
        z = cvx.Variable(w.size)  # TODO pass index
        wplus = w.values + z

        alpha_term = self.alpha_model.weight_expr(t, wplus)
        assert(alpha_term.is_concave())

        costs, constraints = [], []

        for cost in self.costs:
            cost_expr, const_expr = cost.weight_expr(t, wplus, z, value)
            costs.append(cost_expr)
            constraints += const_expr

        constraints += [item for item in (con.weight_expr(t, wplus, z, value)
                                          for con in self.constraints)]

        for el in costs:
            assert (el.is_convex())

        for el in constraints:
            assert (el.is_dcp())

        prob = cvx.Problem(
            cvx.Maximize(alpha_term - sum(costs)),
            [cvx.sum_entries(z) == 0] + constraints)
        try:
            prob.solve(solver=self.solver, **self.solver_opts)

            if prob.status == 'unbounded':
                logging.error(
                    'The problem is unbounded. Defaulting to no trades')
                return self._nulltrade(portfolio)

            if prob.status == 'infeasible':
                logging.error(
                    'The problem is infeasible. Defaulting to no trades')
                return self._nulltrade(portfolio)

            # TODO will have index
            return pd.Series(index=portfolio.index, data=(z.value.A1 * value))
        except cvx.SolverError:
            logging.error(
                'The solver %s failed. Defaulting to no trades' % self.solver)
            return self._nulltrade(portfolio)

# class LookaheadModel():
#     """Returns the planning periods for multi-period.
#     """
#     def __init__(self, trading_times, period_lens):
#         self.trading_times = trading_times
#         self.period_lens = period_lens
#
#     def get_periods(self, t):
#         """Returns planning periods.
#         """
#         periods = []
#         tau = t
#         for length in self.period_lens:
#             incr = length*pd.Timedelta('1 days')
#             periods.append((tau, tau + incr))
#             tau += incr
#         return periods


class MultiPeriodOpt(SinglePeriodOpt):

    def __init__(self, trading_times,
                 terminal_weights, lookahead_periods=None, *args, **kwargs):
        """
        trading_times: list, all times at which get_trades will be called
        lookahead_periods: int or None. if None uses all remaining periods
        """
        # Number of periods to look ahead.
        self.lookahead_periods = lookahead_periods
        self.trading_times = trading_times
        # Should there be a constraint that the final portfolio is the bmark?
        self.terminal_weights = terminal_weights
        super(MultiPeriodOpt, self).__init__(*args, **kwargs)

    def get_trades(self, portfolio, t):

        value = sum(portfolio)
        assert (value > 0.)
        w = cvx.Constant(portfolio.values/value)

        prob_arr = []
        z_vars = []

        # planning_periods = self.lookahead_model.get_periods(t)
        for tau in \
                self.trading_times[self.trading_times.index(t):
                                   self.trading_times.index(t) +
                                   self.lookahead_periods]:
            # delta_t in [pd.Timedelta('%d days' % i) for i in
            # range(self.lookahead_periods)]:

            #            tau = t + delta_t
            z = cvx.Variable(*w.size)
            wplus = w + z
            obj = self.alpha_model.weight_expr_ahead(t, tau, wplus)

            costs, constr = [], []
            for cost in self.costs:
                cost_expr, const_expr = cost.weight_expr_ahead(
                    t, tau, wplus, z, value)
                costs.append(cost_expr)
                constr += const_expr

            obj -= sum(costs)
            constr += [cvx.sum_entries(z) == 0]
            constr += [con.weight_expr(t, wplus, z, value)
                       for con in self.constraints]

            prob = cvx.Problem(cvx.Maximize(obj), constr)
            prob_arr.append(prob)
            z_vars.append(z)
            w = wplus

        # Terminal constraint.
        if self.terminal_weights is not None:
            prob_arr[-1].constraints += [wplus == self.terminal_weights.values]

        sum(prob_arr).solve(solver=self.solver)
        return pd.Series(index=portfolio.index,
                         data=(z_vars[0].value.A1 * value))
