r"""Solves inviscid Burgers with shock formation and merger, using shock fitting.

.. math::

    \partial_t u + \partial_x \frac{u^2}{2} & = 0 \quad \text{in } \Omega \times (0, T) \\
    u & = g \quad \text{on } \partial \Omega \times (0, T) \\
    u & = u_0 \quad \text{on } \Omega \times \{0\}

with :math:`\Omega = (-1, 2)`, :math:`u_0 = 3/4 - \tanh(2x)` for :math:`x < 1.5` and
:math:`-1/2` otherwise (Experiment 4 of Lorin and Novruzi, 2024). The jump at
:math:`x = 1.5` is a shock from :math:`t = 0`; a second shock forms from the smooth
part at :math:`t = 0.5`, and the two merge at :math:`t \approx 1.32`. Two versions
are trained, up to :math:`T = 0.6` and :math:`T = 2`.

The solution is approximated with the non-diffusive neural network (NDNN) method
with two shock lines: three networks in the cells and two for the lines, trained
with Adam then ENG. An attraction term places the first line before its shock
forms. Boundary and initial conditions are weak. The reference is the exact
entropy solution (Hopf-Lax formula).
"""

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.classical_weakform.conservation_law_fluxes import (
    burgers_flux,
)
from scimba_jax.physical_models.temporal_pde.shock_fitting import (
    ShockFittingConservationLaw1D,
)

# One entry per version: final time, collocation points, epochs, slice times
VERSIONS = [
    dict(T=0.6, n_colloc=3000, n_adam=100, n_eng=300, slices=(0.0, 0.3, 0.5, 0.6)),
    dict(T=2.0, n_colloc=6000, n_adam=100, n_eng=300, slices=(0.0, 0.6, 1.3, 2.0)),
]
N_BC_COLLOC = 200
N_IC_COLLOC = 500

ATTRACTION = True  # False: s_1(0) = 0.5 imposed, as in RR-9605

T0 = 0.0
X_L, X_R = -1.0, 2.0
X_JUMP = 1.5  # jump of u0, initial position of s_2
T_FORM = 0.5  # time at which the second shock forms
S0 = [None if ATTRACTION else 0.5, X_JUMP]  # s_k(T0); None: free


def u0(x: jnp.ndarray) -> jnp.ndarray:
    return jnp.where(x < X_JUMP, 0.75 - jnp.tanh(2.0 * x), -0.5)


def primitive_u0(y: jnp.ndarray) -> jnp.ndarray:
    # primitive of u0, vanishing at 0
    left = 0.75 * y - 0.5 * jnp.log(jnp.cosh(2.0 * y))
    at_jump = 0.75 * X_JUMP - 0.5 * jnp.log(jnp.cosh(2.0 * X_JUMP))
    return jnp.where(y < X_JUMP, left, at_jump - 0.5 * (y - X_JUMP))


# step 1e-3, covers the domain of dependence of Omega up to t = 2
Y_COARSE = jnp.linspace(-6.0, 6.0, 12001)
DY_FINE = jnp.linspace(-1e-3, 1e-3, 201)  # step 1e-5 around the coarse minimum


def exact_sol(t: jnp.ndarray, x: jnp.ndarray) -> jnp.ndarray:
    # Hopf-Lax: u = (x - y*) / t, y* minimising U0(y) + (x - y)^2 / (2t)
    t_pos = jnp.maximum(t, 1e-12)

    def cost(y: jnp.ndarray) -> jnp.ndarray:
        return primitive_u0(y) + (x - y) ** 2 / (2.0 * t_pos)

    y_coarse = Y_COARSE[jnp.argmin(cost(Y_COARSE))]
    y_fine = y_coarse + DY_FINE
    y_star = y_fine[jnp.argmin(cost(y_fine))]
    return jnp.where(t > 0.0, (x - y_star) / t_pos, u0(x))


@jax.jit
def exact_on_grid(ts: jnp.ndarray, xs: jnp.ndarray) -> jnp.ndarray:
    return jax.lax.map(lambda t: jax.vmap(exact_sol, in_axes=(None, 0))(t, xs), ts)


def exact_jumps(ts: jnp.ndarray, xs: jnp.ndarray, tol: float) -> jnp.ndarray:
    # midpoints of the drops of the exact solution larger than tol, NaN elsewhere
    drops = -jnp.diff(exact_on_grid(ts, xs), axis=1)
    return jnp.where(drops > tol, 0.5 * (xs[1:] + xs[:-1]), jnp.nan)


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


def run(T: float, n_colloc: int, n_adam: int, n_eng: int, slices: tuple) -> None:
    """Train (or load) the shock-fitting PINN up to ``T``, print errors, plot.

    Args:
        T: The final time.
        n_colloc: The number of interior collocation points.
        n_adam: The number of Adam epochs.
        n_eng: The number of ENG epochs.
        slices: The times of the slices plotted.
    """
    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: two shock lines, the first one attracted to the steep region
    attracted = (0,) if ATTRACTION else ()
    model = ShockFittingConservationLaw1D(
        dx,
        domain_t,
        ConservationLawPDE(burgers_flux),
        n_shocks=2,
        attracted=attracted,
        f_bc_rhs=f_bc,
        f_ic_rhs=u0,
    )

    # create the space: N_0, N_1, N_2 in the cells, s_1, s_2 for the lines
    keys = jax.random.split(jax.random.PRNGKey(0), 5)
    nets = [MLP(in_size=2, out_size=1, hidden_sizes=[20, 20], key=k) for k in keys[:3]]
    lines = [MLP(in_size=1, out_size=1, hidden_sizes=[20, 20], key=k) for k in keys[3:]]
    space = ShockFittingApproximationSpace(
        {"x": 1},
        states=[(m, "scalar", None) for m in nets],
        shocks=lines,
        model_type="t_x",
        t0=T0,
        s0=S0,
    )

    label = dx.get_label()
    # PDE, RH for s_1 and s_2, attraction of s_1
    weights = {
        "boundary": [1.0],
        "ic " + label: [10.0],
        label: [1.0] * (3 + len(attracted)),
    }

    # Adam then ENG; the trained projectors are saved in ~/.scimba and loaded
    # on the next run with the same settings
    postfix = "%s_T%g_n%d_adam%d_eng%d" % (
        "attraction" if ATTRACTION else "no_attraction",
        T,
        n_colloc,
        n_adam,
        n_eng,
    )
    adam = Projector(
        model, space, sampler, optimizer="Adam", learning_rate=2e-3, weights=weights
    )
    pinn = Projector(model, space, sampler, weights=weights)
    n_adam_done, adam = adam.load(__file__, postfix + "_adam_phase")
    n_eng_done, pinn = pinn.load(__file__, postfix + "_eng_phase")

    if n_adam_done == 0 or n_eng_done == 0:
        start = timeit.default_timer()
        key, adam = adam.project(
            jax.random.PRNGKey(1), space, n_adam, n_colloc, N_BC_COLLOC, N_IC_COLLOC
        )
        print(
            "T = %g, %d Adam epochs: %.1f s"
            % (T, n_adam, timeit.default_timer() - start)
        )
        adam.save(__file__, postfix + "_adam_phase")

        pinn = Projector(model, adam.space, sampler, weights=weights)
        start = timeit.default_timer()
        key, pinn = pinn.project(
            key, adam.space, n_eng, n_colloc, N_BC_COLLOC, N_IC_COLLOC
        )
        print(
            "T = %g, %d ENG epochs: %.1f s" % (T, n_eng, timeit.default_timer() - start)
        )
        pinn.save(__file__, postfix + "_eng_phase")

    print("T = %g, ENG best loss: " % T, pinn.best_loss)

    # errors -----------------------------------------------------------------
    variables = pinn.space.create_variables()
    u, shocks = variables[0], variables[4:]
    ts = jnp.linspace(T0, T, 121)
    xs = jnp.linspace(X_L, X_R, 601)
    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_on_grid(ts, xs)
    s_nn = [
        jax.vmap(lambda t, s=s: s(pinn.space, t[None], jnp.zeros(1)))(ts).ravel()
        for s in shocks
    ]

    # exact shocks: drops of more than 0.05 over 5e-4, i.e. slopes above 100
    jumps = exact_jumps(ts, jnp.linspace(X_L, X_R, 6001), tol=0.05)
    # s_1 is a shock line once its shock has formed; s_2 from the start
    for k, (s, t_start) in enumerate(zip(s_nn, [T_FORM, T0])):
        dist = jnp.nanmin(jnp.abs(jumps - s[:, None]), axis=1)
        dist = jnp.where(ts >= t_start, dist, jnp.nan)
        i = jnp.nanargmax(dist)
        print(
            "T = %g, max distance from s_%d to an exact shock: %.2e (t = %.2f)"
            % (T, k + 1, dist[i], ts[i])
        )
    # points within 0.02 (4 cells) of a shock compare the two sides of the jump:
    # left out. Shock cells of the plotting grid: drops above 0.2 over 0.005.
    shock_cells = jnp.pad(~jnp.isnan(exact_jumps(ts, xs, tol=0.2)), ((0, 0), (4, 4)))
    near = jnp.stack([shock_cells[:, d : d + len(xs)] for d in range(8)]).any(axis=0)
    print("T = %g, MSE off the shocks: %.2e" % (T, jnp.mean((u_nn - u_ex)[~near] ** 2)))

    # plots ------------------------------------------------------------------
    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")
        for k, s in enumerate(s_nn):
            ax.plot(s, ts, "w--", lw=1.5, label="shock lines, NDNN" if k == 0 else None)
        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(slices, ["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()

    ax = axes[1, 1]
    found = ~jnp.isnan(jumps)
    t_jumps = jnp.broadcast_to(ts[:, None], jumps.shape)
    ax.plot(jumps[found], t_jumps[found], ".", color="0.6", ms=3, label="exact shocks")
    front = ts <= T_FORM
    ax.plot(0.75 * ts[front], ts[front], "k:", lw=0.8, label="steepening front, 3t/4")
    for k, s in enumerate(s_nn):
        ax.plot(s, ts, "C3--", label="NDNN" if k == 0 else None)
    ax.set(title="Shock lines", xlabel="x", ylabel="t", xlim=(X_L, X_R))
    ax.legend(loc="upper left")

    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("Burgers up to T = %g, shock-fitting PINN (NDNN)" % T)


for version in VERSIONS:
    run(**version)
plt.show()