import pandas as pd
import datetime as dt
import matplotlib.pyplot as plt
import TheoreticallyOptimalStrategy as tos
from util import get_data
from marketsim.marketsim import compute_portvals
from optimize_something.optimization import calculate_stats


def author():
    return "felixm"


def test_optimal_strategy():
    symbol = "JPM"
    start_value = 100000
    sd = dt.datetime(2008, 1, 1)
    ed = dt.datetime(2009, 12, 31)

    orders = tos.testPolicy(symbol=symbol, sd=sd, ed=ed, sv=start_value)
    portvals = compute_portvals(orders, start_value, 0, 0)

    start_date = portvals.index[0]
    end_date = portvals.index[-1]
    cum_ret, avg_daily_ret, \
        std_daily_ret, sharpe_ratio = calculate_stats(portvals, [1])

    d = {"Symbol": [symbol, symbol],
         "Order": ["BUY", "SELL"],
         "Shares": [1000, 1000]}
    orders = pd.DataFrame(data=d, index=[start_date, end_date])
    bench = compute_portvals(orders, start_value, 0, 0)
    cum_ret_bench, avg_daily_ret_bench, \
        std_daily_ret_bench, sharpe_ratio_bench = calculate_stats(bench, [1])

    # Compare portfolio against benchmark
    print(f"Date Range: {start_date} to {end_date}")
    print()
    print(f"Sharpe Ratio of Optimal Strategy: {sharpe_ratio}")
    print(f"Sharpe Ratio of bench: {sharpe_ratio_bench}")
    print()
    print(f"Cumulative Return of Optimal Strategy: {cum_ret}")
    print(f"Cumulative Return of bench: {cum_ret_bench}")
    print()
    print(f"Standard Deviation of Optimal Strategy: {std_daily_ret}")
    print(f"Standard Deviation of bench: {std_daily_ret_bench}")
    print()
    print(f"Average Daily Return of Optimal Strategy: {avg_daily_ret}")
    print(f"Average Daily Return of bench: {avg_daily_ret_bench}")
    print()
    print(f"Final Portfolio Value Optimal: {portvals.iloc[-1][0]}")
    print(f"Final Portfolio Value Bench: {bench.iloc[-1][0]}")

    portvals["Optimal"] = portvals["Portval"]
    portvals["Bench"] = bench["Portval"]
    portvals.drop(columns=["Portval"], inplace=True)

    portvals.plot(title="Optimal strategy versus 1000 shares of JPM")
    plt.savefig('figure_6.png')


def normalize(timeseries):
    return timeseries / timeseries.iloc[0]


def bollinger_band(df, symbol, period=20, m=2):
    boll_sma = df[symbol].rolling(period).mean()
    std = df[symbol].rolling(period).std()
    boll_up = boll_sma + m * std
    boll_lo = boll_sma - m * std
    df[f"{symbol}-Boll({period})-sma"] = boll_sma
    df[f"{symbol}-Boll({period})-up"] = boll_up
    df[f"{symbol}-Boll({period})-lo"] = boll_lo


def sma(df, symbol, period):
    """Adds a new column to the dataframe SMA(period)"""
    df[f"{symbol}-sma({period})"] = df[symbol].rolling(period).mean()


def ema(df, symbol, period):
    """Adds a new column to the dataframe EMA(period)"""
    df[f"{symbol}-ema({period})"] = df[symbol].ewm(span=period).mean()


def price_sma(df, symbol, period):
    """Calculates SMA and adds new column price divided by SMA to the df."""
    sma = df[symbol].rolling(period).mean()
    df[f"{symbol}-price/sma({period})"] = df[symbol] / sma


def rsi(df, symbol, period=14):
    """Calculates relative strength index over given period."""

    def rsi(x):
        pct = x.pct_change()
        avg_gain = pct[pct > 0].mean()
        avg_loss = pct[pct <= 0].abs().mean()
        rsi = 100 - (100 /
                     (1 + ((avg_gain / period) /
                           (avg_loss / period))))
        return rsi

    key = f"{symbol}-rsi({period})"
    # Add one to get 'period' price changes (first change is nan).
    period += 1
    df[key] = df[symbol].rolling(period).apply(rsi)


def macd(df, symbol):
    macd = df[symbol].ewm(span=12).mean() - df[symbol].ewm(span=26).mean()
    df[f"{symbol}-macd"] = macd
    df[f"{symbol}-macd-signal"] = macd.rolling(9).mean()


def price_delta(df, symbol, period=1):
    """Calculate delta between previous day and today."""
    df[f"{symbol}-diff({period})"] = df[symbol].diff(periods=period)


def test_indicators():
    symbol = "JPM"

    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 = normalize(df)

    sma(df, symbol, 21)
    ema(df, symbol, 21)
    df.plot(title="21 SMA and EMA")
    plt.savefig('figure_1.png')

    df = df_orig.copy()
    sma(df, symbol, 8)
    price_sma(df, symbol, 8)
    df.plot(title="SMA and price / SMA", subplots=True)
    plt.savefig('figure_2.png')

    df = df_orig.copy()
    bollinger_band(df, symbol)
    df.plot(title="Bollinger Band")
    plt.savefig('figure_3.png')

    df = df_orig.copy()
    rsi(df, symbol)
    fig, axes = plt.subplots(nrows=2, sharex=True)
    df[symbol].plot(ax=axes[0], title="JPM price action")
    df["JPM-rsi(14)"].plot(ax=axes[1], title="RSI")
    plt.savefig('figure_4.png')

    df = df_orig.copy()
    macd(df, symbol)
    fig, axes = plt.subplots(nrows=2, sharex=True)
    df[symbol].plot(ax=axes[0], title="JPM price action")
    df[["JPM-macd", "JPM-macd-signal"]].plot(ax=axes[1])
    plt.savefig('figure_5.png')


def main():
    test_optimal_strategy()
    test_indicators()


if __name__ == "__main__":
    main()