"""Time-dependent CG-FEM on an unstructured mesh: the heat equation on a disk.

    du/dt - Delta u = 0   on the unit disk,     u = 0 on the boundary
    u(x, y, 0) = 1 - x^2 - y^2

Same point as ``examples_jax/fem/solve/classical_approach/
solve_laplacian_unstructured_2d.py`` (nothing about the discretization knows
the domain isn't a box), extended to a time-dependent solve:
``TimeDiscreteFEscheme`` reuses the exact same mesh, basis and boundary
machinery at every Runge-Kutta stage, so it costs nothing extra either.

There is no closed-form solution on a disk to check against here -- this
simply solves forward in time and plots the state at the final time, sampled
with ``sample_solution`` (cell-by-cell, so the triangulation stays inside the
disk instead of crossing it the way a convex hull of the raw point cloud
would).
"""

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

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.time_discrete_fe_scheme import (
    TimeDiscreteFEscheme,
)
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 macro_mesh_ogrid_disk
from scimba_jax.physical_models.classical_weakform.laplacian_weak_form import (
    LaplacianWeakForm,
)
from scimba_jax.plots.plots_galerkin import sample_solution
from scimba_jax.time_discrete.butcher_tableau import build_implicit_euler_tableau

ORDER = 2
DT = 0.01
NT = 15


def u0(x):
    return 1.0 - x[0] ** 2 - x[1] ** 2


macro = macro_mesh_ogrid_disk(radius=1.0, n=3, order=ORDER)
mesh = UnstructuredMesh(
    nodes=macro.nodes,
    cells=macro.cells,
    ref_quad=UnitSquareTensorized(dim=2, order=2 * ORDER + 2),
    order=ORDER,
)

basis = AnalyticBasis(
    nb_basis=(ORDER + 1) ** 2,
    out_dim=1,
    mesh=mesh,
    local_basis_by_logical=lambda y, i, m: local_lagrange_basis_by_logical(
        y, i, m, order=ORDER, out_dim=1
    ),
    local_basis=lambda y, i, m: local_lagrange_basis(y, i, m, order=ORDER, out_dim=1),
    basis_type="scalar",
)
variables = VariablesFE(basis=basis, nb_variables=1, dof_map=UnstructuredLagrangeDofMap)
scheme = TimeDiscreteFEscheme(
    spatial_weak_form_factory=LaplacianWeakForm(dim=2, f=lambda x: jnp.zeros(())),
    variables=variables,
    butcher_tableau=build_implicit_euler_tableau(),
    dt=DT,
    dirichlet=lambda x: jnp.zeros(1),
)
dofsl_init = scheme.initialize(u0)

print(f"CG-FEM, O-grid disk, {mesh.n_cells_total} isoparametric Q{ORDER} cells")
print(f"solving {NT} implicit-Euler steps of dt={DT} ...")
dofsl_final, _ = scheme.solve(dofsl_init, t0=0.0, nt=NT)
jax.block_until_ready(dofsl_final)
scheme.variables.dofsl = dofsl_final
print(f"done -- max u at t={NT * DT:.2f} is {float(dofsl_final.max()):.4f}")

points, triangles, values = sample_solution(scheme, n_side=3)

figure, ax = plt.subplots(figsize=(6, 5))
drawing = ax.tricontourf(
    points[:, 0], points[:, 1], triangles, values, levels=128, cmap="turbo"
)
ax.triplot(points[:, 0], points[:, 1], triangles, lw=0.25, color="k", alpha=0.5)
figure.colorbar(drawing, ax=ax, fraction=0.046)
ax.set_title(f"CG-FEM heat equation on the disk, t={NT * DT:.2f}")
ax.set_aspect("equal")
ax.set_xlabel("x")
ax.set_ylabel("y")
figure.tight_layout()
plt.show()