r"""Solves the advection of a 2D level-set function.

..math::
    \partial_t u + a \cdot \nabla u = 0

where :math:`a: \mathbb{R}^2 \times (0, T) \to \mathbb{R}^2` is a time-dependent velocity field and :math:`u: \mathbb{R}^2 \times (0, T) \to \mathbb{R}` is the unknown function.

The initial condition is the level-set function of a disk; the disk is deformed by the velocity field, and goes back to the initial configuration at time :math:`T`.

The equation is solved using the neural semi-Lagrangian scheme with a natural gradient projector.
"""

# %%

import matplotlib.pyplot as plt
import torch

from scimba_torch import PI
from scimba_torch.approximation_space.nn_space import NNxSpace
from scimba_torch.domain.meshless_domain.domain_2d import Square2D
from scimba_torch.integration.monte_carlo import DomainSampler, TensorizedSampler
from scimba_torch.integration.monte_carlo_parameters import UniformParametricSampler
from scimba_torch.neural_nets.coordinates_based_nets.mlp import GenericMLP
from scimba_torch.numerical_solvers.temporal_pde.neural_semilagrangian import (
    Characteristic,
    NeuralSemiLagrangian,
)
from scimba_torch.numerical_solvers.temporal_pde.time_discrete import (
    TimeDiscreteNaturalGradientProjector,
)
from scimba_torch.physical_models.temporal_pde.advection_diffusion_equation import (
    AdvectionReactionDiffusionDirichletStrongForm,
)
from scimba_torch.utils.scimba_tensors import LabelTensor

DEFAULT_BOUNDS = [(0, 1), (0, 1)]


def f_ini(x, mu):
    x1, x2 = x.get_components()
    return (x1 - 0.5) ** 2 + (x2 - 0.75) ** 2 - 0.15**2


def exact(t, x):
    return f_ini(x, None)


def a(t, x, mu):
    x1, x2 = x.get_components()
    pi_x1 = PI * x1
    pi_x2 = PI * x2
    pi_t = PI * t / 8
    v1 = torch.sin(pi_x1) ** 2 * torch.sin(2 * pi_x2) * torch.cos(pi_t)
    v2 = torch.sin(pi_x2) ** 2 * torch.sin(2 * pi_x1) * torch.cos(pi_t)
    return torch.cat((-v1, v2), dim=1)


def plot_results(f, sampler, T=0, n_visu=768, iter=None):
    x, mu = sampler(n_visu**2)
    x1_sampled, x2_sampled = x.get_components()

    x1 = torch.linspace(0, 1 - 0, n_visu)
    x2 = torch.linspace(0, 1 - 0, n_visu)
    x1, x2 = torch.meshgrid(x1, x2, indexing="ij")
    x = LabelTensor(torch.stack((x1.flatten(), x2.flatten()), dim=1))

    x1, x2 = x.get_components()
    x1, x2 = x1.detach().cpu(), x2.detach().cpu()

    predictions = f(x, mu).w.detach().cpu()

    fig, ax = plt.subplots(1, 1, figsize=(6, 6), constrained_layout=True)

    x1 = x1.reshape(n_visu, n_visu)
    x2 = x2.reshape(n_visu, n_visu)
    predictions = predictions.reshape(n_visu, n_visu)

    predictions[predictions < 0] = -1
    predictions[predictions > 0] = 1

    ax.contourf(x1, x2, predictions, levels=2)

    to_plot = torch.randperm(x1_sampled.shape[0])[:25_000]
    x1_sampled = x1_sampled[to_plot].detach().cpu()
    x2_sampled = x2_sampled[to_plot].detach().cpu()
    ax.scatter(x1_sampled, x2_sampled, c="black", s=0.1)

    ax.set_title("Predictions")
    ax.set_aspect("equal")

    plt.gca().set_rasterization_zorder(-1)

    if iter:
        plt.savefig(f"{iter}_transport_2D_level_set_large_adaptive.pdf")

    plt.show()


# %%


class NonUniformSampler(DomainSampler):
    def __init__(self, f, domain, **kwargs):
        super().__init__(domain, **kwargs)
        self.f = f

    def sample_around_zeroes(self, n, sigma=0.05, keep=0.75):
        candidates = LabelTensor(self.vol_sampler.sample(10 * n))
        parameters = LabelTensor(torch.zeros((10 * n, 0)))
        f_values = self.f(candidates, parameters).w[:, 0]
        weights = torch.exp(-(f_values**2) / (2 * sigma**2))
        return candidates[weights > keep].x

    def sample_new_points(self, n: int):
        """Samples `n` new points from the volumetric domain and labels them by subdomains.

        Args:
            n (int): Number of points to sample.

        Returns:
            ScimbaTensor: A tensor of sampled points with their subdomain labels.
        """
        pts_1 = self.sample_around_zeroes(n, sigma=0.01)

        if n - pts_1.shape[0] <= 3 * n // 4:
            pts_1 = pts_1[: n // 4]

        pts_2 = self.sample_around_zeroes(n, sigma=0.05)

        if n - pts_2.shape[0] <= n // 2:
            pts_2 = pts_2[: n // 2]

        pts = torch.cat((pts_1, pts_2), dim=0)

        remaining = n - pts.shape[0]
        pts_unif = self.vol_sampler.sample(remaining)

        pts = torch.cat((pts, pts_unif), dim=0)
        labels = torch.zeros(n, dtype=int)

        # Label points based on subdomains
        for i, subdomain in enumerate(self.domain.list_subdomains):
            condition = subdomain.is_inside(pts)
            labels[condition] = i + 1

        # Wrap the points and labels in a ScimbaTensor
        data = LabelTensor(pts, labels)
        data.x.requires_grad_()
        return data


# %%


def solve_with_neural_sl(
    T: float,
    dt: float,
    retrain_initial_condition=False,
    N_c=10_000,
) -> NeuralSemiLagrangian:
    torch.random.manual_seed(0)

    domain_x = Square2D(DEFAULT_BOUNDS, is_main_domain=True)

    sampler_unif = TensorizedSampler(
        [
            DomainSampler(domain_x),
            UniformParametricSampler([]),
        ]
    )

    space = NNxSpace(1, 0, GenericMLP, domain_x, sampler_unif, layer_sizes=[35, 50, 35])
    pde = AdvectionReactionDiffusionDirichletStrongForm(
        space, a=a, u0=f_ini, constant_advection=False, zero_diffusion=True
    )
    characteristic = Characteristic(pde, periodic=True)

    sampler = TensorizedSampler(
        [
            NonUniformSampler(space.evaluate, domain_x),
            UniformParametricSampler([]),
        ]
    )
    space.integrator = sampler
    pde.space = space
    characteristic = Characteristic(pde, periodic=True)

    if retrain_initial_condition:
        projector = TimeDiscreteNaturalGradientProjector(pde, rhs=f_ini)
        scheme = NeuralSemiLagrangian(characteristic, projector)
        scheme.initialization(epochs=250, n_collocation=N_c, verbose=True)
        scheme.projector.save("ini_transport_2D_level_set_large_adaptive")
        plot_results(space.evaluate, sampler.sample)

    projector = TimeDiscreteNaturalGradientProjector(pde, rhs=f_ini)
    projector.load("ini_transport_2D_level_set_large_adaptive")
    projector.space.load_from_best_approx()
    scheme = NeuralSemiLagrangian(characteristic, projector)
    scheme.solve(
        dt=dt,
        final_time=T,
        epochs=250,
        n_collocation=N_c,
        verbose=True,
        plot=plot_results,
        save="transport_2D_level_set_large_adaptive",
    )

    plot_results(space.evaluate, sampler.sample)

    return scheme, sampler


# %%

if __name__ == "__main__":
    scheme, sampler = solve_with_neural_sl(
        8,
        8 / 40,
        N_c=60_000,
        retrain_initial_condition=True,
    )

# %%