import pandas as pd
import datetime as dt
import sys

import util
import indicators
import marketsim.marketsim as marketsim
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from BenchmarkStrategy import BenchmarkStrategy
from ManualStrategy import ManualStrategy
from StrategyLearner import StrategyLearner


def plot_indicators(symbol, df):
    fig, ax = plt.subplots(4, sharex=True)

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

    df[[symbol]].plot(ax=ax[0])
    bb.plot(ax=ax[0])
    price_sma.plot(ax=ax[1])
    macd.plot(ax=ax[2])
    rsi.plot(ax=ax[3])
    for a in ax.flat:
        a.grid()
    plt.show()
    sys.exit(0)


def visualize_correlations(symbol, df):
    indicators.price_sma(df, symbol, [8, 21])
    indicators.price_delta(df, symbol, 5)
    indicators.price_delta(df, symbol, 3)
    indicators.price_delta(df, symbol, 1)
    indicators.macd(df, symbol)
    indicators.rsi(df, symbol)

    # df = df[df['rsi'] > 80]
    fig, ax = plt.subplots(3, 2)  # sharex=True)
    df.plot.scatter(x="price_sma_8", y="pct_5", ax=ax[0, 0])
    df.plot.scatter(x="price_sma_8", y="pct_3", ax=ax[1, 0])
    df.plot.scatter(x="price_sma_8", y="pct_1", ax=ax[2, 0])
    # df.plot.scatter(x="rsi", y="pct_5", ax=ax[0, 1])
    # df.plot.scatter(x="rsi", y="pct_3", ax=ax[1, 1])
    # df.plot.scatter(x="rsi", y="pct_1", ax=ax[2, 1])
    df.plot.scatter(x="macd_diff", y="pct_5", ax=ax[0, 1])
    df.plot.scatter(x="macd_diff", y="pct_3", ax=ax[1, 1])
    df.plot.scatter(x="macd_diff", y="pct_1", ax=ax[2, 1])

    for a in ax.flat:
        a.grid()
    plt.show()
    sys.exit(0)


def compare_manual_strategies(symbol, sv, sd, ed):

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

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

    ms = ManualStrategy()
    orders = ms.testPolicy(symbol, sd, ed, sv, macd_strat=True)
    df["MACD Strat"] = marketsim.compute_portvals(orders, sv)
    df["Orders MACD"] = orders["Shares"]
    # df["Holding Manual"] = orders["Shares"].cumsum()

    orders = ms.testPolicy(symbol, sd, ed, sv)
    df["Three Strat"] = marketsim.compute_portvals(orders, sv)
    df["Orders Three"] = orders["Shares"]

    fig, ax = plt.subplots(3, sharex=True)
    df[[symbol]].plot(ax=ax[0])
    df[["Benchmark", "MACD Strat", "Three Strat"]].plot(ax=ax[1])
    df[["Orders Benchmark", "Orders MACD", "Orders Three"]].plot(ax=ax[2])

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

    # plt.show()
    fig.set_size_inches(10, 8, forward=True)
    plt.savefig('figure_1.png', dpi=fig.dpi)


def experiment1():
    symbol = "JPM"
    sv = 10000
    sd = dt.datetime(2008, 1, 1)  # in-sample
    ed = dt.datetime(2009, 12, 31)  # in-sample
    sd_out = dt.datetime(2010, 1, 1)  # out-sample
    ed_out = dt.datetime(2011, 12, 31)  # out-sample

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

    # visualize_correlations(symbol, df)
    # plot_indicators(symbol, df)
    # compare_manual_strategies(symbol, sv, sd, ed)

    bs = BenchmarkStrategy()
    orders = bs.testPolicy(symbol, sd_out, ed_out, sv)
    df["Benchmark"] = marketsim.compute_portvals(orders, sv)
    df["Orders Benchmark"] = orders["Shares"]

    sl = StrategyLearner(testing=True)
    sl.addEvidence(symbol, sd, ed, sv)
    orders = sl.testPolicy(symbol, sd_out, ed_out, sv)
    df["SL"] = marketsim.compute_portvals(orders, sv)
    df["Orders SL"] = orders["Shares"]

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

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

    # For debugging the classification learner:
    # df["y_train"] = sl.addEvidence(symbol, sd, ed, sv)
    # df["y_query"] = sl.testPolicy(symbol, sd, ed, sv)
    # df[["y_train", "y_query"]].plot(ax=ax[1])



if __name__ == "__main__":
    experiment1()