lda0 = 1.542  # design wavelength (um)
freq0 = td.C_0 / lda0

ldas_wide = np.linspace(1.45, 1.65, 401)  # bus transmission spectrum
freqs_wide = td.C_0 / ldas_wide

ldas = np.linspace(1.522, 1.562, 41)  # far-field band around m = 39
freqs = td.C_0 / ldas

fwidth = 0.7 * (np.max(freqs_wide) - np.min(freqs_wide))

R0 = 3.9  # ring radius, i.e. the waveguide centreline (um)
w_wg = 0.5  # waveguide width
t_si = 0.22  # silicon layer thickness
t_sep = 0.275  # silica between the bus layer and the ring layer
t_top = 0.3  # silica over the ring

m_ring = 39  # ring mode number at lda0
q_values = [38, 37]  # grating element counts -> l = m - q = +1, +2

r_in, r_out = R0 - w_wg / 2, R0 + w_wg / 2

# layer stack, bottom to top: oxide | bus Si | 275 nm oxide | ring Si | 300 nm oxide | air
z_bus = (0.0, t_si)
z_ring = (t_si + t_sep, t_si + t_sep + t_si)
z_ox_top = z_ring[1] + t_top
y_bus = -R0  # the bus runs tangent to the ring, directly beneath the waveguide

tooth_r = 0.06  # inward protrusion of a grating element
tooth_a = 0.25  # azimuthal width of a grating element

si = td.material_library["cSi"]["Li1993_293K"]
sio2 = td.Medium(permittivity=1.444**2)

def grating_ring(q: int) -> td.Structure:
    """Ring annulus plus q grating teeth protruding inward from the inner wall."""
    annulus = td.Cylinder(
        center=(0, 0, np.mean(z_ring)), radius=r_out, length=t_si, axis=2
    ) - td.Cylinder(center=(0, 0, np.mean(z_ring)), radius=r_in, length=2 * t_si, axis=2)

    corners = np.array(
        [
            (r_in - tooth_r, -tooth_a / 2),
            (r_in + 0.01, -tooth_a / 2),
            (r_in + 0.01, tooth_a / 2),
            (r_in - tooth_r, tooth_a / 2),
        ]
    )
    teeth = []
    for j in range(q):
        angle = 2 * np.pi * j / q
        rot = np.array(
            [[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]
        )
        teeth.append(
            td.PolySlab(vertices=corners @ rot.T, slab_bounds=z_ring, axis=2)
        )

    return td.Structure(
        geometry=td.GeometryGroup(geometries=[annulus, *teeth]), medium=si
    )


oxide = td.Structure(
    geometry=td.Box.from_bounds((-td.inf, -td.inf, -100), (td.inf, td.inf, z_ox_top)),
    medium=sio2,
)
bus = td.Structure(
    geometry=td.Box(center=(0, y_bus, np.mean(z_bus)), size=(td.inf, w_wg, t_si)),
    medium=si,
)

Lx = Ly = 17.4  # domain size in x and y (um)
Lz, z_center = 4.5, 1.05
z_nf = 1.6  # collection plane, ~0.6 um above the top oxide
mon_size = 15.0  # collection plane extent (+-7.5 um)
x_port = 6.0
port_size = (0, 2.2, 1.8)
run_time = 15e-12

def make_sim(q: int) -> td.Simulation:
    """Build the vertically coupled ring emitter with a q-element angular grating."""
    mode_source = td.ModeSource(
        center=(-x_port, y_bus, np.mean(z_bus)),
        size=port_size,
        source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth),
        direction="+",
        mode_spec=td.ModeSpec(num_modes=1, target_neff=2.5),
        mode_index=0,
    )
    thru_monitor = td.FluxMonitor(
        center=(x_port, y_bus, np.mean(z_bus)),
        size=port_size,
        freqs=freqs_wide,
        name="thru",
    )
    up_monitor = td.FluxMonitor(
        center=(0, 0, z_nf), size=(mon_size, mon_size, 0), freqs=freqs, name="up"
    )
    field_monitor = td.FieldMonitor(
        center=(0, 0, z_nf),
        size=(mon_size, mon_size, 0),
        freqs=freqs,
        fields=["Ex", "Ey", "Ez"],
        interval_space=(5, 5, 1),
        colocate=True,
        name="near",
    )
    far_monitor = td.FieldProjectionAngleMonitor(
        center=(0, 0, z_nf),
        size=(mon_size, mon_size, 0),
        freqs=freqs,
        theta=list(np.deg2rad(np.linspace(0, 25, 126))),
        phi=list(np.linspace(0, 2 * np.pi, 181)[:-1]),
        proj_distance=1e5,
        far_field_approx=True,
        name="far",
    )
    return td.Simulation(
        center=(0, 0, z_center),
        size=(Lx, Ly, Lz),
        grid_spec=td.GridSpec.auto(min_steps_per_wvl=30, wavelength=lda0),
        structures=[oxide, grating_ring(q), bus],
        sources=[mode_source],
        monitors=[thru_monitor, up_monitor, field_monitor, far_monitor],
        run_time=run_time,
        shutoff=1e-5,
        # The silicon bus crosses the x boundary, as a port must. A dispersive
        # medium inside a PML is numerically unstable, so x uses an Absorber;
        # the y and z boundaries see only oxide and air.
        boundary_spec=td.BoundarySpec(
            x=td.Boundary.absorber(), y=td.Boundary.pml(), z=td.Boundary.pml()
        ),
        medium=td.Medium(),  # air above the top oxide
    )


sims = {f"l{m_ring - q}": make_sim(q) for q in q_values}

fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(15, 4), tight_layout=True)
sim = sims["l1"]
sim.plot(z=np.mean(z_ring), ax=ax1)
ax1.set_title("Ring layer")
sim.plot(x=0, ax=ax2)
ax2.set_title("Vertical stack")
sim.plot(x=0, ax=ax3)
ax3.set_xlim(-5.2, -2.6)
ax3.set_ylim(-0.35, 1.25)
ax3.set_title("Vertical coupler")
plt.show()

def azimuthal_charge(field: np.ndarray, n_modes: int = 1) -> list:
    """Dominant exp(i m phi) harmonics of a far-field component, as (m, weight)."""
    coeff = np.fft.fftshift(np.fft.fft(field, axis=-1), axes=-1) / field.shape[-1]
    orders = np.fft.fftshift(
        np.fft.fftfreq(field.shape[-1], 1 / field.shape[-1])
    ).astype(int)
    weight = (np.abs(coeff) ** 2).sum(axis=0)
    best = np.argsort(weight)[::-1][:n_modes]
    return [(int(orders[i]), weight[i] / weight.sum()) for i in best]


results = {}
for q in q_values:
    ell = m_ring - q
    sim_data = batch_data[f"l{ell}"]
    # monitor data is stored on its own frequency coordinate (ascending), so we
    # read wavelengths from the data rather than from the input array
    f_band = sim_data["up"].flux.f.values
    lda_band = td.C_0 / f_band * 1e3
    up = np.abs(sim_data["up"].flux.values)

    peaks = [
        k
        for k in range(1, len(up) - 1)
        if up[k] >= up[k - 1] and up[k] >= up[k + 1] and up[k] > 0.15 * up.max()
    ]
    charges = []
    for k in peaks:
        e_theta = (
            sim_data["far"]
            .Etheta.sel(f=f_band[k])
            .squeeze(drop=True)
            .transpose("theta", "phi")
            .values
        )
        charges.append((k, azimuthal_charge(e_theta)[0][0]))

    print(f"q = {q}: " + ", ".join(
        f"{lda_band[k]:.0f} nm -> l = {c:+d} (m = {c + q})" for k, c in charges
    ))
    k_res = [k for k, c in charges if c == ell][0]
    results[ell] = dict(
        freq=float(f_band[k_res]), lda=float(lda_band[k_res]), eta=float(up[k_res]),
        data=sim_data
    )

for ell, res in results.items():
    near = res["data"]["near"]
    ex = near.Ex.sel(f=res["freq"]).squeeze(drop=True)
    ey = near.Ey.sel(f=res["freq"]).squeeze(drop=True)
    intensity = (np.abs(ex) ** 2 + np.abs(ey) ** 2).transpose("x", "y").values
    xx, yy = np.meshgrid(ex.x.values, ey.y.values, indexing="ij")
    radius = np.hypot(xx, yy)
    frac = [intensity[radius < r].sum() / intensity.sum() for r in (3, 5, 6, 7, 7.4)]
    print(f"l = {ell:+d}: enclosed intensity at r = 3/5/6/7/7.4 um -> "
          + ", ".join(f"{f:.1%}" for f in frac))

for ell, res in results.items():
    thru_data = res["data"]["thru"].flux
    lda_thru = td.C_0 / thru_data.f.values * 1e3
    thru = np.abs(thru_data.values)
    baseline = np.median(thru)
    window = np.abs(lda_thru - res["lda"]) < 6
    depth = 1 - thru[window] / baseline
    above = lda_thru[window][depth > depth.max() / 2]
    q_factor = res["lda"] / (above.max() - above.min())
    print(f"l = {ell:+d}: resonance {res['lda']:.1f} nm, Q = {q_factor:.0f}, "
          f"dip depth {depth.max():.0%}, upward emission {res['eta']:.1%}")

fig, axes = plt.subplots(2, 2, figsize=(11, 9), tight_layout=True)
for col, ell in enumerate(sorted(results)):
    res = results[ell]
    far = res["data"]["far"]
    theta = far.Etheta.theta.values
    phi = far.Ephi.phi.values
    e_theta = far.Etheta.sel(f=res["freq"]).squeeze(drop=True).transpose("theta", "phi").values
    e_phi = far.Ephi.sel(f=res["freq"]).squeeze(drop=True).transpose("theta", "phi").values
    intensity = np.abs(e_theta) ** 2 + np.abs(e_phi) ** 2

    # close the azimuthal grid so the polar mesh has no seam at phi = 0
    phi_c = np.append(phi, phi[0] + 2 * np.pi)
    intensity_c = np.concatenate([intensity, intensity[:, :1]], axis=1)
    tt, pp = np.meshgrid(np.degrees(theta), phi_c, indexing="ij")
    ax = axes[0, col]
    ax.pcolormesh(tt * np.cos(pp), tt * np.sin(pp), intensity_c / intensity_c.max(),
                  cmap="magma", shading="gouraud")
    ax.set_aspect("equal")
    ax.set_xlabel(r"$\theta\cos\phi$ (deg)")
    ax.set_ylabel(r"$\theta\sin\phi$ (deg)")
    ax.set_title(rf"$q={m_ring - ell}$, $\ell={ell:+d}$, $\lambda={res['lda']:.0f}$ nm")

    ax = axes[1, col]
    ax.plot(np.degrees(theta), intensity.mean(axis=1) / intensity.max(), lw=1.5)
    ax.axvline(np.degrees(np.arcsin(abs(ell) * lda0 / (2 * np.pi * R0))),
               color="C1", ls="--", lw=1, label=r"$\sin^{-1}(\ell / k_0 R_0)$")
    ax.set_xlabel(r"$\theta$ (deg)")
    ax.set_ylabel("Normalized intensity")
    ax.set_ylim(0, 1.05)
    ax.grid(alpha=0.3)
    ax.legend()
plt.show()

print(f"{'l':>4} {'measured l':>11} {'charge(Ex+iEy)':>15} {'charge(Ex-iEy)':>15} {'on-axis':>9}")
for ell in sorted(results):
    res = results[ell]
    far = res["data"]["far"]
    phi = far.Ephi.phi.values
    e_theta = far.Etheta.sel(f=res["freq"]).squeeze(drop=True).transpose("theta", "phi").values
    e_phi = far.Ephi.sel(f=res["freq"]).squeeze(drop=True).transpose("theta", "phi").values
    intensity = np.abs(e_theta) ** 2 + np.abs(e_phi) ** 2

    e_x = e_theta * np.cos(phi)[None, :] - e_phi * np.sin(phi)[None, :]
    e_y = e_theta * np.sin(phi)[None, :] + e_phi * np.cos(phi)[None, :]
    on_axis = intensity[0].mean() / intensity.max()
    print(f"{ell:>+4d} {azimuthal_charge(e_theta)[0][0]:>+11d} "
          f"{azimuthal_charge(e_x + 1j * e_y)[0][0]:>+15d} "
          f"{azimuthal_charge(e_x - 1j * e_y)[0][0]:>+15d} {on_axis:>9.1%}")