"""
Local-mode CMT example for a continuously varying rib-waveguide taper.

This example keeps geometry construction outside `meanas.fdfd.eme`: it samples a
width taper at several cross-sections, solves and normalizes each local mode with
`waveguide_2d`, then asks `eme.get_taper_s(...)` for the bidirectional taper
scattering matrix.
"""

from __future__ import annotations

import numpy
from numpy import pi

from meanas.fdmath import vec
from meanas.fdfd import eme, waveguide_2d


WL = 1310.0
DX = 80.0
TAPER_LENGTH = 4e3
WIDTH_LEFT = 280.0
WIDTH_RIGHT = 520.0
THF = 160.0
THP = 80.0
EPS_SI = 3.51 ** 2
EPS_OX = 1.453 ** 2
MODE_NUMBERS = numpy.array([0])
N_SECTIONS = 7


def build_dxes(shape: tuple[int, int], dx: float = DX) -> list[list[numpy.ndarray]]:
    nx, ny = shape
    return [
        [numpy.full(nx, dx), numpy.full(ny, dx)],
        [numpy.full(nx, dx), numpy.full(ny, dx)],
    ]


def build_cross_section(
        *,
        width: float,
        x: numpy.ndarray,
        y: numpy.ndarray,
        eps_si: float = EPS_SI,
        eps_ox: float = EPS_OX,
        thf: float = THF,
        thp: float = THP,
        ) -> numpy.ndarray:
    epsilon = numpy.full((3, x.size, y.size), eps_ox, dtype=float)
    xx = x[:, None]
    yy = y[None, :]
    slab = (yy >= 0) & (yy <= thp)
    rib = (abs(xx) <= width / 2) & (yy >= 0) & (yy <= thf)
    epsilon[:, slab.repeat(x.size, axis=0)] = eps_si
    epsilon[:, rib] = eps_si
    return epsilon


def solve_cross_section_modes(
        epsilon: numpy.ndarray,
        *,
        omega: float,
        dxes_2d: list[list[numpy.ndarray]],
        mode_numbers: numpy.ndarray = MODE_NUMBERS,
        ) -> tuple[list[tuple[numpy.ndarray, numpy.ndarray]], numpy.ndarray]:
    epsilon_vec = vec(epsilon)
    e_xys, wavenumbers = waveguide_2d.solve_modes(
        epsilon=epsilon_vec,
        omega=omega,
        dxes=dxes_2d,
        mode_numbers=mode_numbers,
    )
    eh_fields = [
        waveguide_2d.normalized_fields_e(
            e_xy,
            wavenumber=wavenumber,
            dxes=dxes_2d,
            omega=omega,
            epsilon=epsilon_vec,
        )
        for e_xy, wavenumber in zip(e_xys, wavenumbers, strict=True)
    ]
    return eh_fields, wavenumbers


def solve_taper_sections() -> tuple[list[eme.ModeSection], list[float], float, list[list[numpy.ndarray]]]:
    omega = 2 * pi / WL
    x = numpy.arange(-480, 480 + DX, DX)
    y = numpy.arange(-240, 400 + DX, DX)
    dxes_2d = build_dxes((x.size, y.size))
    z_samples = numpy.linspace(0, TAPER_LENGTH, N_SECTIONS)
    widths = numpy.linspace(WIDTH_LEFT, WIDTH_RIGHT, N_SECTIONS)

    sections = []
    neffs = []
    for z_coord, width in zip(z_samples, widths, strict=True):
        epsilon = build_cross_section(width=float(width), x=x, y=y)
        modes, wavenumbers = solve_cross_section_modes(epsilon, omega=omega, dxes_2d=dxes_2d)
        sections.append(eme.ModeSection(float(z_coord), modes, wavenumbers))
        neffs.append(float(numpy.real(wavenumbers[0] / omega)))

    return sections, neffs, omega, dxes_2d


def print_summary(
        taper_s: numpy.ndarray,
        abrupt_s: numpy.ndarray,
        neffs: list[float],
        ) -> None:
    n_modes = len(MODE_NUMBERS)
    print('sampled taper effective indices:', ', '.join(f'{value:.5f}' for value in neffs))
    print(f'abrupt endpoint reflection       |S_00|^2          = {abs(abrupt_s[0, 0]) ** 2:.6f}')
    print(f'abrupt endpoint transmission     |S_{n_modes},0|^2 = {abs(abrupt_s[n_modes, 0]) ** 2:.6f}')
    print(f'taper CMT reflection             |S_00|^2          = {abs(taper_s[0, 0]) ** 2:.6f}')
    print(f'taper CMT transmission           |S_{n_modes},0|^2 = {abs(taper_s[n_modes, 0]) ** 2:.6f}')
    print(f'taper CMT total output power                       = {numpy.sum(abs(taper_s[:, 0]) ** 2):.6f}')


def main() -> None:
    sections, neffs, _omega, dxes_2d = solve_taper_sections()
    taper_s = eme.get_taper_s(sections, dxes=dxes_2d)
    abrupt_s = eme.get_s(
        sections[0].modes,
        sections[0].wavenumbers,
        sections[-1].modes,
        sections[-1].wavenumbers,
        dxes=dxes_2d,
    )
    print_summary(taper_s, abrupt_s, neffs)


if __name__ == '__main__':
    main()