"""VHM demo: diagnostics (residual-based fault detection) and prognostics (RUL).

Diagnostics: a CUSUM detector on the innovation residual (measured SOH minus
model estimate). The detector is SELF-BASELINED: residual mean/std are
calibrated on the first BASELINE_N cycles of the monitored battery itself
(assumed healthy), which removes cross-battery model bias. A synthetic
accelerated-fade fault is injected into one held-out battery; false-positive
behavior is checked on all healthy held-out batteries.

Prognostics: the physics network f_phi is an autonomous degradation law
dSOH/dt = f(C-rate, t, SOH). From a chosen in-life cycle it is
forward-integrated (Euler) to the 80% EOL threshold -> predicted RUL,
compared against the true remaining cycles.
"""
import argparse
import json
import os

import jax.numpy as jnp
import numpy as np
import pandas as pd

import data as D
import model as M

BASELINE_N = 60     # healthy cycles used to self-calibrate the detector
CUSUM_K = 5.0       # slack, in sigma units
CUSUM_H = 100.0     # decision threshold, in sigma units
SIGMA_FLOOR = 0.008  # min residual sigma (SOH units), robustifies self-calibration
FAULT_START_FRAC = 0.6
FAULT_RATE = 5e-4   # extra SOH loss per cycle after fault onset


def innovation(params, bat):
    u = np.array(M.u_batch(params["u"], jnp.array(bat["x"]), jnp.array(bat["t"])))
    return bat["soh"] - u, u


def cusum_down(r, mu, sigma, k=CUSUM_K, h=CUSUM_H):
    """One-sided CUSUM for a downward mean shift; returns alarm index or -1."""
    z = (r - mu) / sigma
    s, path, alarm = 0.0, [], -1
    for i, zi in enumerate(z):
        s = max(0.0, s - zi - k)
        path.append(s)
        if s > h and alarm < 0:
            alarm = i
    return alarm, np.array(path)


def monitor(params, bat, soh_override=None):
    """Self-baselined CUSUM monitoring of one battery."""
    b = bat if soh_override is None else {**bat, "soh": soh_override}
    r, u = innovation(params, b)
    mu = float(r[:BASELINE_N].mean())
    sigma = max(float(r[:BASELINE_N].std()), SIGMA_FLOOR)
    alarm, path = cusum_down(r[BASELINE_N:], mu, sigma)
    alarm_cycle = int(b["cycle"][BASELINE_N + alarm]) if alarm >= 0 else None
    return r, u, path, alarm_cycle


def diagnostics(params, test_n, out):
    # false positives on healthy held-out batteries
    fp = {b["name"]: monitor(params, b)[3] for b in test_n}

    # inject accelerated fade into the first held-out battery
    bat = test_n[0]
    n = len(bat["soh"])
    k0 = int(FAULT_START_FRAC * n)
    fault = np.zeros(n, dtype=np.float32)
    fault[k0:] = FAULT_RATE * np.arange(1, n - k0 + 1, dtype=np.float32)
    soh_f = (bat["soh"] - fault).astype(np.float32)
    r_f, u, path, alarm_cycle = monitor(params, bat, soh_override=soh_f)
    onset = int(bat["cycle"][k0])
    delay = alarm_cycle - onset if alarm_cycle is not None else None

    pad = np.full(BASELINE_N, np.nan)
    pd.DataFrame({
        "cycle": bat["cycle"], "soh_healthy": bat["soh"], "soh_faulted": soh_f,
        "soh_model": u, "residual": r_f,
        "cusum": np.concatenate([pad, path]),
    }).to_csv(os.path.join(out, "diagnostics_trace.csv"), index=False)

    return {
        "battery": bat["name"],
        "baseline_cycles": BASELINE_N,
        "fault_onset_cycle": onset,
        "fault_rate_per_cycle": FAULT_RATE,
        "alarm_cycle": alarm_cycle,
        "detection_delay_cycles": delay,
        "false_alarm_cycle_healthy": fp,  # None = no false alarm over full life
        "cusum_k": CUSUM_K, "cusum_h": CUSUM_H,
    }


import functools
import jax


@functools.partial(jax.jit, static_argnums=(4,))
def _integrate(params, c, t0, soh0, n_steps):
    dt = 1.0 / D.T_SCALE

    def step(carry, _):
        t, soh = carry
        rate = M.f_batch(params["f"], c[None], jnp.array([t]), jnp.array([soh]))[0]
        soh2 = jnp.where(soh <= D.EOL_SOH, soh, soh + rate * dt)  # stop at EOL
        return (t + dt, soh2), soh2

    _, sohs = jax.lax.scan(step, (t0, soh0), None, length=n_steps)
    return sohs


def forecast_rul(params, bat, frac, max_cycles=3000):
    """Euler-integrate dSOH/dt = f(C-rate, SOH) from cycle k to EOL."""
    n = len(bat["soh"])
    k = int(frac * n)
    t0 = float(bat["t"][k])
    soh0 = float(M.u_batch(params["u"], jnp.array(bat["x"][k])[None, :],
                           jnp.array([t0]))[0])
    sohs = np.array(_integrate(params, jnp.array(bat["c"][0]), t0, soh0, max_cycles))
    cyc = bat["cycle"][k] + np.arange(1, max_cycles + 1)
    hit = np.argmax(sohs <= D.EOL_SOH)
    end = hit if sohs[hit] <= D.EOL_SOH else max_cycles - 1
    traj = np.column_stack([np.concatenate([[bat["cycle"][k]], cyc[:end + 1]]),
                            np.concatenate([[soh0], sohs[:end + 1]])])
    pred_eol = traj[-1, 0]
    true_eol = float(bat["cycle"][-1])  # test protocol stops at 80% SOH
    cyc_k = float(bat["cycle"][k])
    return {
        "battery": bat["name"], "predict_at_cycle": cyc_k,
        "life_fraction": frac,
        "pred_EOL_cycle": round(pred_eol, 1), "true_EOL_cycle": true_eol,
        "pred_RUL": round(pred_eol - cyc_k, 1), "true_RUL": round(true_eol - cyc_k, 1),
        "RUL_error_cycles": round(pred_eol - true_eol, 1),
    }, traj


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--data_dir", required=True)
    p.add_argument("--run", default="../runs/poc")
    args = p.parse_args()

    params = M.load_params(os.path.join(args.run, "params_pinn.pkl"))
    norm = D.load_norm(os.path.join(args.run, "norm.json"))
    batteries = D.load_dataset(args.data_dir)
    _, val, test = D.split_dataset(batteries)
    test_n = [D.apply_norm(b, norm) for b in test]

    diag = diagnostics(params, test_n, args.run)
    print("diagnostics:", json.dumps(diag, indent=2))

    prog, trajs = [], {}
    for bat in test_n:
        for frac in (0.3, 0.5, 0.7):
            row, traj = forecast_rul(params, bat, frac)
            prog.append(row)
            trajs[f"{bat['name']}@{frac}"] = traj
    pd.DataFrame(prog).to_csv(os.path.join(args.run, "prognostics.csv"), index=False)
    np.savez(os.path.join(args.run, "forecast_trajs.npz"), **trajs)
    print(pd.DataFrame(prog).to_string(index=False))

    with open(os.path.join(args.run, "vhm_results.json"), "w") as f:
        json.dump({"diagnostics": diag, "prognostics": prog}, f, indent=2)


if __name__ == "__main__":
    main()
