"""
Part 3: use the model fitted on healthy data to catch a fault.
Requires: numpy, matplotlib, statsmodels   (pip install statsmodels)
"""
import warnings
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.arima.model import ARIMA

warnings.filterwarnings("ignore")

N_HEALTHY, N_FAULT = 400, 160            # minutes

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:]


def detect(y, n_healthy=N_HEALTHY, verbose=True):
    # model from Part 2: ARMA(1,1), fitted on the healthy data only
    model = ARIMA(y[:n_healthy], order=(1, 0, 1), trend="c").fit()

    # apply the SAME model to all the data
    full = model.apply(y)
    resid = np.asarray(full.resid)
    sigma = resid[5:n_healthy].std()                   # typical error while healthy
    z = resid / sigma                                  # residual in sigmas

    # alarm 1: simple +-3 sigma threshold
    beyond = np.where(np.abs(z[5:]) > 3)[0] + 5

    # alarm 2: CUSUM (watches for a downward shift)
    k, h = 0.5, 8.0
    S = np.zeros(len(z))
    for t in range(6, len(z)):
        S[t] = max(0.0, S[t - 1] - z[t] - k)
    alarms = np.where(S > h)[0]
    alarms = alarms[alarms >= 6]
    first = int(alarms[0]) if len(alarms) else None

    if verbose:
        print(f"sigma = {sigma:.3f} C")
        print("3-sigma exceedances while healthy :", [int(i) for i in beyond if i < n_healthy])
        print("3-sigma exceedances after fault   :", [int(i) for i in beyond if i >= n_healthy][:3], "...")
        print("CUSUM first alarm at minute       :", first,
              f"({first - n_healthy} min after the fault began)" if first else "")
    return dict(z=z, S=S, first_alarm=first, beyond=beyond)


def robustness(n_seeds=30):
    """Is the result a lucky run? Repeat with different random seeds."""
    fa_3sigma, delays = 0, []
    for seed in range(1, n_seeds + 1):
        out = detect(simulate(seed=seed), verbose=False)
        fa_3sigma += any(i < N_HEALTHY for i in out["beyond"])
        if out["first_alarm"] is not None:
            delays.append(out["first_alarm"] - N_HEALTHY)
    print(f"runs with a +-3 sigma false alarm in the healthy period: {fa_3sigma}/{n_seeds}")
    print(f"CUSUM detection delay (minutes): median {np.median(delays):.0f}, range {min(delays)}-{max(delays)}")


if __name__ == "__main__":
    y = simulate()
    out = detect(y)

    fig, ax = plt.subplots(3, 1, figsize=(9, 8), sharex=True)
    ax[0].plot(y);          ax[0].set_ylabel("Temperature (C)")
    ax[1].plot(out["z"]);   ax[1].axhline(3, ls=":"); ax[1].axhline(-3, ls=":"); ax[1].set_ylabel("Residual (sigma)")
    ax[2].plot(out["S"]);   ax[2].axhline(8, ls=":"); ax[2].set_ylabel("CUSUM"); ax[2].set_ylim(-1, 40)
    for a in ax:
        a.axvline(N_HEALTHY, color="red")
    ax[2].set_xlabel("Time (minutes)")
    plt.tight_layout(); plt.show()