"""Vacuum flux of the ITER coils, free boundary, CG-FEM Q2.

The first half of the free-boundary Grad-Shafranov benchmark of Serino et al.
(arXiv:2407.03499): the poloidal flux of the eleven ITER coils alone, no plasma,

    -div( (1/r) grad psi ) = mu_0 J    in the half disk r^2 + z^2 < 16^2, r > 0,
    psi = 0                            on the axis r = 0,
    far-field condition                on the arc,

``J = I_i / |Omega_Ci|`` in coil ``i`` and 0 elsewhere. The arc is not a
boundary of the physics: the flux keeps going beyond it, decaying like a dipole,
and ``psi = 0`` there would be wrong. The exterior is eliminated exactly by the
nonlocal condition :class:`GradShafranovFarField` (a local term and a double
integral over the arc).

The coils are named regions of the mesh, read by ``per_region``. The check
printed at the end compares the flux at a few vacuum points with the exact one,
``mu_0 sum_i I_i / |Omega_Ci| int_{Omega_Ci} G(x, y) dy`` (the Green function of
a current loop, integrated by Gauss on each rectangle). Needs the optional
``mesh`` extra (gmsh).
"""

# %%
import time
from pathlib import Path

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.path import Path as PolygonPath

from scimba_jax.domains.tokamak import DOMAIN_RADIUS, ITER_COILS, iter_first_wall
from scimba_jax.linear_approximation.basis.analytic_bases import (
    local_lagrange_basis,
    local_lagrange_basis_by_logical,
)
from scimba_jax.linear_approximation.basis.dof_map import UnstructuredLagrangeDofMap
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.galerkin.fem.elliptic_fe_scheme import (
    EllipticFEscheme,
)
from scimba_jax.linear_approximation.meshes.unstructured_mesh import UnstructuredMesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.linear_approximation.variables.variables_fe import VariablesFE
from scimba_jax.mapping.macro_mesh import ExactHalfDisk, macro_mesh_from_regions
from scimba_jax.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.abstract_weak_form import AbstractWeakForm
from scimba_jax.physical_models.elliptic_pde.grad_shafranov_free_boundary import (
    GradShafranovFarField,
    gs_green_function,
)
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet

MU_0 = 4e-7 * np.pi
ORDER = 2  # geometry and FE degree

# Mesh size (m), as in examples/examples_jax/mesh/iter_free_boundary_mesh.py:
# fine over the machine, coarsening linearly away from it, capped far away.
H_MACHINE = 0.25
GROWTH = 0.25
H_FAR = 2.0
MACHINE_BOX = (0.5, 13.5, -9.0, 9.0)  # r_min, r_max, z_min, z_max


def mesh_size(r, z):
    r_min, r_max, z_min, z_max = MACHINE_BOX
    dr = max(r_min - r, 0.0, r - r_max)
    dz = max(z_min - z, 0.0, z - z_max)
    return min(H_FAR, H_MACHINE + GROWTH * float(np.hypot(dr, dz)))


def inverse_radius(x):
    return 1.0 / x[0]


def zero(x):
    return jnp.zeros(1)


def lagrange(y, i, mesh):
    return local_lagrange_basis(y, i, mesh, order=ORDER, out_dim=1)


def lagrange_by_logical(y, i, mesh):
    return local_lagrange_basis_by_logical(y, i, mesh, order=ORDER, out_dim=1)


class VacuumGradShafranov(AbstractWeakForm):
    """``(1/r) grad psi . grad v = mu_0 J v``, ``J`` read per coil region.

    Args:
        current_density: ``{coil name: I / |Omega_C|}`` (A/m^2).
    """

    def __init__(self, current_density: dict):
        super().__init__(dim=2)
        self.current_density = current_density
        self.inverse_radius = inverse_radius

    def bilinear_form(self, u, v):
        return self.get_fields()["inverse_radius"] * u.gradient("x").dot(
            v.gradient("x")
        )

    def linear_form(self, v):
        return MU_0 * self.per_region(self.current_density, default=0.0) * v


# %% Mesh: the half disk fragmented by the first wall and the coils.
start = time.perf_counter()
macro, region = macro_mesh_from_regions(
    ExactHalfDisk(DOMAIN_RADIUS),
    [iter_first_wall(), *(coil.corners() for coil in ITER_COILS)],
    order=ORDER,
    mesh_size=mesh_size,
)
mesh = UnstructuredMesh.from_macro_mesh(macro, UnitSquareTensorized(dim=2, order=4))
mesh.label_boundary("axis", lambda x: x[:, 0] < 1e-8)
mesh.label_boundary("farfield", lambda x: x[:, 0] >= 1e-8)
# Region 1 is the limiter (no current in vacuum), regions 2.. the coils.
for k, coil in enumerate(ITER_COILS, start=2):
    mesh.label_cells(coil.name, np.flatnonzero(region == k))
print(f"mesh: {mesh.n_cells_total} cells, {time.perf_counter() - start:.1f} s")

# %% Model and scheme.
model = AbstractPhysicalWeakModel(dim=2)
model.add_weak_form(
    "main", VacuumGradShafranov({c.name: c.current / c.area for c in ITER_COILS})
)
model.add_boundary_condition("axis", Dirichlet(zero))
model.add_boundary_condition("farfield", GradShafranovFarField(DOMAIN_RADIUS))

basis = AnalyticBasis(
    nb_basis=(ORDER + 1) ** 2,
    out_dim=1,
    mesh=mesh,
    local_basis=lagrange,
    local_basis_by_logical=lagrange_by_logical,
    basis_type="scalar",
)
variables = VariablesFE(basis=basis, nb_variables=1, dof_map=UnstructuredLagrangeDofMap)
start = time.perf_counter()
scheme = EllipticFEscheme(model, variables)
n_arc = scheme._nonlocal[0].laplacian.shape[0]
print(
    f"scheme: {variables.dofsl.shape[0]} DOFs, {n_arc} quadrature points on the "
    f"arc, {time.perf_counter() - start:.1f} s"
)
start = time.perf_counter()
scheme = EllipticFEscheme.solve(scheme)
print(f"solve: {time.perf_counter() - start:.1f} s")

nodes = np.asarray(
    variables.boundary_dof_positions(np.arange(variables.dofsl.shape[0]))
)
psi = np.asarray(scheme.variables.dofsl[:, 0])

# %% Check: the exact vacuum flux, sum of the coils' Green integrals.
gauss, gauss_w = np.polynomial.legendre.leggauss(8)
sources, weights = [], []
for coil in ITER_COILS:
    r = coil.r_min + 0.5 * (gauss + 1) * (coil.r_max - coil.r_min)
    z = coil.z_min + 0.5 * (gauss + 1) * (coil.z_max - coil.z_min)
    rr, zz = np.meshgrid(r, z, indexing="ij")
    sources.append(np.stack([rr.ravel(), zz.ravel()], axis=-1))
    w = np.outer(gauss_w, gauss_w).ravel() * 0.25  # sums to 1: the mean over the coil
    weights.append(MU_0 * coil.current * w)
sources, weights = (
    jnp.asarray(np.concatenate(sources)),
    jnp.asarray(np.concatenate(weights)),
)


@jax.jit
def exact_flux(points):
    green = jax.vmap(lambda x: jax.vmap(lambda y: gs_green_function(x, y))(sources))(
        points
    )
    return green @ weights


# Vacuum nodes far enough from every coil for 8x8 Gauss to be exact: inside the
# first wall, and on the far-field arc.
wall = iter_first_wall()
inside_wall = PolygonPath(wall).contains_points(nodes)
on_arc = np.abs(np.linalg.norm(nodes, axis=1) - DOMAIN_RADIUS) < 1e-6
for label, mask in (("inside the first wall", inside_wall), ("far-field arc", on_arc)):
    reference = np.asarray(exact_flux(jnp.asarray(nodes[mask])))
    error = np.max(np.abs(psi[mask] - reference)) / np.max(np.abs(reference))
    print(f"{label:22s}: {mask.sum():5d} nodes, max relative error {error:.2e}")

# %% Plot: the whole domain and the machine.
fig, axes = plt.subplots(1, 2, figsize=(13, 8))
for ax, limits in zip(axes, [(0.0, 16.0, -16.0, 16.0), MACHINE_BOX]):
    contour = ax.tricontourf(nodes[:, 0], nodes[:, 1], psi, levels=40, cmap="RdBu_r")
    ax.tricontour(nodes[:, 0], nodes[:, 1], psi, levels=40, colors="k", linewidths=0.3)
    ax.plot(*np.vstack([wall, wall[:1]]).T, "k-", lw=1.2)
    for coil in ITER_COILS:
        ax.add_patch(
            plt.Rectangle(
                (coil.r_min, coil.z_min),
                coil.r_max - coil.r_min,
                coil.z_max - coil.z_min,
                fill=False,
                edgecolor="k",
            )
        )
    ax.set_xlim(limits[0], limits[1])
    ax.set_ylim(limits[2], limits[3])
    ax.set_aspect("equal")
    ax.set_xlabel("r [m]")
    ax.set_ylabel("z [m]")
fig.colorbar(contour, ax=axes, label=r"$\psi$ [Wb/rad]")
axes[0].set_title("half disk, far-field condition on the arc")
axes[1].set_title("the machine")
out = Path(__file__).with_name("iter_vacuum_field.png")
fig.savefig(out, dpi=140)
print(f"Saved figure to {out}")
plt.show()