"""Train the SOH models.

baseline: solution net u(x,t) trained on data loss only.
pinn: staged physics-informed training —
  stage 1: u initialized from the trained baseline (data fit),
  stage 2: degradation law f(c,t,u) regression-fitted to u's empirical
           cycle-to-cycle SOH rates, with a monotonicity penalty (f <= 0),
  stage 3: joint fine-tune coupling u's increments to f (increment space,
           guarded by best-validation selection).
The result is an accurate estimator u plus an integrable degradation law f.
"""
import argparse
import json
import os
import time

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

import data as D
import model as M


def metrics(y_true, y_pred):
    e = y_pred - y_true
    return {
        "MAPE_%": float(np.mean(np.abs(e / y_true)) * 100),
        "MAE": float(np.mean(np.abs(e))),
        "RMSE": float(np.sqrt(np.mean(e**2))),
    }


def val_mse(params, xv, tv, yv):
    uv = M.u_batch(params["u"], jnp.array(xv), jnp.array(tv))
    return float(jnp.mean((uv - jnp.array(yv)) ** 2))


def train_baseline(pairs, xv, tv, yv, steps, seed, lr=1e-3):
    params = M.init_params(jax.random.PRNGKey(seed))
    x1, t1, y1, x2, t2, y2, c = (jnp.array(a) for a in pairs)
    opt = optax.adam(lr)
    opt_state = opt.init(params["u"])

    def loss(u_params):
        u1 = M.u_batch(u_params, x1, t1)
        u2 = M.u_batch(u_params, x2, t2)
        return 0.5 * jnp.mean((u1 - y1) ** 2) + 0.5 * jnp.mean((u2 - y2) ** 2)

    @jax.jit
    def step(u_params, opt_state):
        l, g = jax.value_and_grad(loss)(u_params)
        upd, opt_state = opt.update(g, opt_state)
        return optax.apply_updates(u_params, upd), opt_state, l

    hist, best, best_u = [], np.inf, params["u"]
    for i in range(steps):
        params["u"], opt_state, l = step(params["u"], opt_state)
        if i % 100 == 0 or i == steps - 1:
            vm = val_mse(params, xv, tv, yv)
            hist.append({"step": i, "loss": float(l), "data": float(l),
                         "phys": 0.0, "mono": 0.0, "val_mse": vm})
            if vm < best:
                best, best_u = vm, jax.device_get(params["u"])
    params["u"] = best_u
    return params, hist


def train_pinn(base_params, pairs, rates, xv, tv, yv, steps_f, steps_ft,
               alpha, beta, seed, lr=1e-3):
    params = M.init_params(jax.random.PRNGKey(seed + 100))
    params["u"] = base_params["u"]
    x1, t1, y1, x2, t2, y2, c = (jnp.array(a) for a in pairs)
    dt = t2 - t1

    # ---- stage 2: fit f to smoothed fleet degradation rates ----
    rc, rt, ru, rr = (jnp.array(a) for a in rates)

    opt_f = optax.adam(lr)
    st_f = opt_f.init(params["f"])

    def loss_f(f_params):
        f = M.f_batch(f_params, rc, rt, ru)
        return jnp.mean((rr - f) ** 2) + beta * jnp.mean(jax.nn.relu(f) ** 2)

    @jax.jit
    def step_f(f_params, st):
        l, g = jax.value_and_grad(loss_f)(f_params)
        upd, st = opt_f.update(g, st)
        return optax.apply_updates(f_params, upd), st, l

    hist = []
    for i in range(steps_f):
        params["f"], st_f, l = step_f(params["f"], st_f)
        if i % 100 == 0 or i == steps_f - 1:
            hist.append({"step": i, "loss": float(l), "data": 0.0,
                         "phys": float(l), "mono": 0.0,
                         "val_mse": val_mse(params, xv, tv, yv)})

    # ---- stage 3: guarded joint fine-tune in increment space ----
    opt_j = optax.adam(lr * 0.1)
    st_j = opt_j.init(params)

    def loss_j(p):
        u1 = M.u_batch(p["u"], x1, t1)
        u2 = M.u_batch(p["u"], x2, t2)
        # stop-gradient: the physics term regularizes u toward the fitted
        # degradation law; it must not drag f toward u's noisy increments
        f = jax.lax.stop_gradient(M.f_batch(p["f"], c, t1, u1))
        l_data = 0.5 * jnp.mean((u1 - y1) ** 2) + 0.5 * jnp.mean((u2 - y2) ** 2)
        l_inc = jnp.mean(((u2 - u1) - f * dt) ** 2)
        l_mono = jnp.mean(jax.nn.relu(f) ** 2)
        return l_data + alpha * l_inc + beta * l_mono, (l_data, l_inc, l_mono)

    @jax.jit
    def step_j(p, st):
        (l, parts), g = jax.value_and_grad(loss_j, has_aux=True)(p)
        upd, st = opt_j.update(g, st)
        return optax.apply_updates(p, upd), st, l, parts

    best, best_p = val_mse(params, xv, tv, yv), jax.device_get(params)
    for i in range(steps_ft):
        params, st_j, l, parts = step_j(params, st_j)
        if i % 100 == 0 or i == steps_ft - 1:
            vm = val_mse(params, xv, tv, yv)
            hist.append({"step": steps_f + i, "loss": float(l),
                         "data": float(parts[0]), "phys": float(parts[1]),
                         "mono": float(parts[2]), "val_mse": vm})
            if vm < best:
                best, best_p = vm, jax.device_get(params)
    return best_p, hist


def evaluate(tag, params, test_n, out, results, extra):
    per_batt, preds = {}, []
    for bat in test_n:
        u = np.array(M.u_batch(params["u"], jnp.array(bat["x"]), jnp.array(bat["t"])))
        preds.append(u)
        per_batt[bat["name"]] = metrics(bat["soh"], u)
        pd.DataFrame({"cycle": bat["cycle"], "soh_true": bat["soh"],
                      f"soh_{tag}": u}).to_csv(
            os.path.join(out, "predictions", f"{bat['name']}_{tag}.csv"), index=False)
    agg = metrics(np.concatenate([b["soh"] for b in test_n]), np.concatenate(preds))
    results[tag] = {"per_battery": per_batt, "aggregate": agg, **extra}
    print(tag, "aggregate:", agg)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--data_dir", required=True)
    p.add_argument("--out", default="../runs/poc")
    p.add_argument("--steps", type=int, default=4000)
    p.add_argument("--steps_f", type=int, default=3000)
    p.add_argument("--steps_ft", type=int, default=1000)
    p.add_argument("--alpha", type=float, default=1000.0)
    p.add_argument("--beta", type=float, default=0.1)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--models", default="baseline,pinn")
    args = p.parse_args()
    os.makedirs(os.path.join(args.out, "predictions"), exist_ok=True)

    batteries = D.load_dataset(args.data_dir)
    train, val, test = D.split_dataset(batteries)
    norm = D.fit_norm(train)
    D.save_norm(norm, os.path.join(args.out, "norm.json"))
    train_n = [D.apply_norm(b, norm) for b in train]
    val_n = [D.apply_norm(b, norm) for b in val]
    test_n = [D.apply_norm(b, norm) for b in test]
    pairs = D.stack_pairs(train_n)
    xv, tv, yv, cv = D.stack(val_n)
    print(f"batteries train/val/test: {len(train)}/{len(val)}/{len(test)}, "
          f"train pairs: {len(pairs[0])}")

    mpath = os.path.join(args.out, "metrics.json")
    results = json.load(open(mpath)) if os.path.exists(mpath) else {}
    for tag in args.models.split(","):
        t0 = time.time()
        if tag == "baseline":
            params, hist = train_baseline(pairs, xv, tv, yv, args.steps, args.seed)
            extra = {"steps": args.steps}
        else:
            base = M.load_params(os.path.join(args.out, "params_baseline.pkl"))
            rates = D.rate_dataset(train_n)
            params, hist = train_pinn(base, pairs, rates, xv, tv, yv, args.steps_f,
                                      args.steps_ft, args.alpha, args.beta, args.seed)
            extra = {"steps_f": args.steps_f, "steps_ft": args.steps_ft,
                     "alpha": args.alpha, "beta": args.beta}
        pd.DataFrame(hist).to_csv(os.path.join(args.out, f"history_{tag}.csv"), index=False)
        M.save_params(params, os.path.join(args.out, f"params_{tag}.pkl"))
        extra["train_seconds"] = round(time.time() - t0, 1)
        evaluate(tag, params, test_n, args.out, results, extra)
        with open(mpath, "w") as f:
            json.dump(results, f, indent=2)


if __name__ == "__main__":
    main()
