r"""Solves a 1D Allen-Cahn equation with periodic boundary conditions.

.. math::::
    \partial_t u - D \partial_{xx} u + k(u - u^3) = 0

The equation is solved using:
- either the Neural Galerkin method (default),
- or an explicit discrete PINN.
Both use a Natural Gradient optimizer where needed.
"""

# %%

import matplotlib.pyplot as plt
import torch
from scipy.interpolate import RegularGridInterpolator

from scimba_torch.approximation_space.nn_space import NNxSpace
from scimba_torch.domain.meshless_domain.domain_1d import Segment1D
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.features import PeriodicMLP
from scimba_torch.numerical_solvers.temporal_pde.discrete_pinn import DiscretePINN
from scimba_torch.numerical_solvers.temporal_pde.neural_galerkin import NeuralGalerkin
from scimba_torch.numerical_solvers.temporal_pde.time_discrete import (
    TimeDiscreteNaturalGradientProjector,
)
from scimba_torch.physical_models.temporal_pde.heat_equation import (
    HeatEquation1DStrongForm,
)
from scimba_torch.plots.plot_time_discrete_scheme import plot_time_discrete_scheme
from scimba_torch.utils.scimba_tensors import LabelTensor


def f_ini_for_tensors(x: torch.Tensor, mu: torch.Tensor | None = None) -> torch.Tensor:
    return x**4 * torch.cos(torch.pi * x**3)


def f_ini(x: LabelTensor, mu: LabelTensor) -> torch.Tensor:
    x_ = x.get_components()
    mu_ = mu.get_components()
    return f_ini_for_tensors(x_, mu_)


class AllenCahnEquation1DStrongForm(HeatEquation1DStrongForm):
    """Implementation of a 1D Allen-Cahn equation with Neumann BC in strong form.

    Args:
        space: The approximation space for the problem
        init: Callable for the initial condition
        f: Source term function
        **kwargs: Additional keyword arguments
    """

    def space_operator(
        self, w, t: LabelTensor, x: LabelTensor, mu: LabelTensor
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """Apply the spatial operator.

        Args:
            w: Solution tensor
            t: Temporal coordinate tensor
            x: Spatial coordinate tensor
            mu: Parameter tensor

        Returns:
            Spatial operator tensor
        """
        d, k = mu.get_components()
        # laplacian term
        u_x = self.grad(w, x)
        u_xx = self.grad(u_x, x)
        laplacian = -d * u_xx
        # reaction term
        u = w.get_components()
        reaction = -k * (u - u**3)
        # sum the two terms
        return laplacian + reaction

    def functional_operator(
        self,
        func,
        x: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        """Compute the functional operator.

        Args:
            func: Callable representing the function
            t: Temporal coordinate tensor
            x: Spatial coordinate tensor
            mu: Parameter tensor
            theta: Additional parameters tensor

        Returns:
            Functional operator tensor
        """
        # laplacian term
        grad_func = torch.func.jacrev(func, 1)
        laplacian = -mu[0] * torch.func.jacrev(grad_func, 1)(x, mu, theta)
        # reaction term
        u = func(x, mu, theta)
        reaction = -mu[1] * (u - u**3)
        # sum the two terms
        return laplacian[0, 0] + reaction


def solve_allen_cahn(
    d: float,
    k: float,
    t_final: float,
    nx: int = 201,
    dt: float = 1e-2,
    xmin: float = -1.0,
    xmax: float = 1.0,
    snapshot_count: int = 0,
) -> tuple[torch.Tensor, torch.Tensor, dict[float, torch.Tensor]]:
    """
    Solve Allen-Cahn: u_t = d u_xx + k(u - u^3)
    using Crank-Nicolson, Newton, sparse Laplacian, periodic BC.
    """

    # -------------------------
    # Grid
    # -------------------------
    x = torch.linspace(xmin, xmax, nx, dtype=torch.float64)
    dx = (xmax - xmin) / (nx - 1)
    nt = int(round(t_final / dt))

    # -------------------------
    # Sparse periodic Laplacian
    # -------------------------
    rows, cols, vals = [], [], []
    inv_dx2 = 1.0 / dx**2

    for i in range(nx):
        rows.append(i)
        cols.append((i - 1) % nx)
        vals.append(inv_dx2)
        rows.append(i)
        cols.append(i)
        vals.append(-2.0 * inv_dx2)
        rows.append(i)
        cols.append((i + 1) % nx)
        vals.append(inv_dx2)

    idx = torch.tensor([rows, cols], dtype=torch.long)
    val = torch.tensor(vals, dtype=torch.float64)
    L = torch.sparse_coo_tensor(idx, val, (nx, nx)).coalesce()
    L = d * L

    Id = torch.eye(nx, dtype=torch.float64)

    # -------------------------
    # Initial condition
    # -------------------------
    u = f_ini_for_tensors(x)

    # -------------------------
    # Newton settings
    # -------------------------
    tol_newton = 1e-10
    max_newton_iter = 25

    # -------------------------
    # CN time-step function
    # -------------------------
    def step(u_n):
        A_n_u = u_n + 0.5 * dt * torch.sparse.mm(L, u_n.unsqueeze(1)).squeeze()
        u_new = u_n.clone()

        for _ in range(max_newton_iter):
            Lu_new = torch.sparse.mm(L, u_new.unsqueeze(1)).squeeze()
            nonlin = k * (u_new - u_new**3)

            F = (u_new - 0.5 * dt * Lu_new) - A_n_u - dt * nonlin
            if torch.norm(F) < tol_newton:
                break

            diag_term = k * (1.0 - 3.0 * u_new**2)
            J = Id - 0.5 * dt * L.to_dense() - dt * torch.diag(diag_term)

            du = torch.linalg.solve(J, -F)
            u_new = u_new + du

        return u_new

    # -------------------------
    # Time loop
    # -------------------------
    save_times = torch.linspace(0, t_final, snapshot_count).tolist()

    snapshots = {}
    if snapshot_count > 1:
        snapshots = {0.0: u.clone()}

    t = 0.0

    for _ in range(1, nt + 1):
        t += dt
        u = step(u)

        for ts in save_times:
            if ts not in snapshots and abs(t - ts) < 0.5 * dt:
                snapshots[ts] = u.clone()

    if snapshot_count > 0:
        snapshots[t_final] = u.clone()

    return x, u, snapshots


def exact(t: LabelTensor, x: LabelTensor, mu: LabelTensor) -> torch.Tensor:
    x = x.get_components().squeeze().detach().cpu().numpy()
    t = t.get_components()
    d, k = mu.get_components()

    # time should have the same value everywhere for the numerical solution
    assert torch.allclose(t, t[0] * torch.ones_like(t)), (
        "The exact solution should not be called with varying t."
    )

    # the parameters should have the same value everywhere for the numerical solution
    assert torch.allclose(d, d[0] * torch.ones_like(d)), (
        "The exact solution should not be called with varying d."
    )
    assert torch.allclose(k, k[0] * torch.ones_like(k)), (
        "The exact solution should not be called with varying k."
    )

    t = t[0].item()
    d = d[0].item()
    k = k[0].item()

    x_ex, u_ex, _ = solve_allen_cahn(d, k, t)

    x_ex = x_ex.detach().cpu().numpy()
    u_ex = u_ex.detach().cpu().numpy()
    u_ex_spline = RegularGridInterpolator(
        (x_ex,), u_ex, method="cubic", bounds_error=False, fill_value=None
    )

    return torch.tensor(u_ex_spline(x)).unsqueeze(1)


# %%

torch.random.manual_seed(0)

domain_x = Segment1D((-1, 1), is_main_domain=True)
domain_mu = [(1e-4, 5e-4), (2, 4)]
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler(domain_mu)]
)

space = NNxSpace(
    1,
    2,
    PeriodicMLP,
    domain_x,
    sampler,
    layer_sizes=[10] * 7,
    parameters_bounds=domain_mu,
)

pde = AllenCahnEquation1DStrongForm(space, init=f_ini)
projector = TimeDiscreteNaturalGradientProjector(pde, rhs=f_ini)

use_discrete_PINN = False
if use_discrete_PINN:
    scheme = DiscretePINN(pde, projector, scheme="rk4")
else:
    scheme = NeuralGalerkin(pde, projector, scheme="rk4")

# %%

scheme.initialization(epochs=250, verbose=False, n_collocation=3000)

scheme.projector.save("allen_cahn_ini")

plot_time_discrete_scheme(scheme, error=exact, solution=exact)
plt.show()

# %%

scheme.projector.load("allen_cahn_ini")

dt = 0.02
T = 1.2
scheme.solve(dt=dt, final_time=T, epochs=15, n_collocation=2000, verbose=False)

scheme.projector.save("allen_cahn_final")

plot_time_discrete_scheme(scheme, error=exact, solution=exact)
plt.show()

# %%