import pandas as pd
import datetime as dt
import marketsim.marketsim as marketsim
import indicators
import util
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from BenchmarkStrategy import BenchmarkStrategy
from ManualStrategy import ManualStrategy


def macd_strat(macd):
    def strat(x):
        print(x)

    macd['macd_trigger'] = macd.rolling(2).apply(strat)

    # for i, row in macd.iterrows():
        # if i == 0:
            # continue
        # print(row)
        # prev_macd, prev_signal, _ = row
        # cur_macd, cur_signal, _ = row
        # if cur_macd < -.5 and (prev_macd < prev_signal) \
                          # and (cur_macd > cur_signal):
            # macd.iloc[i]['macd_buy_sell'] = 1

def experiment1():
    symbol = "JPM"
    start_value = 10000
    sd = dt.datetime(2008, 1, 1)
    ed = dt.datetime(2009, 12, 31)
    df = util.get_data([symbol], pd.date_range(sd, ed))
    df.drop(columns=["SPY"], inplace=True)
    # df = pd.DataFrame(index=df.index)

    bs = BenchmarkStrategy()
    orders = bs.testPolicy(symbol, sd, ed, start_value)
    df["Benchmark"] = marketsim.compute_portvals(orders, start_value)

    ms = ManualStrategy()
    orders = ms.testPolicy(symbol, sd, ed, start_value)
    df["Manual"] = marketsim.compute_portvals(orders, start_value)

    # indicators.price_sma(df, symbol, 21)
    # sma = indicators.sma(df, symbol, [9, 21])
    # rsi = indicators.rsi(df, symbol)
    macd = indicators.macd(df, symbol).copy()
    # macd_strat(macd)

    fig, ax = plt.subplots(2, sharex=True)
    df[symbol].plot(ax=ax[0])
    # sma.plot(ax=ax[0])
    macd.plot(ax=ax[1])
    # macd.iloc[:,0].plot(ax=ax[1])
    # rsi.plot(ax=ax[2])
    # df[["Benchmark", "Manual"]].plot(ax=ax[3])

    # XXX: Plot where we buy and sell.

    for a in ax: a.grid()
    multi = MultiCursor(fig.canvas, ax, color='r', lw=0.5)
    plt.show()

    # df.plot(title="results", subplots=True)
    #sd = dt.datetime(2008, 1, 1)
    #ed = dt.datetime(2009, 12, 31)
    #df = get_data([symbol], pd.date_range(sd, ed))
    #df.drop(columns=["SPY"], inplace=True)
    # df_orig = df.copy()
    #df = indicators.normalize(df)
    #indicators.price_sma(df, symbol, 21)
    #df.plot(title="21 SMA and EMA")
    #plt.show()

    # You may use data from other symbols (such as SPY) to inform both your
    # Manual Learner and Strategy Learner. The in-sample/development period is
    # January 1, 2008 to December 31 2009. The out-of-sample/testing period is
    # January 1, 2010 to December 31 2011.


class BlittedCursor:
    """
    A cross hair cursor using blitting for faster redraw.
    """
    def __init__(self, ax):
        self.ax = ax
        self.background = None
        self.horizontal_line = ax.axhline(color='k', lw=0.8, ls='--')
        self.vertical_line = ax.axvline(color='k', lw=0.8, ls='--')
        # text location in axes coordinates
        self.text = ax.text(0.72, 0.9, '', transform=ax.transAxes)
        self._creating_background = False
        ax.figure.canvas.mpl_connect('draw_event', self.on_draw)

    def on_draw(self, event):
        self.create_new_background()

    def set_cross_hair_visible(self, visible):
        need_redraw = self.horizontal_line.get_visible() != visible
        self.horizontal_line.set_visible(visible)
        self.vertical_line.set_visible(visible)
        self.text.set_visible(visible)
        return need_redraw

    def create_new_background(self):
        if self._creating_background:
            # discard calls triggered from within this function
            return
        self._creating_background = True
        self.set_cross_hair_visible(False)
        self.ax.figure.canvas.draw()
        self.background = self.ax.figure.canvas.copy_from_bbox(self.ax.bbox)
        self.set_cross_hair_visible(True)
        self._creating_background = False

    def on_mouse_move(self, event):
        if self.background is None:
            self.create_new_background()
        if not event.inaxes:
            need_redraw = self.set_cross_hair_visible(False)
            if need_redraw:
                self.ax.figure.canvas.restore_region(self.background)
                self.ax.figure.canvas.blit(self.ax.bbox)
        else:
            self.set_cross_hair_visible(True)
            # update the line positions
            x, y = event.xdata, event.ydata
            self.horizontal_line.set_ydata(y)
            self.vertical_line.set_xdata(x)
            self.text.set_text('x=%1.2f, y=%1.2f' % (x, y))

            self.ax.figure.canvas.restore_region(self.background)
            self.ax.draw_artist(self.horizontal_line)
            self.ax.draw_artist(self.vertical_line)
            self.ax.draw_artist(self.text)
            self.ax.figure.canvas.blit(self.ax.bbox)


if __name__ == "__main__":
    experiment1()