"""Build tracked complex/PML modal bases for a slowly varying taper."""

from __future__ import annotations

import numpy

from meanas.fdmath import vec
from meanas.fdfd import eme, scpml, waveguide_2d_complex


WAVELENGTH = 1.55
OMEGA = 2 * numpy.pi / WAVELENGTH
SHAPE = (17, 17)
SPACING = 0.25
PML_CELLS = 3
EPS_CLADDING = 1.45 ** 2
EPS_CORE = 2.0 ** 2


def build_grids() -> tuple[list[list[numpy.ndarray]], list[list[numpy.ndarray]]]:
    """Build matching physical and complex-stretched transverse Yee grids.

    Returns:
        The real overlap grid followed by the PML-stretched solve grid.
    """
    physical = [[numpy.full(size, SPACING) for size in SHAPE] for _ in range(2)]
    stretched = [[values.copy() for values in family] for family in physical]
    for axis in (0, 1):
        for polarity in (-1, +1):
            stretched = scpml.stretch_with_scpml(
                stretched,
                axis=axis,
                polarity=polarity,
                omega=OMEGA,
                epsilon_effective=EPS_CLADDING,
                thickness=PML_CELLS,
            )
    return physical, stretched


def build_epsilon(edge_fill: float) -> numpy.ndarray:
    """Construct and vectorize one taper cross-section permittivity profile.

    Args:
        edge_fill: Fractional core fill used for the two cells beside the
            central three-cell-wide core.

    Returns:
        Vectorized three-component relative permittivity field.
    """
    epsilon = numpy.full((3, *SHAPE), EPS_CLADDING)
    center_x, center_y = (size // 2 for size in SHAPE)
    epsilon[:, center_x - 1:center_x + 2, center_y] = EPS_CORE
    mixed_epsilon = EPS_CLADDING + edge_fill * (EPS_CORE - EPS_CLADDING)
    epsilon[:, center_x - 1:center_x + 2, center_y - 1] = mixed_epsilon
    epsilon[:, center_x - 1:center_x + 2, center_y + 1] = mixed_epsilon
    return vec(epsilon)


def main() -> None:
    """Solve, track, and summarize the sampled complex modal bases."""
    physical_dxes, solve_dxes = build_grids()
    z_coords = numpy.array([0.0, 5.0, 10.0])
    bases = [
        waveguide_2d_complex.solve_modes_near(
            OMEGA,
            solve_dxes,
            build_epsilon(edge_fill),
            target_wavenumbers=[1.45 * OMEGA],
            modes_per_target=3,
            candidate_margin=3,
            overlap_dxes=physical_dxes,
        )
        for edge_fill in (0.30, 0.35, 0.40)
    ]
    tracked = waveguide_2d_complex.track_mode_path(bases)
    sections = eme.mode_sections_from_bases(z_coords, tracked)

    for z_coord, basis in zip(z_coords, tracked, strict=True):
        print(
            f'z={z_coord:5.1f}: '
            f'beta={basis.wavenumbers[0]:.6g}, '
            f'PML participation={basis.diagnostics.pml_participation[0]:.3f}, '
            f'right residual={basis.diagnostics.right_residuals[0]:.2e}',
        )
    print(f'constructed {len(sections)} tracked EME sections with {len(sections[0].modes)} modes each')


if __name__ == '__main__':
    main()