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
from QLearner import QLearner


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()
    m = MultiCursor(fig.canvas, ax, color='r', lw=0.5)
    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 compare_all_strategies(symbol, sv, sd, ed):
    df = util.get_data([symbol], pd.date_range(sd, ed))
    df.drop(columns=["SPY"], inplace=True)
    normalize = indicators.normalize

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

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

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

    fig, ax = plt.subplots(3, sharex=True)
    df[[symbol]].plot(ax=ax[0])
    df[["Benchmark", "Manual", "Strategy"]].plot(ax=ax[1])
    df[["Orders Benchmark", "Orders Manual", "Orders Strategy"]].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_2.png', dpi=fig.dpi)


def compare_number_trades():
    symbol = "JPM"
    sv = 10000
    sd = dt.datetime(2008, 1, 1)  # in-sample
    ed = dt.datetime(2009, 12, 31)  # in-sample

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

    print(f"| commission | n_orders |")
    print(f"-------------------------")
    for commission in [0, 9.95, 20, 50, 100]:
        ql = QLearner(testing=True, commission=commission, impact=0.005)
        ql.addEvidence(symbol, sd, ed, sv)
        orders = ql.testPolicy(symbol, sd, ed, sv)
        n_orders = orders[orders["Shares"] != 0].shape[0]
        print(f"| {commission} | {n_orders} |")

def compare_q_learners():
    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_out, ed_out))
    df.drop(columns=["SPY"], inplace=True)

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

    ql = QLearner(testing=True, verbose=False)
    ql.addEvidence(symbol, sd, ed, sv)
    orders = ql.testPolicy(symbol, sd_out, ed_out, sv)
    df["QL 5"] = indicators.normalize(marketsim.compute_portvals(orders, sv))
    df["Orders QL 5"] = orders["Shares"]

    ql = QLearner(testing=True, verbose=False, n_bins=4)
    ql.addEvidence(symbol, sd, ed, sv)
    orders = ql.testPolicy(symbol, sd_out, ed_out, sv)
    df["QL 4"] = indicators.normalize(marketsim.compute_portvals(orders, sv))
    df["Orders QL 4"] = orders["Shares"]

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

    for a in ax:
        a.grid()
    m = MultiCursor(fig.canvas, ax, color='r', lw=0.5)
    fig.set_size_inches(10, 8, forward=True)
    plt.savefig('figure_4.png', dpi=fig.dpi)
    sys.exit(0)


def experiment1(create_report=False):
    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_out, ed_out))
    df.drop(columns=["SPY"], inplace=True)

    if create_report:
        compare_manual_strategies(symbol, sv, sd, ed)
        compare_all_strategies(symbol, sv, sd, ed)
        sys.exit(0)

    # visualize_correlations(symbol, df)
    # plot_indicators(symbol, df)
    # compare_number_trades(symbol, sv, sd, ed)
    # compare_q_learners()

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

    ql = QLearner(testing=True, verbose=False)
    ql.addEvidence(symbol, sd, ed, sv)
    orders = ql.testPolicy(symbol, sd_out, ed_out, sv)
    df["QL"] = indicators.normalize(marketsim.compute_portvals(orders, sv))
    df["Orders QL"] = orders["Shares"]

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

    for a in ax:
        a.grid()
    m = MultiCursor(fig.canvas, ax, color='r', lw=0.5)
    plt.show()
    # fig.set_size_inches(10, 8, forward=True)
    # plt.savefig('figure_4.png', dpi=fig.dpi)


if __name__ == "__main__":
    experiment1()