"""The ``-Delta u = f`` data on a mesh of the lopsided domain: sources, a batched P1 solve, the error.

The same data as ``gnn_unet_lopsided_laplacian.py`` (whose own copy stays as it
was measured): ``f`` a random mixture of Gaussians, ``u`` a P1 finite-element
solve with ``u = 0`` on the wall, both read at the cell centres.
"""

from __future__ import annotations

import jax
import jax.numpy as jnp
import numpy as np

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.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.classical_weakform.laplacian_weak_form import (
    LaplacianWeakForm,
)
from scimba_jax.utils.scimba_pytree import ScimbaPytree

N_GAUSSIANS = 3


class GaussianMixture(ScimbaPytree):
    """``f(x) = sum_i A_i exp(-|x - c_i|^2 / 2 sigma_i^2)`` -- a pytree, so it batches.

    Args:
        centers: ``(n_gaussians, 2)``.
        amplitudes: ``(n_gaussians,)``.
        sigmas: ``(n_gaussians,)``.
    """

    def __init__(self, centers, amplitudes, sigmas):
        self.centers = jnp.asarray(centers)
        self.amplitudes = jnp.asarray(amplitudes)
        self.sigmas = jnp.asarray(sigmas)

    def __call__(self, x):
        """The source at one point.

        Args:
            x: a point, ``(2,)``.

        Returns:
            a scalar.
        """
        squared = jnp.sum((x - self.centers) ** 2, axis=1)
        return jnp.sum(self.amplitudes * jnp.exp(-squared / (2.0 * self.sigmas**2)))


class FiniteElementSolver:
    """A batched P1 Laplacian solve on one mesh, read at given points.

    The matrix does not depend on the source: factorised once, only the
    right-hand side is batched.

    Args:
        nodes: ``(n_nodes, 2)``.
        cells: ``(n_cells, 4)``.
        points: where ``f`` and ``u`` are read, ``(n_points, 2)``.
    """

    def __init__(self, nodes, cells, points):
        mesh = UnstructuredMesh(
            nodes=np.asarray(nodes),
            cells=np.asarray(cells),
            ref_quad=UnitSquareTensorized(dim=2, order=4),
            order=1,
        )
        basis = AnalyticBasis(
            nb_basis=4,
            out_dim=1,
            mesh=mesh,
            basis_type="scalar",
            local_basis_by_logical=lambda y, i, m: local_lagrange_basis_by_logical(
                y, i, m, order=1, out_dim=1
            ),
            local_basis=lambda y, i, m: local_lagrange_basis(
                y, i, m, order=1, out_dim=1
            ),
        )
        self.variables = VariablesFE(
            basis=basis, nb_variables=1, dof_map=UnstructuredLagrangeDofMap
        )
        self.points = jnp.asarray(points)
        reference = self._scheme(
            GaussianMixture(
                jnp.zeros((N_GAUSSIANS, 2)),
                jnp.zeros(N_GAUSSIANS),
                jnp.ones(N_GAUSSIANS),
            )
        )
        self._dofs_init = reference._initial_dofs()
        self._factorisation = EllipticFEscheme.factorise(reference, self._dofs_init)
        self._back_solve = EllipticFEscheme._make_back_solve_fn()

    def _scheme(self, source):
        model = AbstractPhysicalWeakModel.from_weak_form(
            LaplacianWeakForm(dim=2, f=source), dirichlet=lambda x: jnp.zeros(1)
        )
        return EllipticFEscheme(model, self.variables)

    def __call__(self, centers, amplitudes, sigmas):
        """``(f, u)`` at the points for a batch of sources.

        Args:
            centers: ``(batch, n_gaussians, 2)``.
            amplitudes: ``(batch, n_gaussians)``.
            sigmas: ``(batch, n_gaussians)``.

        Returns:
            ``(f, u)``, each ``(batch, n_points, 1)``.
        """

        @jax.jit
        def solve(c, a, s):
            def one(c1, a1, s1):
                scheme = self._scheme(GaussianMixture(c1, a1, s1))
                dofs = self._back_solve(
                    scheme,
                    self._dofs_init,
                    self._factorisation.lu,
                    self._factorisation.pivots,
                )
                u = jax.vmap(
                    lambda p: VariablesFE._classical_local_evaluate_pure(
                        self.variables, dofs, p
                    )
                )(self.points)
                f = jax.vmap(GaussianMixture(c1, a1, s1))(self.points)[:, None]
                return f, u

            return jax.vmap(one)(c, a, s)

        return solve(centers, amplitudes, sigmas)


def draw_sources(key, n, candidates, wall_tree, margin=0.25):
    """Random Gaussian mixtures centred away from the wall.

    Args:
        key: a random state.
        n: how many sources.
        candidates: points to draw centres among, ``(m, 2)``.
        wall_tree: a KD-tree on the wall's points.
        margin: the least distance of a centre to the wall.

    Returns:
        ``(centers, amplitudes, sigmas)``.
    """
    keys = jax.random.split(key, 3)
    distance, _ = wall_tree.query(np.asarray(candidates))
    allowed = jnp.asarray(np.nonzero(distance > margin)[0])
    picked = jax.random.choice(keys[0], allowed, (n * N_GAUSSIANS,), replace=True)
    centers = jnp.asarray(candidates)[picked].reshape(n, N_GAUSSIANS, 2)
    amplitudes = jax.random.uniform(keys[1], (n, N_GAUSSIANS), minval=0.5, maxval=2.0)
    sigmas = jax.random.uniform(keys[2], (n, N_GAUSSIANS), minval=0.20, maxval=0.40)
    return centers, amplitudes, sigmas


def relative_error(predicted, truth):
    """Relative L2 error per sample.

    Args:
        predicted: ``(batch, ...)``.
        truth: ``(batch, ...)``.

    Returns:
        ``(batch,)``.
    """
    flat = predicted.reshape(predicted.shape[0], -1) - truth.reshape(truth.shape[0], -1)
    return jnp.linalg.norm(flat, axis=1) / jnp.linalg.norm(
        truth.reshape(truth.shape[0], -1), axis=1
    )