def layer(zc: float, t: float, n: float, w: float = td.inf) -> td.Structure:
    """Return a layer of thickness t and width w, centered at height zc."""
    return td.Structure(geometry=td.Box(center=(0, 0, zc), size=(td.inf, w, t)),
                        medium=td.Medium(permittivity=n**2))


def te0_neff(structures: list, target: float) -> float:
    """Solve the TE0 effective index of a cross-section locally."""
    sim = td.Simulation(
        size=(1, 6, 5), center=(0, 0, -0.6), structures=structures, run_time=1e-13,
        grid_spec=td.GridSpec.auto(min_steps_per_wvl=40, wavelength=lda0),
        boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()),
    )
    modes = td.ModeSimulation.from_simulation(
        simulation=sim, plane=td.Box(center=(0, 0, -0.6), size=(0, 6, 5)),
        mode_spec=td.ModeSpec(num_modes=1, target_neff=target), freqs=[freq0],
    ).run_local().modes
    return float(np.real(modes.n_eff.values.ravel()[0]))


box, buf = layer(-h_sin - h_box / 2, h_box, n_box), layer(h_buf / 2, h_buf, n_buf)
n_slab = te0_neff([box, layer(-h_sin / 2, h_sin, n_sin), buf], n_sin)
n_tooth = te0_neff([box, layer(-h_sin / 2, h_sin, n_sin), buf,
                    layer(h_buf + h_asi / 2, h_asi, n_asi)], n_asi)
n_wg = te0_neff([box, layer(-h_sin / 2, h_sin, n_sin, w_wg), buf], n_sin)
print(f"bare SiN slab      n_eff = {n_slab:.4f}")
print(f"alpha-Si loaded    n_eff = {n_tooth:.4f}")
print(f"feed waveguide TE0 n_eff = {n_wg:.4f}")

def emt_index(w: float) -> float:
    """Second-order Rytov index for E perpendicular to the subwavelength interfaces."""
    f = w / swg_period
    par = f * n_asi**2 + (1 - f)
    perp = 1 / (f / n_asi**2 + (1 - f))
    corr = (
        (np.pi**2 / 3)
        * (swg_period / lda0) ** 2
        * f**2
        * (1 - f) ** 2
        * (1 - n_asi**2) ** 2
        * par
        * (perp / n_asi**2) ** 2
    )
    return np.sqrt(perp * (1 + corr))


# emt_index is non-monotonic: it peaks then falls back, so the inversion is only
# single-valued on the rising branch and brentq must be bracketed there.
_scan = np.linspace(1e-4, swg_period, 2001)
w_peak = _scan[np.argmax([emt_index(w) for w in _scan])]


def emt_width(n_target: float) -> float:
    """Invert the effective medium theory on its rising branch."""
    if n_target >= emt_index(swg_period) - 1e-9:
        return swg_period
    return brentq(lambda w: emt_index(w) - n_target, 1e-4, w_peak)


widths = np.linspace(0.01, swg_period, 200)
plt.plot(widths / swg_period, [emt_index(w) for w in widths], c="crimson", label="second order")
plt.plot(widths / swg_period, [np.sqrt(1 / (w / swg_period / n_asi**2 + 1 - w / swg_period)) for w in widths],
         "--", c="gray", label="zeroth order")
plt.axhline(n_asi, c="k", lw=0.8, ls=":")
plt.xlabel("Duty cycle $w/\\Lambda$")
plt.ylabel("Synthesized index")
plt.legend()
plt.grid(alpha=0.3)
plt.show()

pitch = np.concatenate([np.linspace(0.665, 0.639, n_apod), np.linspace(0.638, 0.603, n_homo)])
n_swg = np.concatenate([np.linspace(swg_n_lo, n_asi, n_apod), np.full(n_homo, n_asi)])
block_w = np.array([emt_width(n) for n in n_swg])
x_start = x_g0 + np.concatenate([[0], np.cumsum(pitch)])[:-1]
tooth_len = duty * pitch

print(f"grating length  = {np.sum(pitch):.2f} um")
print(f"gaps close from {(swg_period - block_w[0]) * 1e3:.0f} nm to {(swg_period - block_w[n_apod - 2]) * 1e3:.0f} nm")

n_g, focus = 1.81, (-10.1, 110)  # in-plane index and equal-path point (um)
phi = np.linspace(-fan_half, fan_half, 481)


def grating_line(x_axis: float, ph: np.ndarray = None) -> tuple:
    """Return the (x, y) curve whose optical path from tip to focus is constant."""
    ph = phi if ph is None else ph
    xf, zf = focus
    const = n_g * (x_axis - x_tip) + np.hypot(x_axis - xf, zf)

    def residual(r, c, s):
        return n_g * r + np.sqrt((x_tip + r * c - xf) ** 2 + zf**2 + (r * s) ** 2) - const

    r = np.array([brentq(residual, 1e-3, 500, args=(np.cos(a), np.sin(a))) for a in ph])
    return x_tip + r * np.cos(ph), r * np.sin(ph)


def tooth(i: int) -> list:
    """Return the polygon(s) of grating period i, segmented if the period is apodized."""
    lead, trail = x_start[i], x_start[i] + tooth_len[i]
    xl, yl = grating_line(lead)
    if block_w[i] >= swg_period:  # uniform period: one solid arc
        xt, yt = grating_line(trail)
        return [list(zip(xl, yl)) + list(zip(xt[::-1], yt[::-1]))]

    s = np.concatenate([[0], np.cumsum(np.hypot(np.diff(xl), np.diff(yl)))])
    s -= np.interp(0, yl, s)  # arc length measured from the symmetry plane
    polys = []
    for m in range(int(s[0] / swg_period) - 1, int(s[-1] / swg_period) + 2):
        a = max(m * swg_period - block_w[i] / 2, s[0])
        b = min(m * swg_period + block_w[i] / 2, s[-1])
        if b - a <= min_feature:  # drop sub-resolution slivers at the sector edge
            continue
        ph = np.concatenate([[np.interp(a, s, phi)], phi[(s > a) & (s < b)], [np.interp(b, s, phi)]])
        x1, y1 = grating_line(lead, ph)
        x2, y2 = grating_line(trail, ph)
        polys.append(list(zip(x1, y1)) + list(zip(x2[::-1], y2[::-1])))
    return polys


teeth = [poly for i in range(len(pitch)) for poly in tooth(i)]
print(f"{len(teeth)} alpha-Si polygons over {len(pitch)} periods")

for poly in teeth:
    v = np.array(poly)
    plt.fill(v[:, 0], v[:, 1], c="teal")
plt.xlabel("x ($\\mu m$)")
plt.ylabel("y ($\\mu m$)")
plt.gca().set_aspect("equal")
plt.show()

fields = sim_data["fiber_plane"].fields_cartesian
xg, yg = np.array(fields.coords["x"]), np.array(fields.coords["y"])
x_org = x_c  # projection coordinates are relative to the monitor center
dA = (xg[1] - xg[0]) * (yg[1] - yg[0])
i_freq = [int(np.argmin(np.abs(np.array(fields.coords["f"]) - f))) for f in freqs]

radiated = lambda i: [np.array(fields[k].isel(f=i_freq[i])).squeeze() for k in ("Ex", "Ey", "Hx", "Hy")]


def fiber_mode(lda: float, tilt: float) -> list:
    """Transverse fields of a tilted SMF-28 mode, sampled on the projection grid."""
    beam = td.GaussianBeamProfile(
        center=(x_org, 0, fiber_h),
        size=(2.4 * np.max(np.abs(xg)), 2.4 * np.max(np.abs(yg)), 0),
        resolution=600, freqs=np.array([td.C_0 / lda]),
        angle_theta=np.deg2rad(tilt), angle_phi=np.pi, pol_angle=np.pi / 2,
        direction="+", waist_radius=mfd / 2,
    ).field_data
    return [beam.field_components[k].interp(x=xg + x_org, y=yg).squeeze().transpose("x", "y").values
            for k in ("Ex", "Ey", "Hx", "Hy")]


def coupling_map(i: int, tilt: float) -> np.ndarray:
    """Coupling efficiency at every lateral fiber offset, as an FFT cross-correlation.

    The launched waveguide mode and the Gaussian beam both carry unit power, so the
    squared overlap integral is already the fraction of input power that couples.
    """
    e1x, e1y, h1x, h1y = radiated(i)
    e2x, e2y, h2x, h2y = fiber_mode(ldas[i], tilt)
    corr = lambda a, b: np.fft.ifft2(np.fft.fft2(a) * np.conj(np.fft.fft2(b)))
    num = corr(e1x, h2y) - corr(e1y, h2x) + corr(h1y, e2x) - corr(h1x, e2y)
    return np.abs(num * dA / 4) ** 2

def unwrap(i: int, n: int) -> int:
    """Convert an FFT correlation index into a signed grid shift."""
    return i if i <= n // 2 else i - n


eta0, tilt0, ix, iy = 0, 0, 0, 0
for tilt in np.arange(11, 17, 0.25):
    m = coupling_map(i0, tilt)
    j, k = np.unravel_index(np.argmax(m), m.shape)
    if m[j, k] > eta0:
        eta0, tilt0, ix, iy = m[j, k], tilt, j, k

x_fiber = x_org + unwrap(ix, len(xg)) * (xg[1] - xg[0])
y_fiber = unwrap(iy, len(yg)) * (yg[1] - yg[0])
print(f"optimum fiber tilt = {tilt0:.2f} deg")
print(f"lateral position   = {x_fiber:.2f} um")
print(f"transverse offset  = {y_fiber:.2f} um (zero, as the symmetry requires)")

ce_db = 10 * np.log10([coupling_map(i, tilt0)[ix, iy] for i in range(len(ldas))])
p_proj = np.array(sim_data["fiber_plane"].flux.sel(f=freqs, method="nearest")).ravel()
print(f"mode overlap at {lda0} um = {10 ** (ce_db[i0] / 10) / p_proj[i0] * 100:.1f} %")

plt.plot(ldas, ce_db, "o-", c="crimson")
plt.xlabel("Wavelength ($\\mu m$)")
plt.ylabel("Coupling efficiency (dB)")
plt.ylim(-10, 0)
plt.grid(alpha=0.3)
plt.show()

peak = ce_db.max()
fine = np.linspace(ldas[0], ldas[-1], 4001)
inside = np.interp(fine, ldas, ce_db) >= peak - 1
print(f"window captures {p_proj[i0] / up[i0] * 100:.1f} % of the upward power")
print(f"peak coupling efficiency = {peak:.2f} dB at {ldas[np.argmax(ce_db)] * 1e3:.0f} nm")
print(f"1 dB bandwidth = {(fine[inside].max() - fine[inside].min()) * 1e3:.0f} nm "
      f"(+/- 2 nm, limited by the 10 nm wavelength sampling)")

xz = sim_data["cross_section"].fields_cartesian
xs, zs = np.array(xz.coords["x"]), np.array(xz.coords["z"])
intensity = sum(np.abs(np.array(xz[k]).squeeze()) ** 2 for k in ("Ex", "Ey", "Ez"))

plt.figure(figsize=(6, 5))
plt.pcolormesh(xs + x_c, zs + z_mon, np.sqrt(intensity / intensity.max()).T, cmap="magma", shading="auto")
plt.colorbar(label="$|E|$ (a.u.)")
plt.gca().add_patch(plt.Circle((x_fiber, fiber_h), mfd / 2, fill=False, ec="w", lw=1.5))
plt.xlabel("x ($\\mu m$)")
plt.ylabel("z ($\\mu m$)")
plt.gca().set_aspect("equal")
plt.tight_layout()
plt.show()