import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.stattools import acf

# Load datasets
df1 = pd.read_csv('../ex_data/ex_data_ewma09_rho0.csv')
df2 = pd.read_csv('../ex_data/ex_data_ewma09_rho01.csv')
df3 = pd.read_csv('../ex_data/ex_data_ewma09_rho06.csv')
df4 = pd.read_csv('../ex_data/ex_data_ewma09_rho09.csv')

labels = [
    r'$\rho=0$',
    r'$\rho=0.1$',
    r'$\rho=0.6$',
    r'$\rho=0.9$',
]

styles = [
    ('#377eb8', 'dotted', '.'),
    ('#ff7f00', 'solid', '.'),
    ('#4daf4a', 'dashed', '.'),
    ('#f781bf', '-.', '.'),
    ('#a65628', 'solid', 'x'),
    ('#984ea3', 'solid', 'x'),
    ('#e41a1c', 'solid', 'x'),
]

feats = ['x1','x3','y']

fig, axes = plt.subplots(1, 3, figsize=(28, 4), sharey=True)
handles = []

for b, target in enumerate(feats):
    ax = axes[b]
    for df, (color, ls, marker), label in zip([df1, df2, df3, df4],
                                              styles, labels):
        series = df[target].dropna()
        autocorr_values = acf(series, nlags=100, fft=True)
        h, = ax.plot(autocorr_values,
                     color=color,
                     linewidth=2,
                     linestyle=ls,
                     marker=marker,
                     markevery=25,
                     label=label)
        if b == 0:
            handles.append(h)

    ax.set_xlim(1, 100)
    ax.set_ylim(-0.2, 1.02)
    ax.grid(True, linestyle='--', alpha=0.6)
    ax.set_xlabel('Lag')

axes[0].set_ylabel(r'Autocorrelation ($\alpha=0.9$)')

plt.savefig('autocorr_plot_ewma09.svg', format='svg', bbox_inches='tight')
plt.show()
plt.close(fig)
