"""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


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,
                max_imag_ratio=1.0,
            )
    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])
    epsilon_samples = [build_epsilon(edge_fill) for edge_fill in (0.30, 0.35, 0.40)]
    solved = [
        waveguide_2d.solve_modes_near(
            OMEGA,
            solve_dxes,
            epsilon,
            target_wavenumbers=[1.45 * OMEGA],
            modes_per_target=3,
            candidate_margin=3,
            overlap_dxes=physical_dxes,
        )
        for epsilon in epsilon_samples
    ]
    bases = [basis for basis, _diagnostics in solved]
    tracked = waveguide_2d.track_mode_path(bases, overlap_dxes=physical_dxes)
    tracked_diagnostics = [
        waveguide_2d.compute_mode_diagnostics(
            basis,
            OMEGA,
            solve_dxes,
            epsilon,
            target_wavenumbers=[1.45 * OMEGA],
            overlap_dxes=physical_dxes,
        )
        for basis, epsilon in zip(tracked, epsilon_samples, strict=True)
    ]
    sections = eme.mode_sections_from_bases(z_coords, tracked)
    prepared = eme.prepare_taper(sections, physical_dxes)
    # Fields and interface overlaps are reused: these are frozen-basis length
    # sweeps. Wavelength-dependent beta overrides would also freeze the fields.
    length_sweep = [eme.evaluate_taper_s(prepared, z_coords=z_coords * scale)
                    for scale in (0.8, 1.0, 1.2)]

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


if __name__ == '__main__':
    main()