"""
Mode-matching / EME example for a straight-to-bent waveguide interface.

This example demonstrates a cylindrical-waveguide EME workflow:

1. build a rib-waveguide cross-section,
2. solve straight port modes with `waveguide_2d`,
3. solve bend modes with `waveguide_cyl`,
4. assemble the straight-to-bend interface scattering matrix with
   `meanas.fdfd.eme.get_s(...)`,
5. optionally cascade a straight section, bend section, and interface pair into
   a compact multiport network using `scikit-rf`.
"""

from __future__ import annotations

import importlib

import numpy
from numpy import pi
from scipy import sparse

import gridlock
from gridlock import Extent

from meanas.fdfd import eme, waveguide_2d, waveguide_cyl
from meanas.fdmath import unvec


WL = 1310.0
DX = 40.0
RADIUS = 6e3
WIDTH = 400.0
THF = 161.0
THP = 77.0
EPS_SI = 3.51 ** 2
EPS_OX = 1.453 ** 2
MODE_NUMBERS = numpy.array([0])
STRAIGHT_SECTION_LENGTH = 12e3
BEND_ANGLE = pi / 2


def require_optional(name: str, package_name: str | None = None):
    package_name = package_name or name
    try:
        return importlib.import_module(name)
    except ImportError as exc:  # pragma: no cover - user environment guard
        raise SystemExit(
            f"This example requires the optional package '{package_name}'. "
            "Install example dependencies with `pip install -e './meanas[examples]'`.",
        ) from exc


def build_geometry(
        *,
        dx: float = DX,
        width: float = WIDTH,
        thf: float = THF,
        thp: float = THP,
        eps_si: float = EPS_SI,
        eps_ox: float = EPS_OX,
        ) -> tuple[gridlock.Grid, numpy.ndarray, list[list[numpy.ndarray]], float]:
    x0 = (width / 2) % dx
    omega = 2 * pi / WL

    grid = gridlock.Grid(
        [
            numpy.arange(-800, 800 + dx, dx),
            numpy.arange(-400, 400 + dx, dx),
            numpy.arange(-2 * dx, 2 * dx + dx, dx),
        ],
        periodic=True,
    )
    epsilon = grid.allocate(eps_ox)

    grid.draw_cuboid(
        epsilon,
        foreground=eps_si,
        x=Extent(center=x0, span=width + 1200),
        y=Extent(min=0, max=thf),
        z=Extent(min=-1e6, max=0),
    )
    grid.draw_cuboid(
        epsilon,
        foreground=eps_ox,
        x=Extent(max=-width / 2, span=300),
        y=Extent(min=thp, max=1e6),
        z=Extent(min=-1e6, center=0),
    )
    grid.draw_cuboid(
        epsilon,
        foreground=eps_ox,
        x=Extent(min=width / 2, span=300),
        y=Extent(min=thp, max=1e6),
        z=Extent(min=-1e6, center=0),
    )

    dxes = [grid.dxyz, grid.autoshifted_dxyz()]
    dxes_2d = [[d[0], d[1]] for d in dxes]
    return grid, epsilon, dxes_2d, omega


def solve_straight_modes(
        epsilon_slice: 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]:
    e_xys, wavenumbers = waveguide_2d.solve_modes(
        epsilon=epsilon_slice.ravel(),
        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_slice.ravel(),
        )
        for e_xy, wavenumber in zip(e_xys, wavenumbers, strict=True)
    ]
    return eh_fields, wavenumbers


def solve_bend_modes(
        epsilon_slice: numpy.ndarray,
        *,
        omega: float,
        dxes_2d: list[list[numpy.ndarray]],
        rmin: float,
        mode_numbers: numpy.ndarray = MODE_NUMBERS,
        ) -> tuple[list[tuple[numpy.ndarray, numpy.ndarray]], numpy.ndarray, numpy.ndarray]:
    e_xys, angular_wavenumbers = waveguide_cyl.solve_modes(
        epsilon=epsilon_slice.ravel(),
        omega=omega,
        dxes=dxes_2d,
        mode_numbers=mode_numbers,
        rmin=rmin,
    )
    linear_wavenumbers = waveguide_cyl.linear_wavenumbers(
        e_xys=e_xys,
        angular_wavenumbers=angular_wavenumbers,
        rmin=rmin,
        epsilon=epsilon_slice.ravel(),
        dxes=dxes_2d,
    )
    eh_fields = [
        waveguide_cyl.normalized_fields_e(
            e_xy,
            angular_wavenumber=angular_wavenumber,
            dxes=dxes_2d,
            omega=omega,
            epsilon=epsilon_slice.ravel(),
            rmin=rmin,
        )
        for e_xy, angular_wavenumber in zip(e_xys, angular_wavenumbers, strict=True)
    ]
    return eh_fields, linear_wavenumbers, angular_wavenumbers


def build_cascaded_network(
        skrf,
        *,
        interface_s: numpy.ndarray,
        straight_wavenumbers: numpy.ndarray,
        bend_angular_wavenumbers: numpy.ndarray,
        n_modes: int,
        ) -> object:
    net_sb = skrf.Network(f=[1 / WL], s=interface_s[numpy.newaxis, ...])
    net_bs = net_sb.copy()
    net_bs.renumber(numpy.arange(2 * n_modes), numpy.roll(numpy.arange(2 * n_modes), n_modes))

    straight_phase = sparse.diags_array(numpy.exp(-1j * straight_wavenumbers[:n_modes] * STRAIGHT_SECTION_LENGTH))
    bend_phase = sparse.diags_array(numpy.exp(-1j * bend_angular_wavenumbers[:n_modes] * BEND_ANGLE))
    zero = numpy.zeros((n_modes, n_modes), dtype=complex)

    straight_s = numpy.block([[zero, straight_phase.toarray()], [straight_phase.toarray(), zero]])
    bend_s = numpy.block([[zero, bend_phase.toarray()], [bend_phase.toarray(), zero]])
    net_straight = skrf.Network(f=[1 / WL], s=straight_s[numpy.newaxis, ...])
    net_bend = skrf.Network(f=[1 / WL], s=bend_s[numpy.newaxis, ...])

    return skrf.network.cascade_list([net_straight, net_sb, net_bend, net_bs, net_straight])


def print_summary(
        interface_s: numpy.ndarray,
        cascaded_s: numpy.ndarray,
        straight_wavenumbers: numpy.ndarray,
        bend_linear_wavenumbers: numpy.ndarray,
        bend_angular_wavenumbers: numpy.ndarray,
        omega: float,
        n_modes: int,
        ) -> None:
    straight_neff = numpy.real(straight_wavenumbers / omega)
    bend_neff = numpy.real(bend_linear_wavenumbers / omega)
    print('straight effective indices:', ', '.join(f'{value:.5f}' for value in straight_neff[:4]))
    print('bend effective indices   :', ', '.join(f'{value:.5f}' for value in bend_neff[:4]))
    print('bend angular wavenumbers :', ', '.join(f'{value:.5e}' for value in bend_angular_wavenumbers[:4]))

    interface_transmission = abs(interface_s[n_modes, 0]) ** 2
    interface_reflection = abs(interface_s[0, 0]) ** 2
    print(f'interface fundamental transmission |S_{n_modes},0|^2 = {interface_transmission:.6f}')
    print(f'interface fundamental reflection   |S_00|^2          = {interface_reflection:.6f}')

    total_cascaded_output = numpy.sum(abs(cascaded_s[:, 0]) ** 2)
    bend_through = abs(cascaded_s[n_modes, 0]) ** 2
    bend_reflection = abs(cascaded_s[0, 0]) ** 2
    print(f'cascaded bend through power        |S_{n_modes},0|^2 = {bend_through:.6f}')
    print(f'cascaded bend reflection           |S_00|^2          = {bend_reflection:.6f}')
    print(f'cascaded left-incident total output power            = {total_cascaded_output:.6f}')


def plot_results(
        *,
        pyplot,
        interface_s: numpy.ndarray,
        cascaded_s: numpy.ndarray,
        straight_mode: tuple[numpy.ndarray, numpy.ndarray],
        bend_mode: tuple[numpy.ndarray, numpy.ndarray],
        shape_2d: tuple[int, int],
        ) -> None:
    fig_s, axes = pyplot.subplots(1, 2, figsize=(12, 4))
    interface_plot = axes[0].imshow(abs(interface_s) ** 2, origin='lower', cmap='magma')
    fig_s.colorbar(interface_plot, ax=axes[0])
    axes[0].set_title('Straight-to-bend interface |S|^2')
    axes[0].set_xlabel('Incident mode index')
    axes[0].set_ylabel('Outgoing mode index')

    cascaded_plot = axes[1].imshow(abs(cascaded_s) ** 2, origin='lower', cmap='magma')
    fig_s.colorbar(cascaded_plot, ax=axes[1])
    axes[1].set_title('Cascaded bend network |S|^2')
    axes[1].set_xlabel('Incident mode index')
    axes[1].set_ylabel('Outgoing mode index')

    straight_e = unvec(straight_mode[0], shape_2d)
    bend_e = unvec(bend_mode[0], shape_2d)
    straight_intensity = numpy.sum(abs(straight_e) ** 2, axis=0).T
    bend_intensity = numpy.sum(abs(bend_e) ** 2, axis=0).T

    fig_modes, axes_modes = pyplot.subplots(1, 2, figsize=(10, 4))
    straight_plot = axes_modes[0].imshow(straight_intensity, origin='lower', cmap='viridis')
    fig_modes.colorbar(straight_plot, ax=axes_modes[0])
    axes_modes[0].set_title('Straight fundamental mode |E|^2')
    bend_plot = axes_modes[1].imshow(bend_intensity, origin='lower', cmap='viridis')
    fig_modes.colorbar(bend_plot, ax=axes_modes[1])
    axes_modes[1].set_title('Bent fundamental mode |E|^2')
    if pyplot.get_backend().lower().endswith('agg'):
        pyplot.close(fig_s)
        pyplot.close(fig_modes)
    else:
        pyplot.show()


def main() -> None:
    pyplot = require_optional('matplotlib.pyplot', package_name='matplotlib')
    skrf = require_optional('skrf', package_name='scikit-rf')

    grid, epsilon, dxes_2d, omega = build_geometry()
    epsilon_slice = epsilon[:, :, :, 2]
    rmin = RADIUS + grid.xyz[0].min()

    straight_modes, straight_wavenumbers = solve_straight_modes(epsilon_slice, omega=omega, dxes_2d=dxes_2d)
    bend_modes, bend_linear_wavenumbers, bend_angular_wavenumbers = solve_bend_modes(
        epsilon_slice,
        omega=omega,
        dxes_2d=dxes_2d,
        rmin=rmin,
    )

    interface_s = eme.get_s(
        straight_modes,
        straight_wavenumbers,
        bend_modes,
        bend_linear_wavenumbers,
        dxes=dxes_2d,
    )
    cascaded = build_cascaded_network(
        skrf,
        interface_s=interface_s,
        straight_wavenumbers=straight_wavenumbers,
        bend_angular_wavenumbers=bend_angular_wavenumbers,
        n_modes=len(MODE_NUMBERS),
    )
    cascaded_s = cascaded.s[0]

    print_summary(
        interface_s,
        cascaded_s,
        straight_wavenumbers,
        bend_linear_wavenumbers,
        bend_angular_wavenumbers,
        omega,
        len(MODE_NUMBERS),
    )
    plot_results(
        pyplot=pyplot,
        interface_s=interface_s,
        cascaded_s=cascaded_s,
        straight_mode=straight_modes[0],
        bend_mode=bend_modes[0],
        shape_2d=grid.shape[:2],
    )


if __name__ == '__main__':
    main()