# ── Simulation parameters ──────────────────────────────────────────────────
wavelength = 1.55  # free-space wavelength (µm)
freq0 = td.C_0 / wavelength

radius = 5.0  # bend radius (µm)
width = 0.6  # waveguide width (µm)
thickness = 0.4  # waveguide thickness (µm)

num_modes = 5  # number of modes in mode solvers
target_neff = 1.65  # target effective index for mode search
plane_size = (9.0, 9.0)  # cross-section window size (µm)
min_steps_per_wvl = 30  # grid resolution (steps per wavelength)

# ── Materials ───────────────────────────────────────────────────────────────
SiO2 = td.material_library["SiO2"]["Palik_Lossless"]
# (1): extraordinary (optic) axis along global y — in the bend plane, the
# orientation-sensitive case. Use (2) to place the optic axis along z instead.
LiNbO3 = td.material_library["LiNbO3"]["Zelmon1997"](1)

ms = ModeSolver(
    simulation=fdtd_sim_test,
    mode_spec=fdtd_sim_test.sources[0].mode_spec,
    plane=fdtd_sim_test.sources[0].geometry,
    freqs=fdtd_sim_test.sources[0].source_time.freq0,
)

ms_data = web.run(ms)

df = ms_data.to_dataframe()
df["Fill factor (%)"] = ms_data.fill_fraction_box.values.squeeze() * 100
df.head(10)

# Side-by-side view of the two most-confined modes (mode_index = 0 and 1).
# TE_fraction tells us which one is quasi-TE (horizontal E in the waveguide
# plane, dominant Ey) and which is quasi-TM (vertical E, dominant Ez).
te_frac = ms_data.TE_fraction.isel(f=0).values  # shape: (num_modes,)

fig, axes = plt.subplots(1, 2, figsize=(11, 4))
for ax, m in zip(axes, [0, 1]):
    label = "quasi-TE" if te_frac[m] > 0.5 else "quasi-TM"
    component = "Ey" if te_frac[m] > 0.5 else "Ez"  # dominant transverse E
    ms.plot_field(component, "real", mode_index=m, ax=ax)
    ax.set_title(f"mode {m} — {label}  (TE frac = {te_frac[m]:.2f}) — {component}")

plt.tight_layout()
plt.show()

mode_index = 0
print("Selected mode index:", mode_index)

fdtd_sim = build_fdtd_bend(
    radius=radius,
    width=width,
    thickness=thickness,
    num_modes=num_modes,
    target_neff=target_neff,
    freq0=freq0,
    plane_size=plane_size,
    min_steps_per_wvl=min_steps_per_wvl,
    mode_index=mode_index,  # fundamental quasi-TE bend mode
)

fdtd_data = web.run(
    fdtd_sim,
    task_name="bent_LiNbO3_FDTD",
    path="FDTD_GT.hdf5",
)

# Transmittance into the fundamental guided mode at the output (|S21(0,0)|²).
amps_out = fdtd_data["mode_out"].amps.sel(direction="-")
T_fdtd = float(np.abs(amps_out.isel(f=0, mode_index=mode_index)) ** 2)
print(f"FDTD transmittance: T = {T_fdtd:.4f}  ({10 * np.log10(T_fdtd):.2f} dB)")

# Field plot
fig, ax = plt.subplots(figsize=(7, 6))
fdtd_data.plot_field("field_xy", "Ey", "real", ax=ax)
ax.set_title("FDTD — Ey (real) at z = 0")
plt.tight_layout()
plt.show()

def build_eme_bend(fdtd_sim, num_cells):
    """Build an EME bend simulation derived from an existing FDTD simulation.

    All parameters (radius, width, thickness, num_modes, target_neff, freq0,
    plane size, and materials) are extracted from fdtd_sim, guaranteeing that
    both simulations model the exact same structure.

    Parameters
    ----------
    fdtd_sim : td.Simulation
        The FDTD bend simulation returned by build_fdtd_bend().
    num_cells : int
        Number of uniform EME cells to divide the 180-degree arc into.
    """
    src = fdtd_sim.sources[0]
    ring_geom = fdtd_sim.structures[0].geometry  # ClipOperation: outer - inner

    # ── Extract all parameters from the FDTD simulation ────────────────────
    freq0 = src.source_time.freq0
    num_modes = src.mode_spec.num_modes
    target_neff = src.mode_spec.target_neff
    radius = abs(src.mode_spec.bend_radius)  # stored as negative in the source
    plane_size = (src.size[1], src.size[2])
    num_pml = src.mode_spec.num_pml
    grid_spec = fdtd_sim.grid_spec

    core_medium = fdtd_sim.structures[0].medium
    background_medium = fdtd_sim.medium

    # Width and thickness from the ring geometry (outer Cylinder - inner Cylinder)
    outer_cyl = ring_geom.geometry_a
    inner_cyl = ring_geom.geometry_b
    width = outer_cyl.radius - inner_cyl.radius
    thickness = outer_cyl.length
    # ───────────────────────────────────────────────────────────────────────

    arc_length = np.pi * radius  # 180-degree arc

    # Sort modes by fill_fraction_box so the well-confined guided mode sits at
    # index 0 at every EME cell boundary. Each cell boundary is a plane normal
    # to x at a different x, so the bounding_box needs to span all of them —
    # use td.inf along x (core centred at y = 0, z = 0 since the EME models a
    # straight waveguide with curvature encoded in bend_radius).
    core_box_eme = td.Box(center=(0, 0, 0), size=(td.inf, width, thickness))
    mode_spec = td.EMEModeSpec(
        num_modes=num_modes,
        target_neff=target_neff,
        bend_radius=radius,
        bend_axis=1,
        num_pml=num_pml,
        increasing_mode_tolerance=1e-3,
        sort_spec=td.ModeSortSpec(sort_key="fill_fraction_box", bounding_box=core_box_eme),
        precision="double",
    )

    # Infinite straight waveguide slab (curvature fully encoded in mode_spec)
    wg = td.Structure(
        geometry=td.PolySlab(
            axis=2,
            slab_bounds=(-thickness / 2, thickness / 2),
            vertices=[
                [-1e15, -width / 2],
                [-1e15, width / 2],
                [1e15, width / 2],
                [1e15, -width / 2],
            ],
        ),
        medium=core_medium,
    )

    return td.EMESimulation(
        size=(arc_length, plane_size[0], plane_size[1]),
        center=(0, 0, 0),
        structures=[wg],
        medium=background_medium,
        axis=0,
        freqs=[freq0],
        eme_grid_spec=td.EMEUniformGrid(num_cells=num_cells, mode_spec=mode_spec),
        grid_spec=grid_spec,
    )

num_cells_list = [1, 2, 3, 4, 5]

eme_sims = [build_eme_bend(fdtd_sim, num_cells=n) for n in num_cells_list]

eme_batch = {f"bent_LiNbO3_EME_N{n:02d}": sim for n, sim in zip(num_cells_list, eme_sims)}
eme_results = web.run(eme_batch)

eme_data_list = [eme_results[f"bent_LiNbO3_EME_N{n:02d}"] for n in num_cells_list]

# S00 is the fundamental-to-fundamental transmission. fill_fraction_box sorts the
# most core-confined mode to index 0, but not by polarization — so here we also
# confirm the index-0 port mode is the quasi-TE fundamental: effective index
# n_eff ~ 1.65 and TE fraction ~ 1, i.e. S00 really is the quasi-TE -> quasi-TE
# channel. (A per-cell check follows below.)
T_eme = []
for n, eme_data in zip(num_cells_list, eme_data_list):
    T = float(eme_data.smatrix.S21.isel(mode_index_in=0, mode_index_out=0, f=0).abs.item() ** 2)
    # Effective index and TE fraction of the index-0 port mode (n_eff averaged
    # over the 2 ports; TE fraction = |Ey|^2 / (|Ey|^2 + |Ez|^2), worst of the 2).
    pm = eme_data.port_modes
    neff0 = float(pm.n_complex.real.isel(mode_index=0).mean().values)
    Ey2 = (abs(pm.Ey) ** 2).sum(("x", "y", "z"))
    Ez2 = (abs(pm.Ez) ** 2).sum(("x", "y", "z"))
    te0 = float((Ey2 / (Ey2 + Ez2)).real.isel(mode_index=0).min().values)
    print(
        f"  N = {n:2d} cells  \u2192  T = {T:.4f}  ({10 * np.log10(T):.2f} dB)   |  mode-0: n_eff \u2248 {neff0:.4f}, TE fraction \u2265 {te0:.3f}"
    )
    T_eme.append(T)

fig, ax = plt.subplots(figsize=(7, 4))

ax.plot(
    num_cells_list,
    T_eme,
    "o-",
    color="steelblue",
    linewidth=1.8,
    markersize=7,
    label="EME",
)
ax.axhline(
    T_fdtd,
    color="crimson",
    linestyle="--",
    linewidth=1.6,
    label=f"FDTD reference  (T = {T_fdtd:.3f})",
)

ax.set_xlabel("Number of EME cells", fontsize=12)
ax.set_ylabel("Transmittance", fontsize=12)
ax.set_title("EME convergence — 180° LiNbO₃ bend  (r = 5 µm, λ = 1.55 µm)", fontsize=12)
ax.legend(fontsize=11)

ax.set_xlim(0, max(num_cells_list) + 1)
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

print(f"\nN={num_cells_list[0]}   overestimates by  {(T_eme[0] - T_fdtd) / T_fdtd * 100:.1f}%")
print(f"N={num_cells_list[-2]} error vs FDTD:    {abs(T_eme[-2] - T_fdtd) / T_fdtd * 100:.1f}%")
print(f"N={num_cells_list[-1]} error vs FDTD:    {abs(T_eme[-1] - T_fdtd) / T_fdtd * 100:.1f}%")

# Reuse the converged 5-cell EME, but for this mode-count study bump the solved
# basis to 20 modes (via the EMEModeSpec), and add a per-cell mode-solver monitor
# and interface-residual diagnostics.
eme_5 = eme_sims[num_cells_list.index(5)]
mode_sweep = [4, 6, 8, 10, 12, 15, 20]

mode_spec_20 = eme_5.eme_grid_spec.mode_spec.updated_copy(num_modes=20)
mode_monitor = td.EMEModeSolverMonitor(
    name="modes",
    size=(td.inf, td.inf, td.inf),
    center=(0, 0, 0),
    num_modes=6,
    eme_cell_interval_space=1,
)
eme_5_diag = eme_5.updated_copy(
    eme_grid_spec=eme_5.eme_grid_spec.updated_copy(mode_spec=mode_spec_20),
    sweep_spec=td.EMEModeSweep(num_modes=mode_sweep),
    monitors=[mode_monitor],
    eme_diagnostics=True,
)
eme_5_data = web.run(eme_5_diag, task_name="bent_LiNbO3_EME_N05_diag")

T_modes = [
    float(
        eme_5_data.smatrix.S21.isel(
            sweep_index=i, mode_index_in=0, mode_index_out=0, f=0
        ).abs.item()
        ** 2
    )
    for i in range(len(mode_sweep))
]

residual = eme_5_data.diagnostics.interface_residuals.normalized_aperture_tangential_E_residual
reduce_dims = [d for d in residual.dims if d not in ("sweep_index", "mode_index")]
residual_fund = (
    residual.isel(mode_index=0).max(dim=reduce_dims).values.ravel()
)  # worst over interfaces/ports

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4))
ax1.plot(mode_sweep, T_modes, "o-", color="steelblue", linewidth=1.8, markersize=7)
ax1.set_xlabel("Number of expansion modes")
ax1.set_ylabel("Transmittance")
ax1.set_title("Mode-count convergence")
ax1.grid(True, alpha=0.3)

ax2.semilogy(mode_sweep, residual_fund, "o-", color="darkorange", linewidth=1.8, markersize=7)
ax2.set_xlabel("Number of expansion modes")
ax2.set_ylabel("fundamental-mode aperture tangential-E residual")
ax2.set_title("Interface residual (excl. PML), incident mode 0")
ax2.grid(True, alpha=0.3, which="both")

plt.tight_layout()
plt.show()

# Mode identity: effective index and TE fraction of the index-0 mode at every cell.
# Index 0 should be the quasi-TE fundamental (TE fraction ~ 1) at every cell.
modes = eme_5_data["modes"]
n_eff = modes.n_complex.real
Ey2 = (abs(modes.Ey) ** 2).sum(("x", "y", "z"))
Ez2 = (abs(modes.Ez) ** 2).sum(("x", "y", "z"))
te_fraction = (Ey2 / (Ey2 + Ez2)).real  # |Ey|² / (|Ey|² + |Ez|²)
identity = n_eff.to_dataframe(name="n_eff").join(te_fraction.to_dataframe(name="TE_fraction"))
print(identity.xs(0, level="mode_index").round(4).to_string())

# Fundamental mode (index 0), |Ey|, across the 5 cells.
fig, axes = plt.subplots(1, 5, figsize=(16, 3.2))
for cell_index, ax in enumerate(axes):
    eme_5_data.plot_field("modes", "Ey", val="abs", eme_cell_index=cell_index, mode_index=0, ax=ax)
    ax.set_title(f"cell {cell_index} — mode 0  |Ey|")
plt.tight_layout()
plt.show()