import pandas as pd
import datetime as dt
import matplotlib.pyplot as plt
from util import get_data


def author():
    return "felixm"


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
    key_sma, key_up, key_lo = "boll_sma", "boll_up", "boll_lo"
    df[key_sma] = boll_sma
    df[key_up] = boll_up
    df[key_lo] = boll_lo
    return df[[key_sma, key_up, key_lo]]


def sma(df, symbol, period):
    """Adds SMA for one or multiple periods to df and returns SMAs"""
    if type(period) is int:
        period = [period]
    keys = []
    for p in period:
        key = f"sma_{p}"
        df[key] = df[symbol].rolling(p).mean()
        keys.append(key)
    return df[keys]


def ema(df, symbol, period):
    """Adds EMA for one or multiple periods to df and returns EMAs"""
    if type(period) is int:
        period = [period]
    keys = []
    for p in period:
        key = f"ema_{p}"
        df[key] = df[symbol].ewm(span=p).mean()
        keys.append(key)
    return df[keys]


def price_sma(df, symbol, period):
    """Calculates SMA and adds new column price divided by SMA to the df."""
    if type(period) is int:
        period = [period]
    keys = []
    for p in period:
        key = f"price_sma_{p}"
        sma = df[symbol].rolling(p).mean()
        df[key] = df[symbol] / sma
        keys.append(key)
    return df[keys]


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 = "rsi"
    # Add one to get 'period' price changes (first change is nan).
    period += 1
    df[key] = df[symbol].rolling(period).apply(rsi)
    return df[[key]]


def macd(df, symbol):
    macd = df[symbol].ewm(span=12).mean() - df[symbol].ewm(span=26).mean()
    k1 = "macd"
    k2 = "macd_signal"
    k3 = "macd_diff"
    df[k1] = macd
    df[k2] = macd.rolling(9).mean()
    df[k3] = df[k1] - df[k2]
    return df[[k1, k2, k3]]


def price_delta(df, symbol, period=1):
    """Calculate percentage change for period."""
    k = f"pct_{period}"
    df[k] = df[symbol].pct_change(periods=period)
    df[k] = df[k].shift(-period)
    return df[k]


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