r"""Solves the dam break problem for 1D shallow water with a shock-fitting PINN.

.. math::

    \partial_t h + \partial_x q & = 0 \\
    \partial_t q + \partial_x \Big(\frac{q^2}{h} + \frac{g h^2}{2}\Big) & = 0

on :math:`\Omega = (-1, 1)`, :math:`g = 9.81`, with water at rest at :math:`t = 0`:
:math:`h = 2` for :math:`x < 0`, :math:`h = 1` for :math:`x > 0`. The exact solution
(Stoker) is a rarefaction moving left and a shock moving right. The computation
starts at :math:`T_0 = 0.02` from the exact solution, which avoids the singular
point :math:`(0, 0)`.

The solution is approximated with the non-diffusive neural network (NDNN) method:
one network for :math:`(h, q)` on each side of the shock, one for its position,
trained with Adam then ENG. Boundary and initial conditions are weak. For
comparison, a classical PINN with about as many parameters and the same training
time is trained after it, with exact Dirichlet conditions.
"""

import math
import timeit

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

from scimba_jax.domains.meshless_domains.domains_1d import Segment1D
from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import (
    ApproximationSpace,
)
from scimba_jax.nonlinear_approximation.approximation_spaces.shock_fitting_approximation_spaces import (
    ShockFittingApproximationSpace,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
    DomainSampler,
    TensorizedSampler,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo_time import (
    UniformTimeSampler,
)
from scimba_jax.nonlinear_approximation.model_class.funcparam import ParamFunction
from scimba_jax.nonlinear_approximation.networks.mlp import MLP
from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector
from scimba_jax.physical_models.abstract_conservative_pde import ConservationLawPDE
from scimba_jax.physical_models.abstract_physical_model import AbstractPhysicalModel
from scimba_jax.physical_models.classical_weakform.shallow_water_weak_form import (
    make_shallow_water_flux_fns,
)
from scimba_jax.physical_models.temporal_pde.conservation_law import ConservationLawND
from scimba_jax.physical_models.temporal_pde.shock_fitting import (
    ShockFittingConservationLaw1D,
)

N_COLLOC = 2000
N_BC_COLLOC = 200
N_IC_COLLOC = 500
N_EPOCHS_ADAM = 100
N_EPOCHS_ENG = 300
# classical PINN: about the same training time as the NDNN
N_EPOCHS_ADAM_PINN = 100
N_EPOCHS_ENG_PINN = 500
PINN_ENG_OPTIONS: dict = {
    "linesearch": "strong-wolfe",
    "adaptive_matrix_regularization": True,
    "matrix_regularization": 1e-1,
}
PINN_STRONG_BC = True  # classical PINN: Dirichlet conditions imposed exactly
PINN_BC_WEIGHT, PINN_IC_WEIGHT = 1.0, 1000.0  # BC weight unused if PINN_STRONG_BC

GRAVITY = 9.81
H_L, H_R = 2.0, 1.0
X_L, X_R = -1.0, 1.0
T0, T = 0.02, 0.2


def star_state() -> tuple[float, float]:
    # (h*, u*) by bisection: left rarefaction u = 2 (c_L - c),
    # right shock u = (h - h_R) sqrt(g (h + h_R) / (2 h h_R))
    c_l = math.sqrt(GRAVITY * H_L)

    def mismatch(h: float) -> float:
        u_rarefaction = 2.0 * (c_l - math.sqrt(GRAVITY * h))
        u_shock = (h - H_R) * math.sqrt(0.5 * GRAVITY * (h + H_R) / (h * H_R))
        return u_rarefaction - u_shock

    low, high = H_R, H_L
    for _ in range(100):
        mid = 0.5 * (low + high)
        low, high = (mid, high) if mismatch(mid) > 0.0 else (low, mid)
    h_star = 0.5 * (low + high)
    return h_star, 2.0 * (c_l - math.sqrt(GRAVITY * h_star))


H_STAR, U_STAR = star_state()
C_L, C_STAR = math.sqrt(GRAVITY * H_L), math.sqrt(GRAVITY * H_STAR)
SHOCK_SPEED = H_STAR * U_STAR / (H_STAR - H_R)


def exact_sol(t: jnp.ndarray, x: jnp.ndarray) -> jnp.ndarray:
    # Stoker solution (h, q), for t > 0; in the fan,
    # u = 2 (c_L + xi) / 3, c = (2 c_L - xi) / 3
    xi = x / t
    c_fan = (2.0 * C_L - xi) / 3.0
    h_fan, u_fan = c_fan**2 / GRAVITY, 2.0 * (C_L + xi) / 3.0
    h = jnp.where(xi < -C_L, H_L, jnp.where(xi < U_STAR - C_STAR, h_fan, H_STAR))
    u = jnp.where(xi < -C_L, 0.0, jnp.where(xi < U_STAR - C_STAR, u_fan, U_STAR))
    h = jnp.where(xi > SHOCK_SPEED, H_R, h)
    u = jnp.where(xi > SHOCK_SPEED, 0.0, u)
    return jnp.stack([h, h * u], axis=-1)


def f_ic(x: jnp.ndarray) -> jnp.ndarray:
    return exact_sol(T0, x[0])


def f_bc(t: jnp.ndarray, x: jnp.ndarray, n: jnp.ndarray) -> jnp.ndarray:
    return exact_sol(t[0], x[0])


def strong_bc_state(y: jnp.ndarray, t: jnp.ndarray, x: jnp.ndarray) -> jnp.ndarray:
    # W_g(x) + (x - x_L)(x_R - x) y, W_g linear between the boundary states
    w_left, w_right = jnp.array([H_L, 0.0]), jnp.array([H_R, 0.0])
    w_g = w_left + (x - X_L) / (X_R - X_L) * (w_right - w_left)
    return w_g + (x - X_L) * (X_R - x) * y


def offset_state(y: jnp.ndarray, t: jnp.ndarray, x: jnp.ndarray) -> jnp.ndarray:
    # the depth starts away from 0, so that q^2 / h stays finite
    return y + jnp.array([1.0, 0.0])


def train(
    model: AbstractPhysicalModel,
    space: ApproximationSpace,
    weights: dict,
    n_adam: int,
    n_eng: int,
    name: str,
    eng_options: dict | None = None,
) -> tuple[Projector, Projector]:
    """Adam then ENG, or load both projectors saved by a previous run.

    Args:
        model: The physical model.
        space: The approximation space.
        weights: The residual weights.
        n_adam: The number of Adam epochs.
        n_eng: The number of ENG epochs; 0: Adam only.
        name: The name of the method, part of the checkpoint name.
        eng_options: Keyword arguments of the ENG optimizer.

    Returns:
        The Adam and the final projectors (the same one if ``n_eng`` is 0).
    """
    eng_options = {} if eng_options is None else eng_options
    postfix = "%s_adam%d_eng%d" % (name, n_adam, n_eng)
    adam = Projector(
        model, space, sampler, optimizer="Adam", learning_rate=2e-3, weights=weights
    )
    n_adam_done, adam = adam.load(__file__, postfix + "_adam_phase")
    if n_eng == 0 and n_adam_done > 0:
        return adam, adam
    pinn = Projector(model, space, sampler, weights=weights, **eng_options)
    n_eng_done, pinn = pinn.load(__file__, postfix + "_eng_phase")
    if n_adam_done > 0 and n_eng_done > 0:
        return adam, pinn

    start = timeit.default_timer()
    key, adam = adam.project(
        jax.random.PRNGKey(1), space, n_adam, N_COLLOC, N_BC_COLLOC, N_IC_COLLOC
    )
    adam.save(__file__, postfix + "_adam_phase")
    pinn = adam
    if n_eng > 0:
        pinn = Projector(model, adam.space, sampler, weights=weights, **eng_options)
        key, pinn = pinn.project(
            key, adam.space, n_eng, N_COLLOC, N_BC_COLLOC, N_IC_COLLOC
        )
        pinn.save(__file__, postfix + "_eng_phase")
    elapsed = timeit.default_timer() - start
    print("%s: %d Adam + %d ENG epochs, %.1f s" % (name, n_adam, n_eng, elapsed))
    return adam, pinn


domain_t = (T0, T)
dx = Segment1D((X_L, X_R), is_main_domain=True)
sampler = TensorizedSampler(
    [UniformTimeSampler(domain_t), DomainSampler(dx)],
    model_type="t_x",
    bc=True,
    ic=True,
)
label = dx.get_label()
shallow_water = ConservationLawPDE(make_shallow_water_flux_fns(GRAVITY)[0])

# NDNN: N_minus, N_plus for (h, q), and the shock s(t), with s(T0) = S T0 imposed
key_m, key_p, key_s = jax.random.split(jax.random.PRNGKey(0), 3)
ndnn_space = ShockFittingApproximationSpace(
    {"x": 1},
    states=[
        (MLP(in_size=2, out_size=2, hidden_sizes=[20, 20], key=key_m), "vec", 2),
        (MLP(in_size=2, out_size=2, hidden_sizes=[20, 20], key=key_p), "vec", 2),
    ],
    shocks=[MLP(in_size=1, out_size=1, hidden_sizes=[20, 20], key=key_s)],
    model_type="t_x",
    t0=T0,
    s0=SHOCK_SPEED * T0,
    post_processing=offset_state,
)
ndnn_model = ShockFittingConservationLaw1D(
    dx, domain_t, shallow_water, n_vars=2, f_bc_rhs=f_bc, f_ic_rhs=f_ic
)
ndnn_weights = {
    "boundary": [1.0, 1.0],
    "ic " + label: [10.0, 10.0],
    label: [1.0] * 4,  # PDE (h, q), RH (h, q)
}

# classical PINN: one network for (h, q)
pinn_space = ApproximationSpace(
    {"x": 1},
    [(MLP(in_size=2, out_size=2, hidden_sizes=[36, 36], key=key_m), "vec", 2)],
    model_type="t_x",
    post_processing=strong_bc_state if PINN_STRONG_BC else offset_state,
)
pinn_model = ConservationLawND(
    dx,
    domain_t,
    shallow_water,
    n_vars=2,
    bc="strong" if PINN_STRONG_BC else "weak",
    f_bc_rhs=f_bc,
    f_ic_rhs=f_ic,
)
pinn_weights = {
    "boundary": [PINN_BC_WEIGHT] * 2,
    "ic " + label: [PINN_IC_WEIGHT] * 2,
    label: [1.0, 1.0],  # PDE (h, q)
}

print("parameters: NDNN %d, classical PINN %d" % (ndnn_space.ndof, pinn_space.ndof))
ndnn_adam, ndnn = train(
    ndnn_model, ndnn_space, ndnn_weights, N_EPOCHS_ADAM, N_EPOCHS_ENG, "ndnn"
)
pinn_name = "pinn_%s_bc%g_ic%g%s" % (
    "strongbc" if PINN_STRONG_BC else "weakbc",
    PINN_BC_WEIGHT,
    PINN_IC_WEIGHT,
    "".join("_%s%s" % (k, v) for k, v in sorted(PINN_ENG_OPTIONS.items())),
)
pinn_adam, pinn = train(
    pinn_model,
    pinn_space,
    pinn_weights,
    N_EPOCHS_ADAM_PINN,
    N_EPOCHS_ENG_PINN,
    pinn_name,
    PINN_ENG_OPTIONS,
)

# ---------------------------------------------------------------------------
# Errors and plots
# ---------------------------------------------------------------------------

ts = jnp.linspace(T0, T, 91)
xs = jnp.linspace(X_L, X_R, 801)
tt, xx = jnp.meshgrid(ts, xs, indexing="ij")
w_ex = exact_sol(tt, xx)


def on_grid(w: ParamFunction, space: ApproximationSpace) -> jnp.ndarray:
    """``(h, q)`` on the grid ``ts x xs``."""
    return jax.vmap(lambda t, x: w(space, t[None], x[None]))(
        tt.ravel(), xx.ravel()
    ).reshape(*tt.shape, 2)


w_ndnn, _, _, s = ndnn.space.create_variables()
w_ndnn = on_grid(w_ndnn, ndnn.space)
w_pinn = on_grid(pinn.space.create_variables()[0], pinn.space)
s_ndnn = jax.vmap(lambda t: s(ndnn.space, t[None], jnp.zeros(1)))(ts).ravel()

# the classical PINN's shock: where h crosses (h_* + h_R) / 2, right of the star
# state; its width at T: from 90% to 10% of the jump
h_mid = 0.5 * (H_STAR + H_R)
right_part = xx > 0.5 * SHOCK_SPEED * tt
s_pinn = xs[jnp.argmax(right_part & (w_pinn[..., 0] < h_mid), axis=1)]
h_at_T = w_pinn[-1, :, 0]
below_90 = right_part[-1] & (h_at_T < H_R + 0.9 * (H_STAR - H_R))
below_10 = right_part[-1] & (h_at_T < H_R + 0.1 * (H_STAR - H_R))

print("h* = %.4f, u* = %.4f, shock speed %.4f" % (H_STAR, U_STAR, SHOCK_SPEED))
# points within 0.02 of the shock compare the two sides of the jump: left out
off_shock = jnp.abs(xx - SHOCK_SPEED * tt) > 0.02
for name, w_nn, s_nn in [("NDNN", w_ndnn, s_ndnn), ("classical PINN", w_pinn, s_pinn)]:
    print("%s: max |s - S t| %.2e" % (name, jnp.max(jnp.abs(s_nn - SHOCK_SPEED * ts))))
    for i, var in enumerate(["h", "q"]):
        err = (w_nn[..., i] - w_ex[..., i])[off_shock]
        print(
            "%s, %s: MSE off the shock %.2e, max error off the shock %.2e"
            % (name, var, jnp.mean(err**2), jnp.max(jnp.abs(err)))
        )
if below_90.any() and below_10.any():
    width = xs[jnp.argmax(below_10)] - xs[jnp.argmax(below_90)]
    print("classical PINN: shock width at T (90%% to 10%% of the jump) %.3f" % width)
else:
    print("classical PINN: at T, h does not fall to 10% of the jump inside the domain")

extent = [X_L, X_R, T0, T]
err_ndnn = jnp.abs(w_ndnn - w_ex).sum(axis=-1)
err_pinn = jnp.abs(w_pinn - w_ex).sum(axis=-1)
err_max = float(jnp.max(jnp.stack([err_ndnn, err_pinn])))
fig, axes = plt.subplots(2, 4, figsize=(19, 8.4), constrained_layout=True)
for ax, field, title, s_line in [
    (axes[0, 0], w_ndnn[..., 0], "NDNN depth h(t, x)", s_ndnn),
    (axes[0, 1], w_pinn[..., 0], "Classical PINN depth h(t, x)", None),
]:
    im = ax.imshow(
        field, origin="lower", extent=extent, aspect="auto", vmin=H_R, vmax=H_L
    )
    if s_line is not None:
        ax.plot(s_line, ts, "w--", lw=1.5, label="s(t), NDNN")
        ax.legend(loc="upper left")
    ax.set(title=title, xlabel="x", ylabel="t")
    fig.colorbar(im, ax=ax)
for ax, err, title in [
    (axes[0, 2], err_ndnn, "NDNN: |h - h_exact| + |q - q_exact|"),
    (axes[0, 3], err_pinn, "Classical PINN: |h - h_exact| + |q - q_exact|"),
]:
    im = ax.imshow(
        err, origin="lower", extent=extent, aspect="auto", cmap="magma", vmax=err_max
    )
    ax.set(title=title, xlabel="x", ylabel="t")
    fig.colorbar(im, ax=ax)

for j, var in enumerate(["h", "q"]):
    for t_slice, color in zip([T0, 0.1, T], ["C0", "C1", "C2"]):
        i = int(round((t_slice - T0) / (T - T0) * (len(ts) - 1)))
        axes[1, j].plot(xs, w_ex[i, :, j], color=color, lw=3, alpha=0.35)
        axes[1, j].plot(
            xs, w_ndnn[i, :, j], color=color, ls="--", label=f"t = {ts[i]:.2f}"
        )
        axes[1, j].plot(xs, w_pinn[i, :, j], color=color, ls=":")
    axes[1, j].set(
        title=f"{var}: exact (thick), NDNN (dashed), PINN (dotted)",
        xlabel="x",
        ylabel=var,
    )
    axes[1, j].legend()

axes[1, 2].plot(ts, SHOCK_SPEED * ts, color="k", lw=3, alpha=0.35, label="exact, S t")
axes[1, 2].plot(ts, s_ndnn, "C3--", label="NDNN, s(t)")
axes[1, 2].plot(ts, s_pinn, "C0:", label="PINN, h = (h_* + h_R) / 2")
axes[1, 2].set(title="Shock position", xlabel="t", ylabel="x")
axes[1, 2].legend()

for adam, eng, color, name in [
    (ndnn_adam, ndnn, "C3", "NDNN"),
    (pinn_adam, pinn, "C0", "classical PINN"),
]:
    histories = [adam] if eng is adam else [adam, eng]  # Adam only: the same one
    loss = jnp.concatenate(
        [p.losses.losses_history["total"].ravel() for p in histories]
    )
    axes[1, 3].semilogy(loss, color=color, label=name)
axes[1, 3].axvline(N_EPOCHS_ADAM, color="grey", lw=0.8, ls=":")
axes[1, 3].set(
    title="Training loss (Adam, then ENG)", xlabel="epoch", ylabel="total loss"
)
axes[1, 3].legend()

fig.suptitle("Dam break, shallow water: shock-fitting PINN (NDNN) and classical PINN")
plt.show()