import numpy as np
from matplotlib import pyplot as plt
import pandas as pd

df = get_pricing(symbols('SPY'), start_date='2018-01-01', end_date='2020-01-01')

Price = df.close_price
Length = len(df)
InputPrice = Price.copy()
batch = 10

# Initialize output before the algorithm
Filt = np.array(InputPrice)


# sequencially calculate all variables and the output
for i in range(0, Length, 1):
    # If there's not enough data, Filt is the price - whitch it already is, so just skip
    if i < 2 * batch:
        continue
    # take 2 batches of the input
    v1 = InputPrice[i-2*batch:i - batch]
    v2 = InputPrice[i - batch:i]

    # for the 1st batch calculate N1
    H1 = np.max(v1)
    L1 = np.min(v1)
    N1 = (H1 - L1) / batch

    # for the 2nd batch calculate N2
    H2 = np.max(v2)
    L2 = np.min(v2)
    N2 = (H2 - L2) / batch

    # for both batches calculate N3
    H = np.max([H1, H2])
    L = np.min([L1, L2])
    N3 = (H - L) / (2 * batch)

    # calculate fractal dimension
    Dimen = 0
    if N1 > 0 and N2 > 0 and N3 > 0:
        Dimen = (np.log(N1 + N2) - np.log(N3)) / np.log(2)

    # calculate lowpass filter factor
    alpha = np.exp(-4.6 * (Dimen - 1)) # this was wrongly calculated as np.exp(-4.6 * (Dimen) - 1)
    alpha = np.max([alpha, 0.01]) # here it was 0.1 but Ehler clips at 0.01
    alpha = np.min([alpha, 1])

    # filter the input data
    Filt[i] = alpha * InputPrice[i] + (1 - alpha) * Filt[i-1]
    # if currentBar < 2*batch + 1: <--- i dont get what these 2 lines do
    # Filt = InputPrice[i]

def pd_frama(c, per):
    c = c.copy()
    window = per * 2

    hh = c.rolling(per).max()
    ll = c.rolling(per).min()

    n1 = (hh - ll) / per
    n2 = n1.shift(per)

    hh2 = c.rolling(window).max()
    ll2 = c.rolling(window).min()
    n3 = (hh2 - ll2) / window

    D = (np.log(n1 + n2) - np.log(n3)) / np.log(2)
    alp = np.exp(-4.6 * (D - 1))
    alp = np.clip(alp, .01, 1).values

    filt = c.values
    for i, x in enumerate(alp):
        cl = c.values[i]
        if i < window:
            continue
        filt[i] = cl * x + (1 - x) * filt[i - 1]

    return filt

df['filt'] = Filt
df.filt.plot(color='y')
Price.plot(color='k')

filt2 = pd_frama(df.close_price, batch)

filt2_df = pd.DataFrame(filt2, index=df.index)
filt2_df.plot(color='y')
Price.plot(color='k')

df2 = get_pricing(symbols(['SPY', 'QQQ']), start_date='2018-01-01', end_date='2020-01-01', fields='close_price')

df2.head()

filt3 = pd_frama(df2, 20)
filt_df = pd.DataFrame(filt3, index=df2.index, columns=df2.columns)