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


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

    # this method should create a QLearner, and train it for trading
    def addEvidence(self, symbol="IBM",
                    sd=dt.datetime(2008, 1, 1),
                    ed=dt.datetime(2009, 1, 1),
                    sv=10000):

        # add your code to do learning here

        # example usage of the old backward compatible util function
        syms = [symbol]
        dates = pd.date_range(sd, ed)
        prices_all = util.get_data(syms, dates)  # automatically adds SPY
        prices = prices_all[syms]  # only portfolio symbols
        # prices_SPY = prices_all['SPY']  # only SPY, for comparison later
        if self.verbose:
            print(prices)

        # example use with new colname
        # automatically adds SPY
        volume_all = util.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 macd_strat(self, macd, orders):
        """Strategy based on MACD cross."""

        def strat(ser):
            m = macd.loc[ser.index]
            prev_macd, prev_signal, _ = m.iloc[0]
            cur_macd, cur_signal, _ = m.iloc[1]
            shares = 0
            if cur_macd < -1 and prev_macd < prev_signal \
                             and cur_macd > cur_signal:
                if self.holding == 0:
                    shares = 1000
                elif self.holding == -1000:
                    shares = 2000
            elif cur_macd > 1 and prev_macd > prev_signal \
                              and cur_macd < cur_signal:
                if self.holding == 0:
                    shares = -1000
                elif self.holding == 1000:
                    shares = -2000
            self.holding += shares
            return shares

        orders['Shares'] = orders['Shares'].rolling(2).apply(strat)

    def three_indicator_strat(self, macd, rsi, price_sma, orders):
        """Strategy based on three indicators. Thresholds selected based on
        scatter plots."""
        def strat(row):
            shares = 0
            _, _, macd_diff = macd.loc[row.name]
            cur_rsi = rsi.loc[row.name][0]
            cur_price_sma = price_sma.loc[row.name][0]
            if self.holding == -1000 and cur_price_sma < 0.9:
                shares = 2000
            elif self.holding == 0 and cur_price_sma < 0.9:
                shares = 1000
            elif self.holding == -1000 and cur_rsi > 80:
                shares = 2000
            elif self.holding == 0 and cur_rsi > 80:
                shares = 1000
            elif self.holding == -1000 and macd_diff < -0.5:
                shares = 2000
            elif self.holding == 0 and macd_diff < -0.5:
                shares = 1000
            elif self.holding == 1000 and cur_price_sma > 1.1:
                shares = -2000
            elif self.holding == 0 and cur_price_sma > 1.1:
                shares = -1000
            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, macd_strat=False):

        self.holding = 0
        df = util.get_data([symbol], pd.date_range(sd, ed))
        df.drop(columns=["SPY"], inplace=True)

        orders = pd.DataFrame(index=df.index)
        orders["Symbol"] = symbol
        orders["Order"] = ""
        orders["Shares"] = 0

        macd = indicators.macd(df, symbol)
        rsi = indicators.rsi(df, symbol)
        price_sma = indicators.price_sma(df, symbol, [8])

        if macd_strat:
            self.macd_strat(macd, orders)
        else:
            self.three_indicator_strat(macd, rsi, price_sma, orders)
        return orders