# %%

import copy
import timeit

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

from scimba_jax.domains.meshless_domains.base import VolumetricDomain
from scimba_jax.domains.meshless_domains.domains_2d import Disk2D
from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import (
    ApproximationSpace,
)
from scimba_jax.nonlinear_approximation.approximation_spaces.densityflow_approximation_spaces import (
    DensityFlowApproximationSpace,
    DensityFlowInvertibleApproximationSpace,
)
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_vectorial import (
    ParamScalarFunction,
    ParamVecFunction,
)
from scimba_jax.nonlinear_approximation.networks.mlp import MLP
from scimba_jax.nonlinear_approximation.networks.structure_preserving_nets.affine_ode_layers import (  # noqa: E501
    AffineFlowLayer,
)
from scimba_jax.nonlinear_approximation.networks.structure_preserving_nets.coupling_layers import (  # noqa: E501
    CouplingLayer,
)
from scimba_jax.nonlinear_approximation.networks.structure_preserving_nets.invertible_nn import (  # noqa: E501
    InvertibleNet,
)
from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector
from scimba_jax.physical_models.abstract_physical_model import (
    PHYSICAL_RESIDUALS_TYPE,
    AbstractPhysicalModel,
)
from scimba_jax.physical_models.abstract_residuals import (
    DOMAIN_TYPE,
    NDARRAYS_FUNC_TYPE,
    PARAM_FUNC_TYPE,
    InitialResidual,
    InteriorResidual,
)


class RotatingTransportResidualInvertibleLossFlow(InteriorResidual):
    def __init__(
        self,
        domain: VolumetricDomain,
        time_domain: tuple[float, float],
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
        Tf: float = 1.0,
    ):
        super().__init__(
            domain=domain,
            size=2,
            model_type="x_t",
            f_rhs=f_rhs,
            time_domain=time_domain,
        )
        self.Tf = Tf

    def velocity(self, x):
        v1 = -2.0 * jnp.pi * x[1]
        v2 = 2.0 * jnp.pi * x[0]
        return jnp.array([v1, v2])

    def construct_residual(
        self, phi: PARAM_FUNC_TYPE, inv_phi: PARAM_FUNC_TYPE, rho: PARAM_FUNC_TYPE
    ) -> PARAM_FUNC_TYPE:
        phi1, phi2 = inv_phi.components()
        dx_phi1 = phi1.gradient("x")
        dx_phi2 = phi2.gradient("x")
        dt_phi1 = phi1.d_t()
        dt_phi2 = phi2.d_t()
        return ParamVecFunction.cat(
            [dt_phi1 + dx_phi1.dot(self.velocity), dt_phi2 + dx_phi2.dot(self.velocity)]
        )


class RotatingTransportResidualMLPLossFlow(InteriorResidual):
    def __init__(
        self,
        domain: VolumetricDomain,
        time_domain: tuple[float, float],
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
        Tf: float = 1.0,
    ):
        super().__init__(
            domain=domain,
            size=2,
            model_type="x_t",
            f_rhs=f_rhs,
            time_domain=time_domain,
        )
        self.Tf = Tf

    def velocity(self, x):
        v1 = -2.0 * jnp.pi * x[1]
        v2 = 2.0 * jnp.pi * x[0]
        return jnp.array([v1, v2])

    def construct_residual(
        self, inv_phi: PARAM_FUNC_TYPE, rho: PARAM_FUNC_TYPE
    ) -> PARAM_FUNC_TYPE:
        phi1, phi2 = inv_phi.components()
        dx_phi1 = phi1.gradient("x")
        dx_phi2 = phi2.gradient("x")
        dt_phi1 = phi1.d_t()
        dt_phi2 = phi2.d_t()
        return ParamVecFunction.cat(
            [dt_phi1 + dx_phi1.dot(self.velocity), dt_phi2 + dx_phi2.dot(self.velocity)]
        )


class RotatingTransportResidualInvertibleLossDensity(InteriorResidual):
    def __init__(
        self,
        domain: VolumetricDomain,
        time_domain: tuple[float, float],
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
        Tf: float = 1.0,
    ):
        super().__init__(
            domain=domain,
            size=1,
            model_type="x_t",
            f_rhs=f_rhs,
            time_domain=time_domain,
        )
        self.Tf = Tf

    def velocity(self, x):
        v1 = -2.0 * jnp.pi * x[1]
        v2 = 2.0 * jnp.pi * x[0]
        return jnp.array([v1, v2])

    def construct_residual(
        self, phi: PARAM_FUNC_TYPE, inv_phi: PARAM_FUNC_TYPE, rho: PARAM_FUNC_TYPE
    ) -> PARAM_FUNC_TYPE:
        dx_rho = rho.gradient("x")
        dt_rho = rho.d_t()
        return ParamVecFunction.cat([dt_rho + dx_rho.dot(self.velocity)])


class RotatingTransportResidualMLPLossDensity(InteriorResidual):
    def __init__(
        self,
        domain: VolumetricDomain,
        time_domain: tuple[float, float],
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
        Tf: float = 1.0,
    ):
        super().__init__(
            domain=domain,
            size=1,
            model_type="x_t",
            f_rhs=f_rhs,
            time_domain=time_domain,
        )
        self.Tf = Tf

    def velocity(self, x):
        v1 = -2.0 * jnp.pi * x[1]
        v2 = 2.0 * jnp.pi * x[0]
        return jnp.array([v1, v2])

    def construct_residual(
        self, inv_phi: PARAM_FUNC_TYPE, rho: PARAM_FUNC_TYPE
    ) -> PARAM_FUNC_TYPE:
        dx_rho = rho.gradient("x")
        dt_rho = rho.d_t()
        return ParamVecFunction.cat([dt_rho + dx_rho.dot(self.velocity)])


class ProjectionResidualInvertibleLossFlow(InitialResidual):
    def __init__(
        self,
        domain: DOMAIN_TYPE,
        time_domain: float | tuple[float] = (0.0,),
        size: int = 2,
        model_type: str = "x_t",
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
    ):
        super().__init__(
            domain=domain,
            time_domain=(
                (time_domain,) if isinstance(time_domain, float) else time_domain
            ),
            size=size,
            model_type=model_type,
            f_rhs=f_rhs,
        )

    def construct_residual(
        self, phi: PARAM_FUNC_TYPE, inv_phi: PARAM_FUNC_TYPE, rho: PARAM_FUNC_TYPE
    ) -> PARAM_FUNC_TYPE:
        assert isinstance(rho, ParamScalarFunction) or isinstance(rho, ParamVecFunction)
        return inv_phi.set_t_0(self.time_domain[0])


class ProjectionResidualMLPLossFlow(InitialResidual):
    def __init__(
        self,
        domain: DOMAIN_TYPE,
        time_domain: float | tuple[float] = (0.0,),
        size: int = 2,
        model_type: str = "x_t",
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
    ):
        super().__init__(
            domain=domain,
            time_domain=(
                (time_domain,) if isinstance(time_domain, float) else time_domain
            ),
            size=size,
            model_type=model_type,
            f_rhs=f_rhs,
        )

    def construct_residual(
        self, inv_phi: PARAM_FUNC_TYPE, rho: PARAM_FUNC_TYPE
    ) -> PARAM_FUNC_TYPE:
        assert isinstance(rho, ParamScalarFunction) or isinstance(rho, ParamVecFunction)
        return inv_phi.set_t_0(self.time_domain[0])


class ProjectionResidualInvertibleLossDensity(InitialResidual):
    def __init__(
        self,
        domain: DOMAIN_TYPE,
        time_domain: float | tuple[float] = (0.0,),
        size: int = 1,
        model_type: str = "x_t",
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
    ):
        super().__init__(
            domain=domain,
            time_domain=(
                (time_domain,) if isinstance(time_domain, float) else time_domain
            ),
            size=size,
            model_type=model_type,
            f_rhs=f_rhs,
        )

    def construct_residual(
        self, phi: PARAM_FUNC_TYPE, inv_phi: PARAM_FUNC_TYPE, rho: PARAM_FUNC_TYPE
    ) -> PARAM_FUNC_TYPE:
        assert isinstance(rho, ParamScalarFunction) or isinstance(rho, ParamVecFunction)
        return rho.set_t_0(self.time_domain[0])


class ProjectionResidualMLPLossDensity(InitialResidual):
    def __init__(
        self,
        domain: DOMAIN_TYPE,
        time_domain: float | tuple[float] = (0.0,),
        size: int = 1,
        model_type: str = "x_t",
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
    ):
        super().__init__(
            domain=domain,
            time_domain=(
                (time_domain,) if isinstance(time_domain, float) else time_domain
            ),
            size=size,
            model_type=model_type,
            f_rhs=f_rhs,
        )

    def construct_residual(
        self, inv_phi: PARAM_FUNC_TYPE, rho: PARAM_FUNC_TYPE
    ) -> PARAM_FUNC_TYPE:
        assert isinstance(rho, ParamScalarFunction) or isinstance(rho, ParamVecFunction)
        return rho.set_t_0(self.time_domain[0])


class RotatingTransportResidualClassic(InteriorResidual):
    """PDE residual for a classic (non-flow) PINN: dt(rho) + v . grad_x(rho) = 0."""

    def __init__(
        self,
        domain: VolumetricDomain,
        time_domain: tuple[float, float],
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
    ):
        super().__init__(
            domain=domain,
            size=1,
            model_type="x_t",
            f_rhs=f_rhs,
            time_domain=time_domain,
        )

    def velocity(self, x):
        v1 = -2.0 * jnp.pi * x[1]
        v2 = 2.0 * jnp.pi * x[0]
        return jnp.array([v1, v2])

    def construct_residual(self, rho: PARAM_FUNC_TYPE) -> PARAM_FUNC_TYPE:
        dx_rho = rho.gradient("x")
        dt_rho = rho.d_t()
        return ParamVecFunction.cat([dt_rho + dx_rho.dot(self.velocity)])


class ProjectionResidualClassic(InitialResidual):
    def __init__(
        self,
        domain: DOMAIN_TYPE,
        time_domain: float | tuple[float] = (0.0,),
        size: int = 1,
        model_type: str = "x_t",
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
    ):
        super().__init__(
            domain=domain,
            time_domain=(
                (time_domain,) if isinstance(time_domain, float) else time_domain
            ),
            size=size,
            model_type=model_type,
            f_rhs=f_rhs,
        )

    def construct_residual(self, rho: PARAM_FUNC_TYPE) -> PARAM_FUNC_TYPE:
        assert isinstance(rho, ParamScalarFunction)
        return rho.set_t_0(self.time_domain[0])


class RotatingTransportClassicModel(AbstractPhysicalModel):
    """Classic PINN model: a single scalar field rho(x, t) advected by the
    rotation velocity field, trained with the PDE residual dt(rho) + v.grad_x(rho) = 0
    and an initial condition rho(x, 0) = f_ic_rhs(x)."""

    def __init__(
        self,
        main_domain: VolumetricDomain,
        time_domain: tuple[float, float],
        f_ic_rhs: NDARRAYS_FUNC_TYPE | None = None,
    ):
        super().__init__(main_domain=main_domain, time_domain=time_domain)
        self.physical_residuals: PHYSICAL_RESIDUALS_TYPE = {
            self.main_domain.get_label(): RotatingTransportResidualClassic(
                domain=main_domain, time_domain=time_domain
            ),
            "ic " + self.main_domain.get_label(): ProjectionResidualClassic(
                domain=main_domain,
                time_domain=time_domain,
                size=1,
                model_type="x_t",
                f_rhs=f_ic_rhs,
            ),
        }

    def renew_time_domain(self, new_time_domain: tuple[float, float]):
        """Renew the time domain of the physical model and of all its residuals."""
        self.time_domain = new_time_domain
        for label in self.physical_residuals:
            self.physical_residuals[label].time_domain = new_time_domain


class RotatingTransportTimeInterval(AbstractPhysicalModel):
    def __init__(
        self,
        main_domain: VolumetricDomain,
        time_domain: tuple[float, float],
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
        bc: str = "strong",
        ic: str = "weak",
        f_ic_rhs: NDARRAYS_FUNC_TYPE | None = None,
        type_loss: str = "flow",
        type_net: str = "MLP",
        Tf=8.0,
    ):
        super().__init__(main_domain=main_domain, time_domain=time_domain)
        if type_loss == "flow":
            if type_net == "MLP":
                residual = RotatingTransportResidualMLPLossFlow
                ic_residual = ProjectionResidualMLPLossFlow
            else:
                residual = RotatingTransportResidualInvertibleLossFlow
                ic_residual = ProjectionResidualInvertibleLossFlow
            size = 2
        else:
            if type_net == "MLP":
                residual = RotatingTransportResidualMLPLossDensity
                ic_residual = ProjectionResidualMLPLossDensity
            else:
                residual = RotatingTransportResidualInvertibleLossDensity
                ic_residual = ProjectionResidualInvertibleLossDensity
            size = 1

        self.physical_residuals: PHYSICAL_RESIDUALS_TYPE = {
            self.main_domain.get_label(): residual(
                domain=main_domain,
                time_domain=time_domain,
                f_rhs=f_rhs,
                Tf=Tf,
            ),
        }

        if ic == "weak":
            label = "ic " + self.main_domain.get_label()
            self.physical_residuals[label] = ic_residual(
                domain=main_domain,
                time_domain=time_domain,
                size=size,
                model_type="x_t",
                f_rhs=f_ic_rhs,
            )

    def renew_time_domain(self, new_time_domain: tuple[float, float]):
        """Renew the time domain of the physical model and of all its residuals."""
        self.time_domain = new_time_domain
        for label in self.physical_residuals:
            self.physical_residuals[label].time_domain = new_time_domain


def create_net(type_net, size, conditional_size, key, hidden_size=14, num_layers=8):
    """Crée un InvertibleNet avec des CouplingLayers."""
    if type_net == "invertible":
        layers_list = [
            CouplingLayer(
                size=size,
                conditional_size=conditional_size,
                num_splits=2,
                ode_layer_type=AffineFlowLayer,
                hidden_sizes=[hidden_size],
                activation="tanh",
                key=key,
            )
            for _ in range(num_layers)
        ]
        model = InvertibleNet(
            size=size,
            conditional_size=conditional_size,
            layers_list=layers_list,
        )
    else:
        model = MLP(
            in_size=dim + time_and_params_dim,
            hidden_sizes=[hidden_size] * num_layers,
            out_size=dim,
            activation="tanh",
        )

    return model


def rotate_point(x, theta):
    """Rotate a 2D point x by angle theta around the origin."""
    cos_t, sin_t = jnp.cos(theta), jnp.sin(theta)
    return jnp.array([cos_t * x[0] - sin_t * x[1], sin_t * x[0] + cos_t * x[1]])


def make_exact_density(init_density_fn):
    """Build the exact density at time t from an initial density.

    Under the rigid rotation v = (-2*pi*x2, 2*pi*x1), any initial shape is
    transported without distortion: x(t) = R(2*pi*t) x(0), so
    rho(t, x) = rho_0(R(-2*pi*t) x).
    """

    def exact_density(t, x):
        theta = -2.0 * jnp.pi * t[0]
        return init_density_fn(rotate_point(x, theta))

    return exact_density


def initial_solution_flow(x):
    return x


def initial_density_gaussian(x):
    nu = 0.08
    c = 0.3
    r2 = (x[0] - c) ** 2 + x[1] ** 2
    return jnp.exp(-r2 / (2 * nu**2))[None]


def initial_density_cross(x):
    cx, cy = 0.3, 0.0
    half_len, half_w = 0.18, 0.05
    dxc, dyc = x[0] - cx, x[1] - cy
    in_h = (jnp.abs(dxc) < half_len) & (jnp.abs(dyc) < half_w)
    in_v = (jnp.abs(dxc) < half_w) & (jnp.abs(dyc) < half_len)
    return jnp.where(in_h | in_v, 1.0, 0.0)[None]


def initial_density_cone_slot(x):
    # cone, centered above the rotation arc
    cx_cone, cy_cone, r_cone = 0.3, 0.18, 0.13
    dist_cone = jnp.sqrt((x[0] - cx_cone) ** 2 + (x[1] - cy_cone) ** 2)
    cone = jnp.clip(1.0 - dist_cone / r_cone, 0.0, 1.0)

    # slotted cylinder ("crochet"), centered below the rotation arc
    cx_slot, cy_slot, r_slot = 0.3, -0.18, 0.13
    dist_slot = jnp.sqrt((x[0] - cx_slot) ** 2 + (x[1] - cy_slot) ** 2)
    in_disk = dist_slot < r_slot
    in_notch = (jnp.abs(x[0] - cx_slot) < 0.035) & (x[1] - cy_slot > -0.02)
    slot = jnp.where(in_disk & ~in_notch, 1.0, 0.0)

    return jnp.maximum(cone, slot)[None]


def plot_losses_and_density_single_pinn(
    domain_x, pinn, flow_number, type_network=None, n_visu=256
):
    for label in pinn.losses.losses_history:
        if (label == "total") or (pinn.losses.losses_history[label].shape[1] == 1):
            plt.semilogy(pinn.losses.losses_history[label], label=label + " loss")
        else:
            for i in range(pinn.losses.losses_history[label].shape[1]):
                plt.semilogy(
                    pinn.losses.losses_history[label][:, i],
                    label=label + "%d loss" % i,
                )

    plt.title(f"loss history for flow {flow_number}")
    plt.legend()
    plt.show(block=True)

    if type_network == "MLP":
        id_density = 2
    else:
        id_density = 4

    sh = (n_visu, n_visu)

    x1_lin = jnp.linspace(-1, 1, n_visu)
    x2_lin = jnp.linspace(-1, 1, n_visu)
    x1, x2 = jnp.meshgrid(x1_lin, x2_lin)
    x = jnp.stack([x1.ravel(), x2.ravel()], axis=-1)

    r2 = x1**2 + x2**2
    mask = r2 > 1.0

    def plot_rho(a, title, t_val):
        t = jnp.ones_like(x[:, 0:1]) * t_val
        print(
            f"in plot: pinn.space.idx_current_flow={pinn.space.idx_current_flow}, t_val={t_val}"
        )
        rho = pinn.evaluate(x, t)[:, id_density : id_density + 1].reshape(*sh)
        rho = jnp.where(mask, jnp.nan, rho)

        im = a.contourf(x1, x2, rho, levels=64, cmap="turbo")
        plt.colorbar(im, ax=a)
        a.set_aspect("equal")
        a.contour(x1, x2, rho, levels=10, colors="white", linewidths=0.5)
        a.set_aspect("equal")
        min_rho = float(rho.min())
        max_rho = float(rho.max())
        title = f"flow {flow_number}: {title} (min={min_rho:.2f}, max={max_rho:.2f})"
        a.set_title(title)

    fig, ax = plt.subplots(3, 2, figsize=(12, 18))
    plot_rho(ax[0, 0], "t=0.0", 0.0)
    plot_rho(ax[0, 1], "t=0.25", 0.25)
    plot_rho(ax[1, 0], "t=0.45", 0.45)
    plot_rho(ax[1, 1], "t=0.49", 0.49)
    plot_rho(ax[2, 0], "t=0.75", 0.75)
    plot_rho(ax[2, 1], "t=1.0", 1.0)

    plt.tight_layout()
    plt.show(block=True)


def plot_losses_and_density(domain_x, list_of_pinns, type_network=None, n_visu=256):
    for i, pinn in enumerate(list_of_pinns):
        pinn.space.inference_mode = True
        old_idx = pinn.space.idx_current_flow
        pinn.space.idx_current_flow = i
        plot_losses_and_density_single_pinn(
            domain_x, pinn, flow_number=i, type_network=type_network, n_visu=n_visu
        )
        pinn.space.idx_current_flow = old_idx
        pinn.space.inference_mode = False


def plot_single_flow(domain_x, pinn, flow_number, type_network=None, n_visu=256):
    sh = (n_visu, n_visu, 2)

    x1_lin = jnp.linspace(-1, 1, n_visu)
    x2_lin = jnp.linspace(-1, 1, n_visu)
    x1, x2 = jnp.meshgrid(x1_lin, x2_lin)
    x = jnp.stack([x1.ravel(), x2.ravel()], axis=-1)

    r2 = x1**2 + x2**2
    mask = (r2 > 1.0)[:, :, None]

    if type_network == "MLP":
        id_bwd = 0
    else:
        id_bwd = 2

    def plot_flow(ax, title, t_val):
        t = jnp.ones_like(x[:, 0:1]) * t_val
        flow_bwd = pinn.evaluate(x, t)[:, id_bwd : id_bwd + 2].reshape(*sh)
        flow_bwd = jnp.where(mask, jnp.nan, flow_bwd)

        norm_flow_bwd = jnp.linalg.norm(flow_bwd, axis=-1)

        im = ax[0].contourf(x1, x2, norm_flow_bwd, levels=64, cmap="turbo")
        plt.colorbar(im, ax=ax[0])
        ax[0].set_aspect("equal")
        ax[0].set_title(f"flow {flow_number}, {title}: (norm)")

        im = ax[1].contourf(x1, x2, flow_bwd[..., 0], levels=64, cmap="turbo")
        plt.colorbar(im, ax=ax[1])
        ax[1].set_aspect("equal")
        ax[1].set_title(f"flow {flow_number}, {title}: (x)")

        im = ax[2].contourf(x1, x2, flow_bwd[..., 1], levels=64, cmap="turbo")
        plt.colorbar(im, ax=ax[2])
        ax[2].set_aspect("equal")
        ax[2].set_title(f"flow {flow_number}, {title}: (y)")

    fig, ax = plt.subplots(5, 3, figsize=(9, 6))
    plot_flow(ax[0], "t=0.0", 0.0)
    plot_flow(ax[1], "t=0.25", 0.25)
    plot_flow(ax[2], "t=0.5", 0.5)
    plot_flow(ax[3], "t=0.75", 0.75)
    plot_flow(ax[4], "t=1.0", 1.0)

    plt.tight_layout()
    plt.show(block=True)


def plot_flows(domain_x, pinn, type_network=None, n_visu=256):
    for i, pinn in enumerate(pinn):
        old_idx = pinn.space.idx_current_flow
        pinn.space.idx_current_flow = i
        plot_single_flow(
            domain_x, pinn, flow_number=i, type_network=type_network, n_visu=n_visu
        )
        pinn.space.idx_current_flow = old_idx


def plot_loss_history(pinn, title):
    for label in pinn.losses.losses_history:
        if (label == "total") or (pinn.losses.losses_history[label].shape[1] == 1):
            plt.semilogy(pinn.losses.losses_history[label], label=label + " loss")
        else:
            for i in range(pinn.losses.losses_history[label].shape[1]):
                plt.semilogy(
                    pinn.losses.losses_history[label][:, i],
                    label=label + "%d loss" % i,
                )

    plt.title(title)
    plt.legend()
    plt.show(block=True)


def plot_comparison(case_name, classic_pinn, flow_pinn, init_density_fn, n_visu=128):
    """Compare, at 3 time instants, the classic PINN and the flow PINN
    densities, together with their pointwise errors against the exact
    solution. All panels share the same colormap and the same contourf/contour
    style as the loss/density plots (turbo filled contours with white
    contour lines), each with its own color scale."""
    times = [0.0, 0.5, 1.0]

    flow_pinn.space.initial_density = init_density_fn
    flow_pinn.space.inference_mode = True
    flow_pinn.space.idx_current_flow = 0

    x1_lin = jnp.linspace(-1, 1, n_visu)
    x2_lin = jnp.linspace(-1, 1, n_visu)
    x1, x2 = jnp.meshgrid(x1_lin, x2_lin)
    x = jnp.stack([x1.ravel(), x2.ravel()], axis=-1)
    sh = (n_visu, n_visu)

    r2 = x1**2 + x2**2
    mask = r2 > 1.0

    exact_density = jax.vmap(make_exact_density(init_density_fn))
    cmap = "turbo"

    fig, axes = plt.subplots(len(times), 4, figsize=(18, 4.2 * len(times)))

    for row, t_val in enumerate(times):
        t = jnp.ones_like(x[:, 0:1]) * t_val

        rho_classic = classic_pinn.evaluate(x, t)[:, 0:1].reshape(*sh)
        rho_flow = flow_pinn.evaluate(x, t)[:, 2:3].reshape(*sh)
        rho_exact = exact_density(t, x).reshape(*sh)

        rho_classic = jnp.where(mask, jnp.nan, rho_classic)
        rho_flow = jnp.where(mask, jnp.nan, rho_flow)
        rho_exact = jnp.where(mask, jnp.nan, rho_exact)
        err_classic = jnp.abs(rho_classic - rho_exact)
        err_flow = jnp.abs(rho_flow - rho_exact)

        columns = [
            (rho_classic, "classic PINN"),
            (rho_flow, "flow PINN"),
            (err_classic, "|classic - exact|"),
            (err_flow, "|flow - exact|"),
        ]

        for col, (field, name) in enumerate(columns):
            ax = axes[row, col]
            im = ax.contourf(x1, x2, field, levels=64, cmap=cmap)
            plt.colorbar(im, ax=ax)
            ax.contour(x1, x2, field, levels=10, colors="white", linewidths=0.5)
            ax.set_aspect("equal")
            ax.set_title(f"{name}, t={t_val:.2f}")

    fig.suptitle(f"Rotating transport - {case_name}")
    plt.tight_layout()
    plt.show(block=True)

    flow_pinn.space.inference_mode = False


# %%
N_COLLOC = 8_000
N_IC_COLLOC = 4_000
N_EPOCHS = 100

domain_x = [(-1.0, 1.0), (-1.0, 1.0)]
dx = Disk2D([0.0, 0.0], 1.0, is_main_domain=True)

key = jax.random.PRNGKey(0)

dim = 2
params_dim = 0
nb_models = 2

keys = jax.random.split(key, nb_models)
time_and_params_dim = 1 + params_dim

############# important
type_loss = "flow"
type_net = "MLP"
models = [
    create_net(
        type_net=type_net,
        size=dim,
        conditional_size=time_and_params_dim,
        key=keys[0],
        hidden_size=14,
        num_layers=5,
    )
    for _ in range(nb_models)
]

Tf = 1.0
Deltat = Tf / nb_models
time_intervals = [(Deltat * i, Deltat * (i + 1)) for i in range(nb_models)]

if type_net == "invertible":
    space = DensityFlowInvertibleApproximationSpace(
        dim=dim,
        params_dim=params_dim,
        models=models,
        model_type="x_t",
        initial_density=initial_density_gaussian,
        time_intervals=time_intervals,
    )
else:
    space = DensityFlowApproximationSpace(
        dim=dim,
        params_dim=params_dim,
        models=models,
        model_type="x_t",
        initial_density=initial_density_gaussian,
        time_intervals=time_intervals,
    )
if type_loss == "flow":
    weights = {"interior": [0.3, 0.3], "ic interior": [5.0, 5.0]}
else:
    weights = {"interior": [0.3], "ic interior": [5.0]}

initial_domain_t = (0.0, time_intervals[0][1])

sampler = TensorizedSampler(
    [
        DomainSampler(dx),
        UniformTimeSampler(initial_domain_t),
    ],
    bc=False,
    ic=True,
    model_type="x_t",
)

model = RotatingTransportTimeInterval(
    main_domain=dx,
    time_domain=initial_domain_t,
    f_ic_rhs=initial_solution_flow,
    Tf=Tf,
    type_loss=type_loss,
    type_net=type_net,
)


list_of_pinns = []

for i in range(nb_models):
    space.idx_current_flow = i
    domain_t = time_intervals[i]
    print(
        f"Training flow {i + 1}/{nb_models} on time interval {domain_t} with {type_loss} loss and {type_net} network"
    )
    sampler.renew_sampler(UniformTimeSampler(domain_t), "t")
    model.renew_time_domain(domain_t)

    pinn = Projector(
        model,
        space,
        sampler,
        optimizer="ENG",
        weights=weights,
        one_loss_per_residual=True,
        matrix_regularization=1e-6 * (i + 1),
        gram_assembly="gemm",
    )

    start = timeit.default_timer()
    key, pinn = pinn.project(
        key, space, n_epochs=N_EPOCHS, n_colloc=N_COLLOC, n_ic_colloc=N_IC_COLLOC
    )
    end = timeit.default_timer()

    list_of_pinns.append(copy.deepcopy(pinn))
    space = pinn.space

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


# %%
plot_losses_and_density(domain_x, list_of_pinns, type_network=type_net, n_visu=256)
# %%
# Classic PINN: a single MLP rho(x, t) trained with the PDE residual
# dt(rho) + v.grad_x(rho) = 0 and an IC residual rho(x, 0) = initial_density(x).
#
# The flow PINN above is trained ONCE: its IC residual targets the identity
# map (f_ic_rhs=initial_solution_flow), independent of the density's initial
# condition. It is therefore reused for all test cases below by swapping
# space.initial_density before evaluating the density via the change of
# variables.
#
# Use the LAST entry, not the first: list_of_pinns[i] is a snapshot taken
# right after flow i finishes training, so only models[0..i] are trained --
# models[i+1:] are still at their random initialization (frozen, unmodified
# by later stages, cf. how the training loop keeps `space` across
# iterations). Only list_of_pinns[-1] has every flow segment trained, which
# inference mode needs since it picks the segment from t (not from a fixed
# idx_current_flow) and therefore can select any of them. Using an earlier
# snapshot silently evaluates an untrained network for t past its segment
# (harmless-looking for the smooth Gaussian, but an exact 0 for
# compact-support shapes like "cross"/"cone_slot" since a random flow field
# scatters points outside their tiny support).
flow_pinn = list_of_pinns[-1]

test_cases = [
    ("gaussian", initial_density_gaussian),
    ("cross", initial_density_cross),
    ("cone_slot", initial_density_cone_slot),
]

classic_weights = {"interior": [0.3], "ic interior": [5.0]}

classic_pinns = []

for case_idx, (case_name, init_density_fn) in enumerate(test_cases):
    print(f"\n=== Test case: {case_name} ===")

    classic_key = jax.random.fold_in(key, case_idx)
    nn_classic = MLP(
        in_size=dim + time_and_params_dim,
        hidden_sizes=[20] * 4,
        out_size=1,
        activation="tanh",
        key=classic_key,
    )
    space_classic = ApproximationSpace(
        {"x": dim}, [(nn_classic, "scalar", None)], model_type="x_t"
    )

    sampler_classic = TensorizedSampler(
        [
            DomainSampler(dx),
            UniformTimeSampler(initial_domain_t),
        ],
        bc=False,
        ic=True,
        model_type="x_t",
    )

    model_classic = RotatingTransportClassicModel(
        main_domain=dx,
        time_domain=initial_domain_t,
        f_ic_rhs=init_density_fn,
    )

    pinn_classic = Projector(
        model_classic,
        space_classic,
        sampler_classic,
        optimizer="ENG",
        weights=classic_weights,
        one_loss_per_residual=True,
        matrix_regularization=1e-6,
        gram_assembly="gemm",
    )

    start = timeit.default_timer()
    key, pinn_classic = pinn_classic.project(
        key,
        space_classic,
        n_epochs=N_EPOCHS,
        n_colloc=N_COLLOC,
        n_ic_colloc=N_IC_COLLOC,
    )
    end = timeit.default_timer()

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

    classic_pinns.append(copy.deepcopy(pinn_classic))

# %%
if plot := True:
    plot_losses_and_density(domain_x, list_of_pinns, type_network=type_net, n_visu=256)
    plot_flows(domain_x, list_of_pinns, type_network=type_net, n_visu=256)

    for (case_name, init_density_fn), pinn_classic in zip(test_cases, classic_pinns):
        plot_loss_history(pinn_classic, title=f"classic PINN losses ({case_name})")
        plot_comparison(case_name, pinn_classic, flow_pinn, init_density_fn, n_visu=128)

# %%