r"""Solves a Riemann problem for the inviscid Burgers equation with 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, 1)`, :math:`T = 1`, :math:`u_0 = 1` for :math:`x < 0` and
:math:`0` otherwise. The exact solution is a shock moving at speed :math:`1/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 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.classical_weakform.conservation_law_fluxes import (
    burgers_flux,
)
from scimba_jax.physical_models.temporal_pde.shock_fitting import (
    ShockFittingConservationLaw1D,
)

N_COLLOC = 400
N_BC_COLLOC = 100
N_IC_COLLOC = 200
N_EPOCHS = 300

T0, T = 0.0, 1.0
S0 = 0.0  # initial shock position


def exact_shock(t: jnp.ndarray) -> jnp.ndarray:
    return S0 + 0.5 * t


def exact_sol(t: jnp.ndarray, x: jnp.ndarray) -> jnp.ndarray:
    return jnp.where(x < exact_shock(t), 1.0, 0.0)


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)


domain_t = (T0, T)
dx = Segment1D((-1.0, 1.0), 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(burgers_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)
space = ShockFittingApproximationSpace(
    {"x": 1},
    states=[
        (MLP(in_size=2, out_size=1, hidden_sizes=[10], key=key_m), "scalar", None),
        (MLP(in_size=2, out_size=1, hidden_sizes=[10], key=key_p), "scalar", None),
    ],
    shocks=[MLP(in_size=1, out_size=1, hidden_sizes=[5], 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

pinn = Projector(model, space, sampler, weights=weights)

start = timeit.default_timer()
key, pinn = pinn.project(
    jax.random.PRNGKey(1), space, N_EPOCHS, N_COLLOC, N_BC_COLLOC, N_IC_COLLOC
)
end = timeit.default_timer()

print("best loss: ", pinn.best_loss)
print("time for %d epochs: " % N_EPOCHS, end - start)

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

u, _, _, s = pinn.space.create_variables()
ts = jnp.linspace(T0, T, 101)
xs = jnp.linspace(-1.0, 1.0, 401)
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)
s_nn = jax.vmap(lambda t: s(pinn.space, t[None], jnp.zeros(1)))(ts).ravel()
print("max |s - s_exact|: ", float(jnp.max(jnp.abs(s_nn - exact_shock(ts)))))

fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), constrained_layout=True)
im = axes[0].imshow(
    u_nn, origin="lower", extent=[-1.0, 1.0, T0, T], aspect="auto", cmap="viridis"
)
axes[0].plot(s_nn, ts, "w--", lw=1.5, label="s(t)")
axes[0].set(title="NDNN solution u(t, x)", xlabel="x", ylabel="t")
axes[0].legend(loc="upper left")
fig.colorbar(im, ax=axes[0])

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

axes[2].plot(ts, exact_shock(ts), color="k", lw=3, alpha=0.35, label="exact")
axes[2].plot(ts, s_nn, "C3--", label="NDNN")
axes[2].set(title="Shock position", xlabel="t", ylabel="s(t)")
axes[2].legend()

fig.suptitle("Riemann problem for Burgers, shock-fitting PINN (NDNN)")
plt.show()