r"""Solves a shock and rarefaction problem for a concave flux with a shock-fitting PINN.

.. math::

    \partial_t u + \partial_x f(u) & = 0 \quad \text{in } \Omega \times (T_0, T) \\
    u & = g \quad \text{on } \partial \Omega \times (T_0, T) \\
    u & = u_{T_0} \quad \text{on } \Omega \times \{T_0\}

with :math:`f(u) = 4u(2 - u)`, :math:`\Omega = (-4, 3)`, :math:`T = 0.75` and, at
:math:`t = 0`, :math:`u = 1, 0.5, 1.5` on :math:`x < -2`, :math:`-2 < x < 0`,
:math:`x > 0` (Experiment 1 of Lorin and Novruzi, 2024). A rarefaction leaves
:math:`x = -2` and reaches the shock at :math:`t = 0.5`. The computation starts at
:math:`T_0 = 0.1` from the exact solution, which avoids the singular point
:math:`(0, -2)`.

The solution is approximated with the non-diffusive neural network (NDNN) method:
one network on each side of the shock, one for its position, trained with Adam
then ENG. Boundary and initial conditions are weak.
"""

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.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.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.temporal_pde.shock_fitting import (
    ShockFittingConservationLaw1D,
)

N_COLLOC = 2000
N_BC_COLLOC = 200
N_IC_COLLOC = 500
N_EPOCHS_ADAM = 100
N_EPOCHS_ENG = 500
OUTPUT_SCALE = 0.1  # last layer of N_minus, N_plus scaled: they start close to 0

T0, T = 0.1, 0.75  # T0 = 0.0 for the original discontinuous initial condition
S0 = 0.0  # shock position at T0
X_L, X_R = -4.0, 3.0


def flux(u: jnp.ndarray) -> jnp.ndarray:
    return 4.0 * u * (2.0 - u)


def exact_shock(t: jnp.ndarray) -> jnp.ndarray:
    return jnp.where(t < 0.5, 0.0, 4.0 * jnp.sqrt(2.0 * t) - 4.0 * t - 2.0)


def exact_sol(t: jnp.ndarray, x: jnp.ndarray) -> jnp.ndarray:
    # rarefaction fan f'(u) = (x + 2) / t, then the shock
    fan = jnp.clip(1.0 - (x + 2.0) / (8.0 * jnp.maximum(t, 1e-12)), 0.5, 1.0)
    return jnp.where(x > exact_shock(t), 1.5, fan)


def f_ic(x: jnp.ndarray) -> jnp.ndarray:
    return exact_sol(jnp.full_like(x, T0), x)


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


def scaled_output(net: MLP, scale: float) -> MLP:
    last = jax.tree_util.tree_map(lambda a: scale * a, net.get_layer(-1))
    return MLP.set_layer(last, -1, net)


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,
)

# create the model
model = ShockFittingConservationLaw1D(
    dx, domain_t, ConservationLawPDE(flux), f_bc_rhs=f_bc, f_ic_rhs=f_ic
)

# create the space: N_minus, N_plus, and the shock s(t), with s(T0) = S0 imposed
key_m, key_p, key_s = jax.random.split(jax.random.PRNGKey(0), 3)
nets = [
    scaled_output(
        MLP(in_size=2, out_size=1, hidden_sizes=[20, 20], key=k), OUTPUT_SCALE
    )
    for k in (key_m, key_p)
]
space = ShockFittingApproximationSpace(
    {"x": 1},
    states=[(net, "scalar", None) for net in nets],
    shocks=[MLP(in_size=1, out_size=1, hidden_sizes=[20, 20], key=key_s)],
    model_type="t_x",
    t0=T0,
    s0=S0,
)

label = dx.get_label()
weights = {"boundary": [1.0], "ic " + label: [10.0], label: [1.0, 1.0]}  # PDE, RH

# Adam puts the three networks in the right region, then ENG converges
adam = Projector(
    model, space, sampler, optimizer="Adam", learning_rate=2e-3, weights=weights
)
start = timeit.default_timer()
key, adam = adam.project(
    jax.random.PRNGKey(1), space, N_EPOCHS_ADAM, N_COLLOC, N_BC_COLLOC, N_IC_COLLOC
)
print("Adam best loss: ", adam.best_loss)
print("time for %d Adam epochs: " % N_EPOCHS_ADAM, timeit.default_timer() - start)

pinn = Projector(model, adam.space, sampler, weights=weights)
start = timeit.default_timer()
key, pinn = pinn.project(
    key, adam.space, N_EPOCHS_ENG, N_COLLOC, N_BC_COLLOC, N_IC_COLLOC
)
print("ENG best loss: ", pinn.best_loss)
print("time for %d ENG epochs: " % N_EPOCHS_ENG, timeit.default_timer() - start)

# ---------------------------------------------------------------------------
# Plots
# ---------------------------------------------------------------------------

u, _, _, s = pinn.space.create_variables()
ts = jnp.linspace(T0, T, 131)
xs = jnp.linspace(X_L, X_R, 701)
tt, xx = jnp.meshgrid(ts, xs, indexing="ij")
u_nn = jax.vmap(lambda t, x: u(pinn.space, t[None], x[None]))(
    tt.ravel(), xx.ravel()
).reshape(tt.shape)
u_ex = exact_sol(tt, xx)
s_nn = jax.vmap(lambda t: s(pinn.space, t[None], jnp.zeros(1)))(ts).ravel()
s_ex = exact_shock(ts)

# points on the shock line compare the two sides of the jump: left out of the error
off_shock = jnp.abs(xx - exact_shock(tt)) > 0.02
print("max |s - s_exact|: ", float(jnp.max(jnp.abs(s_nn - s_ex))))
print("MSE off the shock: ", float(jnp.mean((u_nn - u_ex)[off_shock] ** 2)))

extent = [X_L, X_R, T0, T]
fig, axes = plt.subplots(2, 3, figsize=(15, 8.4), constrained_layout=True)
for ax, field, title in [
    (axes[0, 0], u_nn, "NDNN solution u(t, x)"),
    (axes[0, 1], u_ex, "Exact solution"),
]:
    im = ax.imshow(
        field, origin="lower", extent=extent, aspect="auto", vmin=0.5, vmax=1.5
    )
    ax.plot(s_nn, ts, "w--", lw=1.5, label="s(t), NDNN")
    ax.set(title=title, xlabel="x", ylabel="t")
    fig.colorbar(im, ax=ax)
axes[0, 0].legend(loc="upper left")

im = axes[0, 2].imshow(
    jnp.abs(u_nn - u_ex), origin="lower", extent=extent, aspect="auto", cmap="magma"
)
axes[0, 2].set(title="|u - u_exact|", xlabel="x", ylabel="t")
fig.colorbar(im, ax=axes[0, 2])

for t_slice, color in zip([T0, 0.3, 0.5, T], ["C0", "C1", "C2", "C3"]):
    i = int(round((t_slice - T0) / (T - T0) * (len(ts) - 1)))
    axes[1, 0].plot(xs, u_ex[i], color=color, lw=3, alpha=0.35)
    axes[1, 0].plot(xs, u_nn[i], color=color, ls="--", label=f"t = {ts[i]:.2f}")
axes[1, 0].set(title="Slices: exact (thick), NDNN (dashed)", xlabel="x", ylabel="u")
axes[1, 0].legend()

axes[1, 1].plot(ts, s_ex, color="k", lw=3, alpha=0.35, label="exact")
axes[1, 1].plot(ts, s_nn, "C3--", label="NDNN")
axes[1, 1].axvline(0.5, color="grey", lw=0.8, ls=":")
axes[1, 1].set(title="Shock position", xlabel="t", ylabel="s(t)")
axes[1, 1].legend()

loss_adam = adam.losses.losses_history["total"].ravel()
loss_eng = pinn.losses.losses_history["total"].ravel()
axes[1, 2].semilogy(loss_adam, color="C0", label="Adam")
axes[1, 2].semilogy(
    len(loss_adam) + jnp.arange(len(loss_eng)), loss_eng, color="C3", label="ENG"
)
axes[1, 2].set(title="Training loss", xlabel="epoch", ylabel="total loss")
axes[1, 2].legend()

fig.suptitle("Shock and rarefaction, f(u) = 4u(2 - u), shock-fitting PINN (NDNN)")
plt.show()