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 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)

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

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

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

    fig, ax = plt.subplots(4, sharex=True)
    df[symbol].plot(ax=ax[0])
    macd.plot(ax=ax[3])
    df[["Benchmark", "Manual"]].plot(ax=ax[1])
    df[["Orders Benchmark", "Orders Manual"]].plot(ax=ax[2])

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

    # 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.


if __name__ == "__main__":
    experiment1()