"""
Real-valued straight-waveguide FDTD/FDFD comparison.

This example shows the user-facing "compare real FDTD against reconstructed real
FDFD" workflow:

1. build a straight waveguide on a uniform Yee grid,
2. drive it with a real-valued continuous-wave mode source,
3. solve the matching FDFD problem from the analytic source phasor, and
4. compare late real monitor slices against `fdtd.reconstruct_real_e/h(...)`.

Unlike the complex-source examples, this script does not use phasor extraction
as the main output. The comparison target is the real field itself.
"""

import numpy

from meanas import fdfd, fdtd
from meanas.fdfd import functional, scpml, waveguide_3d
from meanas.fdmath import vec, unvec


DT = 0.25
PERIOD_STEPS = 64
OMEGA = 2 * numpy.pi / (PERIOD_STEPS * DT)
CPML_THICKNESS = 3
SHAPE = (3, 37, 13, 13)
SOURCE_SLICES = (slice(5, 6), slice(None), slice(None))
MONITOR_SLICES = (slice(30, 31), slice(None), slice(None))
WARMUP_PERIODS = 16


def build_uniform_dxes(shape: tuple[int, int, int, int]) -> list[list[numpy.ndarray]]:
    return [[numpy.ones(shape[axis + 1]) for axis in range(3)] for _ in range(2)]


def build_epsilon(shape: tuple[int, int, int, int]) -> numpy.ndarray:
    epsilon = numpy.ones(shape, dtype=float)
    y0 = (shape[2] - 3) // 2
    z0 = (shape[3] - 3) // 2
    epsilon[:, :, y0:y0 + 3, z0:z0 + 3] = 12.0
    return epsilon


def build_stretched_dxes(base_dxes: list[list[numpy.ndarray]]) -> list[list[numpy.ndarray]]:
    stretched_dxes = [[dx.copy() for dx in group] for group in base_dxes]
    for axis in (0, 1, 2):
        for polarity in (-1, 1):
            stretched_dxes = scpml.stretch_with_scpml(
                stretched_dxes,
                axis=axis,
                polarity=polarity,
                omega=OMEGA,
                epsilon_effective=1.0,
                thickness=CPML_THICKNESS,
            )
    return stretched_dxes


def build_cpml_params() -> list[list[dict[str, numpy.ndarray | float]]]:
    return [
        [fdtd.cpml_params(axis=axis, polarity=polarity, dt=DT, thickness=CPML_THICKNESS, epsilon_eff=1.0)
         for polarity in (-1, 1)]
        for axis in range(3)
    ]


def main() -> None:
    epsilon = build_epsilon(SHAPE)
    base_dxes = build_uniform_dxes(SHAPE)
    stretched_dxes = build_stretched_dxes(base_dxes)

    source_mode = waveguide_3d.solve_mode(
        0,
        omega=OMEGA,
        dxes=base_dxes,
        axis=0,
        polarity=1,
        slices=SOURCE_SLICES,
        epsilon=epsilon,
    )
    j_mode = waveguide_3d.compute_source(
        E=source_mode['E'],
        wavenumber=source_mode['wavenumber'],
        omega=OMEGA,
        dxes=base_dxes,
        axis=0,
        polarity=1,
        slices=SOURCE_SLICES,
        epsilon=epsilon,
    )

    e_fdfd = unvec(
        fdfd.solvers.generic(
            J=vec(j_mode),
            omega=OMEGA,
            dxes=stretched_dxes,
            epsilon=vec(epsilon),
            matrix_solver_opts={'atol': 1e-10, 'rtol': 1e-7},
        ),
        SHAPE[1:],
    )
    h_fdfd = functional.e2h(OMEGA, stretched_dxes)(e_fdfd)

    update_e, update_h = fdtd.updates_with_cpml(
        cpml_params=build_cpml_params(),
        dt=DT,
        dxes=base_dxes,
        epsilon=epsilon,
    )

    e_field = numpy.zeros_like(epsilon)
    h_field = numpy.zeros_like(epsilon)
    total_steps = (WARMUP_PERIODS + 1) * PERIOD_STEPS
    e_errors: list[float] = []
    h_errors: list[float] = []

    for step in range(total_steps):
        update_e(e_field, h_field, epsilon)

        # Real-valued FDTD uses the real part of the analytic mode source.
        t_half = (step + 0.5) * DT
        j_real = (j_mode.real * numpy.cos(OMEGA * t_half) - j_mode.imag * numpy.sin(OMEGA * t_half)).real
        e_field -= DT * j_real / epsilon

        update_h(e_field, h_field)

        if step >= total_steps - PERIOD_STEPS // 4:
            reconstructed_e = fdtd.reconstruct_real_e(
                e_fdfd[:, MONITOR_SLICES[0], :, :][numpy.newaxis, ...],
                OMEGA,
                DT,
                step + 1,
            )[0]
            reconstructed_h = fdtd.reconstruct_real_h(
                h_fdfd[:, MONITOR_SLICES[0], :, :][numpy.newaxis, ...],
                OMEGA,
                DT,
                step + 1,
            )[0]

            e_monitor = e_field[:, MONITOR_SLICES[0], :, :]
            h_monitor = h_field[:, MONITOR_SLICES[0], :, :]
            e_errors.append(numpy.linalg.norm(e_monitor - reconstructed_e) / numpy.linalg.norm(reconstructed_e))
            h_errors.append(numpy.linalg.norm(h_monitor - reconstructed_h) / numpy.linalg.norm(reconstructed_h))

    print(f'late-cycle monitor E errors: min={min(e_errors):.4f} max={max(e_errors):.4f}')
    print(f'late-cycle monitor H errors: min={min(h_errors):.4f} max={max(h_errors):.4f}')


if __name__ == '__main__':
    main()