"""
run_electricity.py
==================
Electricity demand experiment (UCI Household Power Consumption).
Tests K-R vs Tuned ESN vs DeepESN-2L on a real-world periodic signal.
This is a deliberate negative control confirming the boundary condition.

Data required:
    household_power_consumption.txt
    Download from: https://archive.ics.uci.edu/dataset/235

Run:
    python run_electricity.py

Saves: results_electricity.pkl, figF_electricity.png, figG_elec_bar.png
"""

import numpy as np
import itertools
import pickle
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import stats

from kr_reservoir import KRReservoir
from esn_baseline import StandardESN, DeepESN2Layer


# ── Load and preprocess ────────────────────────────────────────
def load_electricity(path: str = 'household_power_consumption.txt',
                     n_hours: int = 26280) -> np.ndarray:
    """
    Load UCI Household Power Consumption dataset.
    Returns last n_hours of hourly-resampled, normalised Global_active_power.
    """
    try:
        import pandas as pd
        df = pd.read_csv(path, sep=';', na_values=['?'], low_memory=False)
        df['datetime'] = pd.to_datetime(
            df['Date'] + ' ' + df['Time'], dayfirst=True, format='mixed')
        df = df.set_index('datetime')
        df['Global_active_power'] = pd.to_numeric(
            df['Global_active_power'], errors='coerce')
        series = (df['Global_active_power']
                  .dropna()
                  .resample('h')
                  .mean()
                  .dropna()
                  .values)
        series = series[-n_hours:]
        series = (series - series.mean()) / series.std()
        print(f"Loaded {len(series)} hourly readings.")
        return series
    except Exception as e:
        raise FileNotFoundError(
            f"Could not load {path}. "
            "Download from: https://archive.ics.uci.edu/dataset/235\n"
            f"Error: {e}")


# ── Config ────────────────────────────────────────────────────
N        = 200
N_SEEDS  = 10
WASHOUT  = 500
TEST_SZ  = 720    # 30 days
M        = 168    # weekly memory order
AS       = 3.0 / M
VAL_SEED = 99
KF       = 1.5 * (200 / N) ** 0.40
KS       = 1.0 * (200 / N) ** 0.40
DELAYS   = list(range(12, M + 1, 12))  # 12,24,...,168h

RHO_GRID = [0.5, 0.7, 0.9, 0.95, 0.99]
LAM_GRID = [1e-6, 1e-4, 1e-2, 1.0, 10.0]


def gs(u, y, method):
    half = len(u) // 2; best = (np.inf, None)
    for rho, lam in itertools.product(RHO_GRID, LAM_GRID):
        if method == 'std':
            m = StandardESN(N=N, rho=rho, sigma=0.1, lam=lam,
                            seed=VAL_SEED, washout=WASHOUT)
        elif method == 'kr':
            m = KRReservoir(N=N, M=M, rho_f=rho, rho_s=rho,
                            alpha_s=AS, lam=lam, seed=VAL_SEED, washout=WASHOUT)
        elif method == 'deep2':
            m = DeepESN2Layer(N=N, rho1=rho, rho2=rho, lam=lam,
                              seed=VAL_SEED, washout=WASHOUT)
        m.fit(u[:half], y[:half])
        n, _, _ = m.evaluate(u[:half], y[:half], 200)
        if n < best[0]: best = (n, (rho, lam))
    return best[1]  # (rho, lam)


print("=" * 55)
print("ELECTRICITY DEMAND EXPERIMENT")
print("=" * 55)

series = load_electricity()
u_in = series[:-1]; y_out = series[1:]

# Grid search
print("\nGrid search...", flush=True)
sp = gs(u_in, y_out, 'std')
kp = gs(u_in, y_out, 'kr')
dp = gs(u_in, y_out, 'deep2')
print(f"  ESN:   rho={sp[0]:.2f} lam={sp[1]:.0e}")
print(f"  K-R:   rho={kp[0]:.2f} lam={kp[1]:.0e}")
print(f"  Deep2: rho={dp[0]:.2f} lam={dp[1]:.0e}")

# 10-seed evaluation
std_sc = []; kr_sc = []; deep_sc = []
yte_s = yp_s = yp_k = yp_d = None

print("\n10-seed evaluation...")
for seed in range(N_SEEDS):
    ms = StandardESN(N=N, rho=sp[0], sigma=0.1, lam=sp[1],
                     seed=seed, washout=WASHOUT)
    ms.fit(u_in, y_out); ns, yt, yp = ms.evaluate(u_in, y_out, TEST_SZ)

    mk = KRReservoir(N=N, M=M, rho_f=kp[0], rho_s=kp[0],
                     alpha_s=AS, lam=kp[1], seed=seed, washout=WASHOUT)
    mk.fit(u_in, y_out); nk, _, ypk = mk.evaluate(u_in, y_out, TEST_SZ)

    md = DeepESN2Layer(N=N, rho1=dp[0], rho2=dp[0], lam=dp[1],
                       seed=seed, washout=WASHOUT)
    md.fit(u_in, y_out); nd, _, ypd = md.evaluate(u_in, y_out, TEST_SZ)

    std_sc.append(ns); kr_sc.append(nk); deep_sc.append(nd)
    if seed == 0:
        yte_s = yt; yp_s = yp; yp_k = ypk; yp_d = ypd
    print(f"  seed {seed}: ESN={ns:.4f}  K-R={nk:.4f}  Deep={nd:.4f}")

sm, ss = np.mean(std_sc), np.std(std_sc)
km, ks = np.mean(kr_sc),  np.std(kr_sc)
dm, ds = np.mean(deep_sc), np.std(deep_sc)

_, p1 = stats.ttest_ind(std_sc, kr_sc)
_, p2 = stats.ttest_ind(std_sc, deep_sc)

def sig(p): return '***' if p < 0.001 else('**' if p < 0.01 else('*' if p < 0.05 else 'ns'))

print(f"\nTuned ESN:   {sm:.4f} ± {ss:.4f}")
print(f"DeepESN-2L:  {dm:.4f} ± {ds:.4f}  vs ESN: {(sm-dm)/sm*100:+.1f}%  {sig(p2)}")
print(f"K-R:         {km:.4f} ± {ks:.4f}  vs ESN: {(sm-km)/sm*100:+.1f}%  {sig(p1)}")
print("\n⇒ ESN wins: both K-R and DeepESN degrade on periodic real-world data.")
print("  This confirms the boundary condition.")


# ── Figures ──────────────────────────────────────────────────
# Figure F: Prediction traces (1-week window)
fig, axes = plt.subplots(3, 1, figsize=(9, 6), sharex=True)
sl = slice(0, 168); ta = np.arange(168)

for ax, yp, label, color, nrmse in zip(
        axes,
        [yp_s, yp_d, yp_k],
        ['Tuned ESN', 'DeepESN (2-layer)', 'K-R (proposed)'],
        ['#4C72B0', '#9467BD', '#2CA02C'],
        [sm, dm, km]):
    ax.plot(ta, yte_s[sl], 'k-', lw=1.2, alpha=0.7, label='Ground truth')
    ax.plot(ta, yp[sl], '--', color=color, lw=1.0,
            label=f'{label}  NRMSE={nrmse:.4f}')
    ax.legend(fontsize=7.5, loc='upper right')
    ax.grid(alpha=0.3, lw=0.5)
    ax.set_ylabel('Norm. power', fontsize=8)

axes[-1].set_xlabel('Hours in test window (1 week)', fontsize=9)
fig.suptitle('Electricity Demand: 1-Week Ahead Prediction\n'
             '(Periodic signal — ESN captures daily/weekly cycle best)',
             fontsize=10, fontweight='bold')
plt.tight_layout()
plt.savefig('figF_electricity.png', dpi=180, bbox_inches='tight', facecolor='white')
plt.close()

# Figure G: Bar chart
fig, ax = plt.subplots(figsize=(5, 4))
vals = [sm, dm, km]; errs = [ss, ds, ks]
colors = ['#4C72B0', '#9467BD', '#D62728']
ax.bar([0, 1, 2], vals, yerr=errs, capsize=5,
       color=colors, edgecolor='#1a1a1a', lw=0.8, width=0.55)
for i, (v, e) in enumerate(zip(vals, errs)):
    ax.text(i, v + e + 0.004, f'{v:.4f}', ha='center', va='bottom',
            fontsize=9, fontweight='bold')
ymax = max(v + e for v, e in zip(vals, errs)) + 0.03
ax.plot([0, 2], [ymax, ymax], 'k-', lw=1.0)
ax.text(1, ymax + 0.005, sig(p1), ha='center', fontsize=11, fontweight='bold')
ax.set_xticks([0, 1, 2])
ax.set_xticklabels(['Tuned ESN', 'DeepESN\n(2-layer)', 'K-R\n(proposed)'], fontsize=9)
ax.set_ylabel('NRMSE (lower = better)', fontsize=9)
ax.set_title('Electricity Demand (N=200, 10 seeds)\n'
             'Periodic signal — boundary condition confirmed', fontsize=9)
ax.grid(axis='y', alpha=0.3)
plt.tight_layout()
plt.savefig('figG_elec_bar.png', dpi=180, bbox_inches='tight', facecolor='white')
plt.close()

print("\nFigures saved: figF_electricity.png, figG_elec_bar.png")

with open('results_electricity.pkl', 'wb') as f:
    pickle.dump(dict(sm=sm, ss=ss, km=km, ks=ks, dm=dm, ds=ds,
                     std_sc=std_sc, kr_sc=kr_sc, deep_sc=deep_sc), f)
print("Results saved: results_electricity.pkl")