"""
Part 2: choose an ARMA model on healthy data and check that it is good enough.
Requires: numpy, pandas, matplotlib, statsmodels   (pip install statsmodels)
"""
import warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.stattools import adfuller
from statsmodels.stats.diagnostic import acorr_ljungbox
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf

warnings.filterwarnings("ignore")        # hide convergence chatter in the model grid

N_HEALTHY, N_FAULT = 400, 160            # minutes

# ---------------------------------------------------------------- 1. data

def simulate(seed=8, n_healthy=N_HEALTHY, n_fault=N_FAULT, slope=0.003):
    """Thermostat temperature: an ARMA(2,1) process around 21 C.
    After minute `n_healthy` the heater loses power: a heating shortfall
    that grows by `slope` degrees every minute."""
    rng = np.random.default_rng(seed)
    n, burn = n_healthy + n_fault, 300
    e = rng.normal(0, 0.10, n + burn)
    x = np.zeros(n + burn)
    for t in range(n + burn):
        k = t - burn - n_healthy
        fault = -slope * k if k > 0 else 0.0
        ar = 0.55 * x[t - 1] + 0.30 * x[t - 2] if t >= 2 else 0.0
        ma = 0.50 * e[t - 1] if t >= 1 else 0.0
        x[t] = ar + e[t] + ma + fault
    return 21.0 + x[burn:]

if __name__ == "__main__":
    y = simulate()
    y_healthy = y[:N_HEALTHY]            # Part 2 uses the healthy data only

    # ------------------------------------------------------------ 2. choose d
    adf_p = adfuller(y_healthy)[1]
    print(f"ADF p-value: {adf_p:.3f}   (p < 0.05 -> stationary -> d = 0)")

    # ------------------------------------------------------------ 3. first guess for p and q
    fig, ax = plt.subplots(1, 2, figsize=(10, 3.5))
    plot_acf(y_healthy, lags=20, ax=ax[0], zero=False)
    plot_pacf(y_healthy, lags=20, ax=ax[1], zero=False, method="ywm")
    plt.show()

    # ------------------------------------------------------------ 4. compare models with AIC / BIC
    rows = []
    for p in range(4):
        for q in range(4):
            res = ARIMA(y_healthy, order=(p, 0, q), trend="c").fit()
            lb = acorr_ljungbox(res.resid[5:], lags=[10], model_df=p + q)
            rows.append(dict(p=p, q=q, k=p + q + 2, aic=res.aic, bic=res.bic,
                             lb_pvalue=float(lb["lb_pvalue"].iloc[0])))
    table = pd.DataFrame(rows).sort_values("aic").reset_index(drop=True)
    print(table.head(8).round(3).to_string(index=False))

    # ------------------------------------------------------------ 5. final model + whiteness check
    model = ARIMA(y_healthy, order=(1, 0, 1), trend="c").fit()
    print(model.summary().tables[1])     # const = mean, ar.L1, ma.L1, sigma2
    lb = acorr_ljungbox(model.resid[5:], lags=[10], model_df=2)
    print("Ljung-Box p-value:", float(lb["lb_pvalue"].iloc[0]), "(above 0.05 -> residuals look white)")