"""Physics-informed NN for battery SOH (PINN4SOH-style, JAX).

Two networks:
  u_theta(x, t)      -> SOH estimate (solution network)
  f_phi(x, t, u)     -> dSOH/dt      (learned degradation dynamics)

Physics residual (enforced by loss, and reused at runtime as a health signal):
  R = du/dt - f_phi(x, t, u)  -> 0 on healthy data
plus a monotonicity prior relu(du/dt)^2 (capacity should not grow).
"""
import pickle

import jax
import jax.numpy as jnp


def init_mlp(key, sizes):
    params = []
    for i in range(len(sizes) - 1):
        key, k = jax.random.split(key)
        w = jax.random.normal(k, (sizes[i], sizes[i + 1])) * jnp.sqrt(2.0 / sizes[i])
        params.append({"w": w, "b": jnp.zeros(sizes[i + 1])})
    return params


def mlp(params, x):
    for layer in params[:-1]:
        x = jnp.tanh(x @ layer["w"] + layer["b"])
    return x @ params[-1]["w"] + params[-1]["b"]


def init_params(key, n_features=16):
    ku, kf = jax.random.split(key)
    return {
        "u": init_mlp(ku, [n_features + 1, 64, 64, 32, 1]),
        "f": init_mlp(kf, [3, 32, 32, 1]),  # inputs: (c, t, u)
    }


def u_scalar(u_params, x, t):
    return mlp(u_params, jnp.concatenate([x, jnp.array([t])]))[0]


u_batch = jax.vmap(u_scalar, in_axes=(None, 0, 0))
dudt_batch = jax.vmap(jax.grad(u_scalar, argnums=2), in_axes=(None, 0, 0))


def f_batch(f_params, c, t, u):
    """Degradation law: dSOH/dt = f(C-rate, t, SOH).

    Excludes the cycle features so it can be forward-integrated for
    prognosis without assuming future feature values. Time is included
    because the early-life SOH plateau is not identifiable from state
    alone; the cost is that forecasts lean toward the fleet-average
    trajectory for heterogeneous groups.
    """
    inp = jnp.stack([c, t, u], axis=1)
    # structural monotonicity: degradation rate is non-positive by
    # construction (softplus output negated), not merely loss-penalized
    return -jax.nn.softplus(mlp(f_params, inp)[:, 0])


def loss_fn(params, pairs, alpha=1.0, beta=0.1):
    """Data + discrete physics (state-space) + monotonicity losses.

    pairs = (x1, t1, y1, x2, t2, y2, c): consecutive cycles of one battery.
    Physics: the model's SOH increment between consecutive cycles must equal
    the learned degradation law, in rate units:
        (u(x2,t2) - u(x1,t1)) / (t2-t1) = f(c, t1, u1)
    This captures the TOTAL cycle-to-cycle derivative (feature-driven and
    explicit-time), so f can be forward-integrated for prognosis.
    """
    x1, t1, y1, x2, t2, y2, c = pairs
    u1 = u_batch(params["u"], x1, t1)
    u2 = u_batch(params["u"], x2, t2)
    f = f_batch(params["f"], c, t1, u1)
    dt = t2 - t1
    rate = (u2 - u1) / dt
    sg = jax.lax.stop_gradient
    l_data = 0.5 * jnp.mean((u1 - y1) ** 2) + 0.5 * jnp.mean((u2 - y2) ** 2)
    # asymmetric coupling: f is fitted hard to the model's empirical SOH rates
    # (stop-gradient, weight 1) so it stays integrable; u receives only a
    # light physics pull toward f (weight alpha) as a regularizer.
    l_fit_f = jnp.mean((sg(rate) - f) ** 2)
    l_reg_u = jnp.mean((rate - sg(f)) ** 2)
    l_phys = l_fit_f + alpha * l_reg_u
    l_mono = jnp.mean(jax.nn.relu(f) ** 2)  # degradation rate must be <= 0
    total = l_data + l_phys + beta * l_mono
    return total, {"data": l_data, "phys": l_phys, "mono": l_mono}


def physics_residual(params, c, t, u_meas_dudt, u_meas):
    """Residual between an observed SOH rate and the learned dynamics law."""
    return u_meas_dudt - f_batch(params["f"], c, t, u_meas)


def save_params(params, path):
    with open(path, "wb") as fh:
        pickle.dump(jax.device_get(params), fh)


def load_params(path):
    with open(path, "rb") as fh:
        return pickle.load(fh)
