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

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

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

The initial condition is the level-set function of an ellipsoid with two parameters (the last two input components of :math:`\mu`); the ellipsoid 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 numpy as np
import torch

from scimba_torch import PI
from scimba_torch.approximation_space.nn_space import NNxSpace
from scimba_torch.domain.meshless_domain.domain_3d import Cube3D
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), (0, 1)]
DEFAULT_PARAMETER_BOUNDS = [(0.8, 1.25), (0, 0.25)]
DEFAULT_PARAMETER_BOUNDS_ = torch.tensor(DEFAULT_PARAMETER_BOUNDS)


def f_ini(x, mu):
    # Get coordinates from LabelTensor
    x1, x2, x3 = x.get_components()
    a1, eps = mu.get_components()

    # Define parameters for the ellipsoid
    a2 = 1 / a1
    a3 = 1.2
    # Normalize to ellipsoidal coordinates
    x_ = (x1 - 0.35) / a1
    y_ = (x2 - 0.35) / a2
    z_ = (x3 - 0.35) / a3
    r2 = x_**2 + y_**2 + z_**2

    # Use spherical angles derived from normalized coordinates
    phi = torch.arctan2(torch.sqrt(x_**2 + y_**2), z_)

    # Define perturbation function
    pert = eps * torch.sin(9 * phi) * torch.exp(-1.5 * ((phi - PI / 2) ** 2))

    # Defines a perturbed ellipsoid
    return 2 * (r2 * (1 + pert) - 0.15**2)


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


def plot_implicit(fn, bbox=(0, 1), iter=None):
    """Create a plot of an implicit function.

    Args:
        fn: Implicit function (plot where fn==0).
        bbox: The x, y, and z limits of plotted interval.
        iter: Iteration number (optional).
    """

    def make_xyz(x, y, z):
        if not isinstance(x, np.ndarray):
            x = x * torch.ones_like(torch.tensor(y).flatten().unsqueeze(1))
            y = torch.tensor(y).flatten().unsqueeze(1)
            z = torch.tensor(z).flatten().unsqueeze(1)
        if not isinstance(y, np.ndarray):
            y = y * torch.ones_like(torch.tensor(z).flatten().unsqueeze(1))
            x = torch.tensor(x).flatten().unsqueeze(1)
            z = torch.tensor(z).flatten().unsqueeze(1)
        if not isinstance(z, np.ndarray):
            z = z * torch.ones_like(torch.tensor(x).flatten().unsqueeze(1))
            x = torch.tensor(x).flatten().unsqueeze(1)
            y = torch.tensor(y).flatten().unsqueeze(1)
        x_ = torch.cat((x, y, z), dim=1)
        mu = torch.zeros((x_.shape[0], 0))
        return LabelTensor(x_), LabelTensor(mu)

    xmin, xmax, ymin, ymax, zmin, zmax = bbox * 3
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")
    A = np.linspace(xmin, xmax, 100)  # resolution of the contour
    B = np.linspace(xmin, xmax, 15)  # number of slices
    A1, A2 = np.meshgrid(A, A)  # grid on which the contour is plotted

    for z in B:  # plot contours in the XY plane
        X, Y = A1, A2
        Z = fn(*make_xyz(X, Y, z)).w[:, 0].detach().cpu().reshape((100, 100))
        ax.contour(X, Y, Z + z, [z], zdir="z")
        # [z] defines the only level to plot for this contour for this value of z

    for y in B:  # plot contours in the XZ plane
        X, Z = A1, A2
        Y = fn(*make_xyz(X, y, Z)).w[:, 0].detach().cpu().reshape((100, 100))
        ax.contour(X, Y + y, Z, [y], zdir="y")

    for x in B:  # plot contours in the YZ plane
        Y, Z = A1, A2
        X = fn(*make_xyz(x, Y, Z)).w[:, 0].detach().cpu().reshape((100, 100))
        ax.contour(X + x, Y, Z, [x], zdir="x")

    # must set plot limits because the contour will likely extend
    # way beyond the displayed level.  Otherwise matplotlib extends the plot limits
    # to encompass all values in the contour.
    ax.set_zlim3d(zmin, zmax)
    ax.set_xlim3d(xmin, xmax)
    ax.set_ylim3d(ymin, ymax)

    ax.set_aspect("equal")

    plt.gca().set_rasterization_zorder(-1)

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

    plt.show()


def plot_results(f, sampler, T=0, n_visu=85, iter=None):
    plot_implicit(f)


# %%


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 = (
            torch.rand(10 * n, 2)
            * (DEFAULT_PARAMETER_BOUNDS_[:, 1] - DEFAULT_PARAMETER_BOUNDS_[:, 0])
            + DEFAULT_PARAMETER_BOUNDS_[:, 0]
        )
        parameters = LabelTensor(parameters)
        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.002)

        if pts_1.shape[0] >= n // 4:
            pts_1 = pts_1[: n // 4]

        pts_2 = self.sample_around_zeroes(n, sigma=0.01)
        pts_2 = torch.cat((pts_1, pts_2), dim=0)

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

        pts_3 = self.sample_around_zeroes(n, sigma=0.05)
        pts_3 = torch.cat((pts_2, pts_3), dim=0)

        if pts_3.shape[0] >= 3 * n // 4:
            pts_3 = pts_3[: 3 * n // 4]

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

        pts = torch.cat((pts_3, 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 = Cube3D(DEFAULT_BOUNDS, is_main_domain=True)

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

    space = NNxSpace(
        1, 2, GenericMLP, domain_x, sampler_unif, layer_sizes=[35, 50, 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(DEFAULT_PARAMETER_BOUNDS),
        ]
    )
    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=500, n_collocation=N_c, verbose=True)
        scheme.projector.save(
            "ini_transport_3D_parametric_level_set_larger_truly_adaptive"
        )
        # plot_results(space.evaluate, sampler.sample)

    projector = TimeDiscreteNaturalGradientProjector(pde, rhs=f_ini)
    projector.load("ini_transport_3D_parametric_level_set_larger_truly_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_3D_parametric_level_set_larger_truly_adaptive_large_time_step",
        additional_epochs_for_first_iteration=250,
    )

    # plot_results(space.evaluate, sampler.sample)

    return scheme, sampler


# %%

if __name__ == "__main__":
    scheme, sampler = solve_with_neural_sl(
        3,
        3 / 10,
        N_c=40**3,
        retrain_initial_condition=True,
    )

# # %%

# from skimage import measure


# def f_ini_np_3d(x1, x2, x3):
#     a1 = 1.2
#     a2 = 1 / a1
#     a3 = 1.2
#     # Normalize to ellipsoidal coordinates
#     x_ = (x1 - 0.35) / a1
#     y_ = (x2 - 0.35) / a2
#     z_ = (x3 - 0.35) / a3
#     r2 = x_**2 + y_**2 + z_**2

#     # Use spherical angles derived from normalized coordinates
#     theta = np.arctan2(y_, x_)  # azimuth
#     phi = np.arctan2(np.sqrt(x_**2 + y_**2), z_)  # polar

#     eps = 0.25
#     pert = eps * np.sin(9 * phi) * np.exp(-1.5 * ((phi - np.pi / 2) ** 2))

#     # This defines a perturbed ellipsoid
#     return 2 * (r2 * (1 + pert) - 0.15**2)


# # # %%

# # Create a 3D grid
# n = 100
# x = np.linspace(0.1, 0.6, n)
# y = np.linspace(0.1, 0.6, n)
# z = np.linspace(0.1, 0.6, n)
# X, Y, Z = np.meshgrid(x, y, z, indexing="ij")

# # Evaluate level set function
# phi = f_ini_np_3d(X, Y, Z)

# # # %%

# # # Extract the zero level set using marching cubes
# # verts, faces, _, _ = measure.marching_cubes(
# #     phi, level=0.0, spacing=(x[1] - x[0], y[1] - y[0], z[1] - z[0])
# # )

# # # Plot the zero level set
# # fig = plt.figure(figsize=(8, 6))
# # ax = fig.add_subplot(111, projection="3d")
# # ax.plot_trisurf(
# #     verts[:, 0],
# #     verts[:, 1],
# #     faces,
# #     verts[:, 2],
# #     cmap="viridis",
# #     lw=0.5,
# #     edgecolor="k",
# #     alpha=0.9,
# # )
# # ax.set_title("Zero Level Set of Perturbed 3D Ellipsoid")
# # ax.set_xlabel("x")
# # ax.set_ylabel("y")
# # ax.set_zlabel("z")
# # ax.view_init(elev=30, azim=45)
# # plt.tight_layout()
# # plt.show()

# # # %%

# # Compute the center indices for slicing
# ix = np.abs(x - 0.35).argmin()
# iy = np.abs(y - 0.35).argmin()
# iz = np.abs(z - 0.35).argmin()

# # Plot cross-sectional slices of the level set function
# fig, axs = plt.subplots(1, 3, figsize=(18, 5))

# # XY plane at z = 0.5
# im0 = axs[0].contourf(x, y, phi[:, :, iz].T, levels=150, cmap="RdBu_r")
# axs[0].contour(x, y, phi[:, :, iz].T, levels=[0], colors="k", linewidths=1.5)
# axs[0].set_title("XY slice at z=0.35")
# axs[0].set_xlabel("x")
# axs[0].set_ylabel("y")
# fig.colorbar(im0, ax=axs[0])

# # XZ plane at y = 0.75
# im1 = axs[1].contourf(x, z, phi[:, iy, :].T, levels=150, cmap="RdBu_r")
# axs[1].contour(x, z, phi[:, iy, :].T, levels=[0], colors="k", linewidths=1.5)
# axs[1].set_title("XZ slice at y=0.35")
# axs[1].set_xlabel("x")
# axs[1].set_ylabel("z")
# fig.colorbar(im1, ax=axs[1])

# # YZ plane at x = 0.5
# im2 = axs[2].contourf(y, z, phi[ix, :, :].T, levels=150, cmap="RdBu_r")
# axs[2].contour(y, z, phi[ix, :, :].T, levels=[0], colors="k", linewidths=1.5)
# axs[2].set_title("YZ slice at x=0.35")
# axs[2].set_xlabel("y")
# axs[2].set_ylabel("z")
# fig.colorbar(im2, ax=axs[2])

# plt.tight_layout()
# plt.show()

# # %%