import datetime as dt
import pandas as pd
import util
import indicators
from RTLearner import RTLearner


class StrategyLearner(object):

    def __init__(self, verbose=False, impact=0.0, commission=0.0, testing=False):
        self.verbose = verbose
        self.impact = impact
        self.commission = commission
        self.testing = testing

    def _get_volume(self):
        """For reference."""
        volume_all = ut.get_data(syms, dates, colname="Volume")
        volume = volume_all[syms]  # only portfolio symbols
        # volume_SPY = volume_all['SPY']  # only SPY, for comparison later
        if self.verbose:
            print(volume)

    def _add_indicators(self, df, symbol):
        """Add indicators for learning to DataFrame."""
        df.drop(columns=["SPY"], inplace=True)
        indicators.macd(df, symbol)
        indicators.rsi(df, symbol)
        indicators.price_sma(df, symbol, [8])
        indicators.price_delta(df, symbol, 3)
        df.dropna(inplace=True)

    def addEvidence(self, symbol="IBM",
                    sd=dt.datetime(2008, 1, 1),
                    ed=dt.datetime(2009, 1, 1),
                    sv=10000):

        self.y_threshold = 0.2
        self.indicators = ['macd_diff', 'rsi', 'price_sma_8']
        df = util.get_data([symbol], pd.date_range(sd, ed))
        self._add_indicators(df, symbol)

        def classify_y(row):
            if row > self.y_threshold:
                return 1
            elif row < -self.y_threshold:
                return -1
            else:
                pass
            return 0

        def set_y_threshold(pct):
            if max(pct) < 0.2:
                self.y_threshold = 0.02

        self.learner = RTLearner(leaf_size = 5)
        # self.learner = BagLearner(RTLearner, 3, {'leaf_size': 5})
        data_x = df[self.indicators].to_numpy()
        pct = df['pct_3']

        # This is a hack to get a low enough buy/sell threshold for the
        # cyclic the test 'ML4T-220' where the max pct_3 is 0.0268.
        set_y_threshold(pct)
        y = pct.apply(classify_y)

        self.learner.addEvidence(data_x, y.to_numpy())
        return y

    def strat(self, data_y, orders):
        self.holding = 0

        def strat(row):
            y = int(data_y.loc[row.name][0])
            shares = 0
            if self.holding == 0 and y == 1:
                shares = 1000
            elif self.holding == -1000 and y == 1:
                shares = 2000
            elif self.holding == 0 and y == -1:
                shares = -1000
            elif self.holding == 1000 and y == -1:
                shares = -2000
            self.holding += shares
            return shares

        orders["Shares"] = orders.apply(strat, axis=1)

    def testPolicy(self, symbol="IBM",
                   sd=dt.datetime(2009, 1, 1),
                   ed=dt.datetime(2010, 1, 1),
                   sv=10000):
        df = util.get_data([symbol], pd.date_range(sd, ed))
        self._add_indicators(df, symbol)
        data_x = df[self.indicators].to_numpy()
        data_y = pd.DataFrame(index=df.index, data=self.learner.query(data_x))

        orders = pd.DataFrame(index=df.index)
        orders["Symbol"] = symbol
        orders["Order"] = ""
        orders["Shares"] = 0
        self.strat(data_y, orders)
        if self.testing:
            return orders
        else:
            return orders[["Shares"]]