r"""Solves the Grad-Shafranov equation in 2D on a tokamak-like geometry using PINNs.

.. math::

    \partial_{xx} \psi + \partial_{yy} \psi - \frac 1 {x_1} \partial_x \psi & = (k_1^2 x_1^2 + k_2^2) in \Omega \\
    \psi & = 0 on \partial \Omega

where :math:`x = (x_1, x_2) \in \Omega` (with :math:`\Omega` resembling a tokamak cross-section), and where :math:`k_1 = 1.05`, :math:`k_2 = 1`.

The boundary conditions are enforced weakly.

The neural network used is a simple MLP (Multilayer Perceptron); the optimization is done using Adam.

Two tokamak-like geometries are considered: one with a hole and one without hole.
"""

import matplotlib.pyplot as plt
import torch

from scimba_torch.approximation_space.abstract_space import AbstractApproxSpace
from scimba_torch.approximation_space.nn_space import NNxSpace
from scimba_torch.domain.meshless_domain.domain_2d import Disk2D
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.elliptic_pde.pinns import PinnsElliptic
from scimba_torch.optimizers.losses import GenericLosses
from scimba_torch.optimizers.optimizers_data import OptimizerData
from scimba_torch.physical_models.elliptic_pde.abstract_elliptic_pde import (
    StrongFormEllipticPDE,
)
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces
from scimba_torch.utils.mapping import Mapping
from scimba_torch.utils.scimba_tensors import LabelTensor

shape = {}
shape["amin"] = 0.6
shape["R_axis"] = 0.93
shape["Z_axis"] = 0.0
shape["GS_shift"] = 0.06
shape["ellip"] = 1.8
shape["rect_u"] = -0.2
shape["tria_u"] = 0.5
shape["rect_l"] = -0.2
shape["tria_l"] = 0.5
shape["symmetry"] = 0.0


def mapping_tokamak(x):
    x1, x2 = x[:, 0], x[:, 1]
    R0 = shape["R_axis"] - shape["GS_shift"]
    Z0 = shape["Z_axis"]
    theta = torch.atan2(x2 - Z0, x1 - R0)

    r = ((x1 - R0) ** 2 + (x2 - Z0) ** 2) ** 0.5

    R1 = (
        r
        * torch.cos(
            theta
            + shape["tria_u"] * torch.sin(theta)
            - shape["rect_u"] * torch.sin(2.0 * theta)
        )
        + R0
    )

    R2 = (
        r
        * torch.cos(
            theta
            + shape["tria_l"] * torch.sin(theta)
            - shape["rect_l"] * torch.sin(2.0 * theta)
        )
        + R0
    )

    R = torch.where(theta < torch.pi, R1, R2)

    Z = shape["ellip"] * r * torch.sin(theta) + Z0

    res = torch.cat([R[:, None], Z[:, None]], dim=1)
    return res


class GradShafranovStrongForm(StrongFormEllipticPDE):
    def __init__(self, space: AbstractApproxSpace, **kwargs):
        r"""Initialize the 2D Laplacian problem with Dirichlet boundary conditions.

        :param space: The approximation space for the problem.
        :type space: AbstractApproxSpace
        :param f: Callable representing the source term \( f(x, \mu) \).
        :type f: Callable
        :param g: Callable representing the Dirichlet boundary condition \( g(x, \mu) \).
        :type g: Callable
        :param kwargs: Additional keyword arguments.
        :type kwargs: dict
        """
        super().__init__(space, residual_size=1, bc_residual_size=1, **kwargs)

    def rhs(self, w, x, mu):
        # psi = w.get_components()
        x1, _ = x.get_components()
        k1 = 1.05
        k2 = 1
        return k1**2 * x1**2 + k2**2

    def operator(self, w, x, mu):
        x1, _ = x.get_components()
        psi = w.get_components()
        # alpha = mu.get_components()
        psi_x, psi_y = self.grad(psi, x)
        psi_xx, _ = self.grad(psi_x, x)
        _, psi_yy = self.grad(psi_y, x)
        return psi_xx + psi_yy - (1 / x1) * psi_x

    def bc_rhs(self, w, x, n, mu):
        psi = w.get_components()
        # psi_0 = 0
        return psi * 0

    def bc_operator(self, w, x, n, mu):
        u = w.get_components()
        return u


R0 = 2.0
mapping = Mapping(2, 2, map=mapping_tokamak)
domain_x = Disk2D(
    (shape["R_axis"] - shape["GS_shift"], 0.0),
    shape["amin"],
    is_main_domain=True,
)
# add a hole
domain_x.add_hole(Disk2D([1.0, 0.0], 0.3))
domain_x.set_mapping(mapping, bounds_postmap=[(0.25, 1.5), (-1.1, 1.1)], to_holes=True)

sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler([(1.0, 1.000001)])]
)

domain_x_no_hole = Disk2D(
    (shape["R_axis"] - shape["GS_shift"], 0.0),
    shape["amin"],
    is_main_domain=True,
)
# add a hole
domain_x_no_hole.set_mapping(mapping, bounds_postmap=[(0.25, 1.5), (-1.1, 1.1)])

sampler_no_hole = TensorizedSampler(
    [DomainSampler(domain_x_no_hole), UniformParametricSampler([(1.0, 1.000001)])]
)


def post_processing(inputs: torch.Tensor, x: LabelTensor, mu: LabelTensor):
    R, Z = x.get_components()
    # r = torch.sqrt((x1 - R0) ** 2.0 + x2**2.0)
    psi_0 = 0.5
    k1 = 0.2
    k2 = 1
    R0 = shape["R_axis"] - shape["GS_shift"]
    Z0 = shape["Z_axis"]
    return (
        inputs * psi_0 * torch.cos(k1 / 2 * (R**2 - R0**2) * torch.cos(k2 * (Z - Z0)))
    )


space = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[40] * 3,
    # post_processing=post_processing,
)
space_no_hole = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x_no_hole,
    sampler_no_hole,
    layer_sizes=[40] * 3,
    # post_processing=post_processing,
)
pde = GradShafranovStrongForm(space)
pde_no_hole = GradShafranovStrongForm(space_no_hole)
losses1 = GenericLosses(
    [
        ("residual", torch.nn.MSELoss(), 1.0),
        ("bc", torch.nn.MSELoss(), 80.0),
    ],
)
losses_no_hole = GenericLosses(
    [
        ("residual", torch.nn.MSELoss(), 1.0),
        ("bc", torch.nn.MSELoss(), 80.0),
    ],
)
opt_1 = {
    "name": "adam",
    "optimizer_args": {"lr": 2.0e-2, "betas": (0.9, 0.999)},
}
opt = OptimizerData(opt_1)
pinns = PinnsElliptic(pde, bc_type="weak", optimizers=opt, losses=losses1)
pinns.solve(epochs=400, n_collocation=5000, n_bc_collocation=2000, verbose=True)
pinns.save("grad_shafranov_with_hole")

pinns_no_hole = PinnsElliptic(
    pde_no_hole, bc_type="weak", optimizers=OptimizerData(opt_1), losses=losses_no_hole
)
pinns_no_hole.solve(epochs=400, n_collocation=5000, n_bc_collocation=2000, verbose=True)
pinns_no_hole.save("grad_shafranov")

# pinns.load("grad_shafranov_with_hole")
# pinns.space.load_from_best_approx()

# pinns_no_hole.load("grad_shafranov")
# pinns_no_hole.space.load_from_best_approx()

# plot_AbstractApproxSpaces(
#     pinns.space,
#     domain_x,
#     [[1.0, 1.000001]],
#     loss=pinns.losses,
#     residual=pde,
#     cuts=[
#         ([0.0, 0.5], [-0.5, 0.5]),
#         # ([0.0, 0.5], [-0.5, 0.5]),
#         ([0.0, 0.5], [0.5, 0.5]),
#     ],
#     draw_contours=True,
#     n_drawn_contours=20,
#     parameters_value="random",
# )

# plt.show()

plot_abstract_approx_spaces(
    (pinns.space, pinns_no_hole.space),
    (domain_x, domain_x_no_hole),
    [(1.0, 1.000001)],
    loss=(pinns.losses, pinns_no_hole.losses),
    residual=(pde, pde_no_hole),
    cuts=[
        ([0.0, 0.5], [-0.5, 0.5]),
    ],
    draw_contours=True,
    n_drawn_contours=20,
    parameters_values="random",
)

plt.show()