# Calibration constants (tuned so default geometry hits 50/50 at 1550 nm)
# pi * Delta_n * L / lambda = pi/4  =>  Delta_n = 0.25 * 1.55 / 4.7 ≈ 0.0825
_DELTA_N_REF = 0.0825  # supermode index split at reference geometry
_GAP_REF = 0.20  # um, reference gap
_GAP_DECAY = 1.85  # 1/um, exponential decay vs gap
_WIDTH_REF = 0.45  # um
_WIDTH_SLOPE = -0.55  # narrower waveguide -> more coupling
_THICKNESS_REF = 0.220  # um
_THICKNESS_SLOPE = -0.85  # thinner core -> mode less confined -> more coupling
_SIDEWALL_SLOPE = 0.0035  # per degree, asymmetry-induced loss
_DILATION_SLOPE = 1.20  # mask dilation widens core, equivalent to width change
_LAMBDA_DISP = -0.18  # 1/um, dn/dlambda (group index dispersion proxy)
_LAMBDA_REF = 1.55  # um


def supermode_split(
    *,
    core_width=0.45,
    gap=0.20,
    core_thickness=0.220,
    mask_dilation=0.0,
    wavelength=1.55,
):
    """Effective-index split (n_s - n_a) of the two supermodes."""
    eff_width = core_width + _DILATION_SLOPE * mask_dilation
    eff_gap = max(gap - _DILATION_SLOPE * mask_dilation, 1e-3)
    return (
        _DELTA_N_REF
        * np.exp(-_GAP_DECAY * (eff_gap - _GAP_REF))
        * (1.0 + _WIDTH_SLOPE * (eff_width - _WIDTH_REF))
        * (1.0 + _THICKNESS_SLOPE * (core_thickness - _THICKNESS_REF))
        * (1.0 + _LAMBDA_DISP * (np.asarray(wavelength) - _LAMBDA_REF))
    )


def analytical_s_matrix(
    *,
    coupling_length=4.7,
    core_width=0.45,
    gap=0.20,
    core_thickness=0.220,
    sidewall_angle=0.0,
    mask_dilation=0.0,
    excess_loss_db=0.05,
    freqs=None,
    wavelengths=None,
):
    """PhotonForge-style 2x2 directional coupler S-matrix dict.

    Ports follow PhotonForge's `s_bend_coupler` convention:
        P0 ---\\        /--- P2     through path: P0 -> P2
                \\======/             cross   path: P0 -> P3
                /======\\
        P1 ---/        \\--- P3
    """
    if wavelengths is None:
        wavelengths = (
            (C0 / np.asarray(freqs)) * 1e6
            if freqs is not None
            else np.linspace(1.50, 1.60, 51)
        )
    wl = np.atleast_1d(np.asarray(wavelengths, dtype=float))

    delta_n = supermode_split(
        core_width=core_width,
        gap=gap,
        core_thickness=core_thickness,
        mask_dilation=mask_dilation,
        wavelength=wl,
    )
    sidewall_penalty = np.clip(1.0 - _SIDEWALL_SLOPE * abs(sidewall_angle), 0.5, 1.0)
    kappa = np.clip(
        np.sin(np.pi * delta_n * coupling_length / wl) ** 2 * sidewall_penalty, 0.0, 1.0
    )

    loss_lin = 10 ** (-excess_loss_db / 20)
    t = np.sqrt(1.0 - kappa) * loss_lin  # through amplitude
    k = 1j * np.sqrt(kappa) * loss_lin  # cross amplitude (90° phase)

    S = {}
    S[("P0@0", "P2@0")] = t
    S[("P0@0", "P3@0")] = k  # P0 input
    S[("P1@0", "P3@0")] = t
    S[("P1@0", "P2@0")] = k  # P1 input (mirror)
    S[("P2@0", "P0@0")] = t
    S[("P3@0", "P0@0")] = k  # reciprocity
    S[("P3@0", "P1@0")] = t
    S[("P2@0", "P1@0")] = k
    for p in ("P0@0", "P1@0", "P2@0", "P3@0"):
        S[(p, p)] = np.full_like(wl, 1e-3, dtype=complex)
    return S


# Sanity check: nominal geometry should give ~50:50 at 1550 nm
S_nom = analytical_s_matrix(wavelengths=np.array([1.55]))
print(f"Nominal cross-coupling at 1550 nm: {abs(S_nom[('P0@0', 'P3@0')][0]) ** 2:.3f}")

freqs = C0 / np.linspace(1.50e-6, 1.60e-6, 41)  # 1500–1600 nm, 41 points

study = StudySpec(
    name="directional_coupler_50_50",
    freqs=freqs,
    foms=default_dc_foms(),
    params=[
        # --- Design knobs (we choose these) ---
        Param(
            "coupling_length",
            4.7,
            kind="design",
            value_range=(2.5, 7.5),
            units="um",
            label="coupling length",
        ),
        Param(
            "gap",
            0.20,
            kind="design",
            value_range=(0.15, 0.28),
            units="um",
            label="gap",
        ),
        # --- Process variations ---
        Param(
            "core_thickness", 0.220, stdev=0.0025, units="um", label="core thickness"
        ),
        Param("core_width", 0.45, stdev=0.006, units="um", label="core width"),
        Param("sidewall_angle", 0.0, stdev=2.0, units="deg", label="sidewall angle"),
        Param("mask_dilation", 0.0, stdev=0.005, units="um", label="mask dilation"),
    ],
)

print("Design knobs:")
for p in study.design_params:
    print(
        f"   {p.name:<18s} nominal={p.value:>6.3f} {p.units:<3s}  range={p.value_range}"
    )
print("\nProcess variations:")
for p in study.process_params:
    print(f"   {p.name:<18s} mean={p.value:>6.3f} {p.units:<3s}  1σ={p.stdev}")
print("\nFigures of merit:")
for f in study.foms:
    if f.spec == "window":
        print(f"   {f.name:<18s} {f.target} ± {f.tolerance} {f.units}")
    else:
        print(f"   {f.name:<18s} {f.spec:>3s} {f.target} {f.units}")

result = run_mc(
    study, analytical_s_matrix, n_samples=500, seed=42, keep_s_matrices=True
)
df = result.to_dataframe()
print(f"{len(df)} samples, {len(df.columns)} columns.")
df.describe().round(4)

@dataclass
class YieldResult:
    yield_total: float
    yield_per_fom: dict
    pass_mask: np.ndarray
    fom_pass_masks: dict


def compute_yield(result, foms):
    fom_masks = {
        f.name: np.asarray(f.passes(result.foms[f.name]))
        for f in foms
        if f.name in result.foms
    }
    if not fom_masks:
        return YieldResult(0.0, {}, np.array([], dtype=bool), {})
    total = np.ones(len(next(iter(fom_masks.values()))), dtype=bool)
    for m in fom_masks.values():
        total &= m
    return YieldResult(
        yield_total=float(total.mean()),
        yield_per_fom={k: float(v.mean()) for k, v in fom_masks.items()},
        pass_mask=total,
        fom_pass_masks=fom_masks,
    )


yr = compute_yield(result, study.foms)
print(f"Overall yield: {100 * yr.yield_total:.1f}%\n")
for name, y in yr.yield_per_fom.items():
    print(f"  {name:<18s} {100 * y:>5.1f}%")

def plot_fom_distributions(result, study, yr, bins=40):
    foms = study.foms
    n = len(foms)
    fig, axes = plt.subplots(1, n, figsize=(3.4 * n, 3.4))
    if n == 1:
        axes = [axes]
    for ax, fom in zip(axes, foms):
        y = result.foms[fom.name]
        passes = np.asarray(fom.passes(y))
        ax.hist(y[passes], bins=bins, color=C_PASS, alpha=0.85, label="pass")
        ax.hist(y[~passes], bins=bins, color=C_FAIL, alpha=0.85, label="fail")
        if fom.spec == "window":
            ax.axvspan(
                fom.target - fom.tolerance,
                fom.target + fom.tolerance,
                color="black",
                alpha=0.05,
            )
            ax.axvline(fom.target, color="black", ls="--", lw=1)
        else:
            ax.axvline(fom.target, color="black", ls="--", lw=1, label="spec")
        title = fom.label or fom.name
        if fom.name in yr.yield_per_fom:
            title += f"\nyield {100 * yr.yield_per_fom[fom.name]:.1f}%"
        ax.set_title(title, fontsize=10)
        ax.set_xlabel(f"{fom.name} [{fom.units}]" if fom.units else fom.name)
        ax.set_ylabel("count")
    axes[-1].legend(loc="upper right", fontsize=8, frameon=False)
    fig.tight_layout()
    return fig


plot_fom_distributions(result, study, yr)
plt.show()

def plot_correlation_heatmap(result, study):
    df = result.to_dataframe()
    var_cols = [c for c in df.columns if df[c].std() > 1e-12]
    corr = df[var_cols].corr()
    fig, ax = plt.subplots(figsize=(0.5 * len(var_cols) + 2, 0.5 * len(var_cols) + 1.5))
    im = ax.imshow(corr.values, cmap="RdBu_r", vmin=-1, vmax=1)
    ax.set_xticks(range(len(var_cols)))
    ax.set_yticks(range(len(var_cols)))
    ax.set_xticklabels(var_cols, rotation=45, ha="right", fontsize=8)
    ax.set_yticklabels(var_cols, fontsize=8)
    for i in range(len(var_cols)):
        for j in range(len(var_cols)):
            v = corr.values[i, j]
            ax.text(
                j,
                i,
                f"{v:.2f}",
                ha="center",
                va="center",
                color="white" if abs(v) > 0.5 else "black",
                fontsize=7,
            )
    fig.colorbar(im, ax=ax, label="Pearson r", shrink=0.7)
    fig.tight_layout()
    return fig


plot_correlation_heatmap(result, study)
plt.show()

def plot_scatter_matrix(result, study, fom_color=None):
    df = result.to_dataframe()
    var_params = [p for p in study.params if df[p.name].std() > 1e-12]
    P = len(var_params)
    if P < 2:
        fig, ax = plt.subplots(figsize=(5, 4))
        ax.text(0.5, 0.5, "Need >=2 varied params", ha="center")
        return fig

    fig, axes = plt.subplots(P, P, figsize=(2.0 * P, 2.0 * P))
    c = result.foms[fom_color] if fom_color else C_DESIGN
    cmap = "viridis" if fom_color else None

    for i, pi in enumerate(var_params):
        for j, pj in enumerate(var_params):
            ax = axes[i, j]
            if i == j:
                ax.hist(df[pi.name], bins=25, color=C_DESIGN, alpha=0.7)
                ax.set_yticks([])
            else:
                ax.scatter(
                    df[pj.name],
                    df[pi.name],
                    c=c,
                    cmap=cmap,
                    s=8,
                    alpha=0.7,
                    edgecolors="none",
                )
            label_i = f"{pi.label}\n[{pi.units}]" if pi.units else pi.label
            label_j = f"{pj.label}\n[{pj.units}]" if pj.units else pj.label
            if i == P - 1:
                ax.set_xlabel(label_j, fontsize=8)
            else:
                ax.set_xticklabels([])
            if j == 0:
                ax.set_ylabel(label_i, fontsize=8)
            else:
                ax.set_yticklabels([])
            ax.tick_params(labelsize=7)
    fig.suptitle("Parameter scatter matrix (coloured by coupling ratio)", fontsize=12)
    fig.tight_layout()
    return fig


plot_scatter_matrix(result, study, fom_color="coupling_ratio")
plt.show()

def plot_spectral_envelope(
    result, study, port_pair=("P0@0", "P3@0"), n_show=80, label="cross coupling"
):
    if result.s_matrices is None:
        fig, ax = plt.subplots(figsize=(5, 3))
        ax.text(0.5, 0.5, "Re-run with keep_s_matrices=True", ha="center")
        return fig
    lam_nm = (C0 / np.asarray(study.freqs)) * 1e9
    fig, ax = plt.subplots(figsize=(7, 4))
    for S in result.s_matrices[:n_show]:
        ax.plot(lam_nm, np.abs(S[port_pair]) ** 2, color=C_DESIGN, alpha=0.15, lw=0.8)
    mean = np.mean([np.abs(S[port_pair]) ** 2 for S in result.s_matrices], axis=0)
    ax.plot(lam_nm, mean, color=C_FAIL, lw=2.0, label="mean across MC")
    ax.axhspan(0.45, 0.55, color="black", alpha=0.05, label="±5% window")
    ax.set_xlabel("wavelength [nm]")
    ax.set_ylabel(f"|S{port_pair[0]} → {port_pair[1]}|²  ({label})")
    ax.set_title(f"Spectral envelope across {len(result.s_matrices)} MC samples")
    ax.legend(loc="best", fontsize=9)
    fig.tight_layout()
    return fig


plot_spectral_envelope(result, study)
plt.show()

from SALib.sample import saltelli
from SALib.analyze import sobol


def sobol_sensitivity(study, sim_func, n_base=128, seed=0):
    process = study.process_params
    if not process:
        return pd.DataFrame()

    # +/- 3 sigma bounds for Gaussian params
    problem = {
        "num_vars": len(process),
        "names": [p.name for p in process],
        "bounds": [
            list(p.value_range)
            if p.value_range is not None
            else [p.value - 3 * p.stdev, p.value + 3 * p.stdev]
            for p in process
        ],
    }

    rng_state = np.random.get_state()
    np.random.seed(seed)
    try:
        X = saltelli.sample(problem, n_base, calc_second_order=False)
    finally:
        np.random.set_state(rng_state)

    # Pin design params at nominal
    pinned = {p.name: p.value for p in study.design_params}

    Y = {f.name: np.zeros(len(X)) for f in study.foms}
    for i, row in enumerate(X):
        kwargs = dict(zip(problem["names"], row), **pinned)
        S = sim_func(**kwargs, freqs=study.freqs)
        for fom in study.foms:
            Y[fom.name][i] = fom.extract(S, study.freqs)

    rows = []
    for fom_name, y in Y.items():
        if np.std(y) < 1e-12:
            continue
        Si = sobol.analyze(problem, y, calc_second_order=False, print_to_console=False)
        for j, name in enumerate(problem["names"]):
            rows.append(
                {
                    "param": name,
                    "fom": fom_name,
                    "S1": float(Si["S1"][j]),
                    "ST": float(Si["ST"][j]),
                }
            )
    return pd.DataFrame(rows)


sob = sobol_sensitivity(study, analytical_s_matrix, n_base=128, seed=0)
print("Total-effect Sobol indices (rows=parameter, cols=FOM):")
sob.pivot(index="param", columns="fom", values="ST").round(3)

def plot_sobol_tornado(sens_df, fom, value_col="ST", title=None):
    sub = sens_df[sens_df["fom"] == fom].copy()
    if sub.empty:
        fig, ax = plt.subplots(figsize=(5, 3))
        ax.text(0.5, 0.5, f"No data for {fom!r}", ha="center")
        return fig
    sub = sub.sort_values(value_col, ascending=True)
    fig, ax = plt.subplots(figsize=(6, max(2.4, 0.4 * len(sub) + 1)))
    ax.barh(sub["param"], sub[value_col], color=C_DESIGN)
    for i, v in enumerate(sub[value_col]):
        ax.text(v, i, f" {v:.3f}", va="center", fontsize=9)
    ax.set_xlabel(f"{value_col} (Sobol sensitivity index)")
    ax.set_xlim(0, max(sub[value_col].max() * 1.15, 0.05))
    ax.set_title(title or f"Sensitivity of {fom}")
    fig.tight_layout()
    return fig


plot_sobol_tornado(
    sob, "coupling_ratio", title="Total-effect Sobol indices — coupling ratio"
)
plt.show()
plot_sobol_tornado(sob, "bw_3db", title="Total-effect Sobol indices — 50:50 bandwidth")
plt.show()

def yield_contour(study, sim_func, x_param, y_param, grid=15, n_inner=80, seed=0):
    px, py = study.param_by_name(x_param), study.param_by_name(y_param)
    xs = np.linspace(*px.value_range, grid)
    ys = np.linspace(*py.value_range, grid)
    yield_grid = np.zeros((grid, grid))

    for i, xv in enumerate(xs):
        for j, yv in enumerate(ys):
            overrides = {x_param: xv, y_param: yv}
            for p in study.design_params:
                if p.name not in overrides:
                    overrides[p.name] = p.value
            sub = run_mc(
                study,
                sim_func,
                n_samples=n_inner,
                seed=seed + i * grid + j,
                design_overrides=overrides,
            )
            yield_grid[j, i] = compute_yield(sub, study.foms).yield_total
    return {
        "x": xs,
        "y": ys,
        "x_param": x_param,
        "y_param": y_param,
        "yield_grid": yield_grid,
    }


# Compute the yield contour over (coupling_length, gap)
print(
    "Computing yield contour (18x18 grid, 80 inner samples each = 25920 evaluations)..."
)
contour = yield_contour(
    study,
    analytical_s_matrix,
    x_param="coupling_length",
    y_param="gap",
    grid=18,
    n_inner=80,
    seed=0,
)
print(
    f"Yield range: {contour['yield_grid'].min():.2f} to {contour['yield_grid'].max():.2f}"
)

def find_nominal_optimum(study, sim_func, objective_fom, target=None):
    fom = next(f for f in study.foms if f.name == objective_fom)
    if target is None:
        target = fom.target
    design = study.design_params
    bounds = [p.value_range for p in design]
    names = [p.name for p in design]

    def obj(x):
        kwargs = dict(zip(names, x))
        for p in study.process_params:
            kwargs[p.name] = p.value
        S = sim_func(**kwargs, freqs=study.freqs)
        return abs(fom.extract(S, study.freqs) - target)

    res = differential_evolution(
        obj, bounds=bounds, seed=0, maxiter=80, tol=1e-5, polish=True, popsize=15
    )
    return {
        "point": dict(zip(names, res.x)),
        "objective": float(res.fun),
        "success": res.success,
    }


nom = find_nominal_optimum(study, analytical_s_matrix, "coupling_ratio")
print("Nominal optimum (best 50:50 ignoring variation):")
for k, v in nom["point"].items():
    print(f"   {k} = {v:.4f}")
print(f"   |coupling - 0.5| at nominal = {nom['objective']:.5f}")

def find_robust_optimum(study, sim_func, n_inner=80, initial_grid=6, seed=0):
    design = study.design_params
    bounds = [p.value_range for p in design]
    names = [p.name for p in design]

    def neg_yield(x):
        for xv, (lo, hi) in zip(x, bounds):
            if xv < lo or xv > hi:
                return 1.0
        overrides = dict(zip(names, x))
        for p in study.design_params:
            if p.name not in overrides:
                overrides[p.name] = p.value
        sub = run_mc(
            study, sim_func, n_samples=n_inner, seed=seed, design_overrides=overrides
        )
        return -compute_yield(sub, study.foms).yield_total

    # Coarse grid to seed the local optimizer
    grids = [np.linspace(lo, hi, initial_grid) for lo, hi in bounds]
    mesh = np.array(np.meshgrid(*grids, indexing="ij")).reshape(len(grids), -1).T
    best_x, best_y = mesh[0], 1.0
    for x in mesh:
        y = neg_yield(x)
        if y < best_y:
            best_x, best_y = x, y

    # Local refinement (Nelder-Mead)
    res = minimize(
        neg_yield,
        best_x,
        method="Nelder-Mead",
        options={"xatol": 1e-3, "fatol": 1e-3, "maxiter": 60},
    )
    return {
        "point": dict(zip(names, res.x)),
        "yield": float(-res.fun),
        "success": res.success,
    }


print("Searching for robust optimum...")
robust = find_robust_optimum(
    study, analytical_s_matrix, n_inner=80, initial_grid=6, seed=0
)
print("\nRobust optimum (max yield under variation):")
for k, v in robust["point"].items():
    print(f"   {k} = {v:.4f}")
print(f"   predicted yield = {100 * robust['yield']:.1f}%")

def plot_yield_contour(contour, nominal_point=None, robust_point=None, study=None):
    fig, ax = plt.subplots(figsize=(7, 5.5))
    cs = ax.contourf(
        contour["x"],
        contour["y"],
        contour["yield_grid"],
        levels=np.linspace(0, 1, 11),
        cmap="viridis",
    )
    ax.contour(
        contour["x"],
        contour["y"],
        contour["yield_grid"],
        levels=[0.5, 0.8, 0.95],
        colors="white",
        linewidths=1.0,
        linestyles=["--", "-.", "-"],
    )
    fig.colorbar(cs, ax=ax, label="Yield (fraction passing all specs)")
    xn, yn = contour["x_param"], contour["y_param"]
    if nominal_point is not None:
        ax.plot(
            nominal_point[xn],
            nominal_point[yn],
            marker="*",
            ms=20,
            color=C_NOMINAL,
            mec="black",
            mew=1.0,
            label="Nominal optimum",
        )
    if robust_point is not None:
        ax.plot(
            robust_point[xn],
            robust_point[yn],
            marker="P",
            ms=15,
            color=C_ROBUST,
            mec="black",
            mew=1.0,
            label="Robust optimum",
        )
    if study is not None:
        px, py = study.param_by_name(xn), study.param_by_name(yn)
        ax.set_xlabel(f"{px.label} [{px.units}]")
        ax.set_ylabel(f"{py.label} [{py.units}]")
    else:
        ax.set_xlabel(xn)
        ax.set_ylabel(yn)
    ax.set_title("Fabrication yield over design space")
    ax.legend(loc="best", framealpha=0.9, fontsize=10)
    fig.tight_layout()
    return fig


plot_yield_contour(
    contour, nominal_point=nom["point"], robust_point=robust["point"], study=study
)
plt.show()

nom_result = run_mc(
    study, analytical_s_matrix, n_samples=1000, seed=99, design_overrides=nom["point"]
)
nom_yield = compute_yield(nom_result, study.foms)

rob_result = run_mc(
    study,
    analytical_s_matrix,
    n_samples=1000,
    seed=99,
    design_overrides=robust["point"],
)
rob_yield = compute_yield(rob_result, study.foms)

summary = pd.DataFrame(
    {
        "design point": ["nominal optimum", "robust optimum"],
        "coupling_length [um]": [
            nom["point"]["coupling_length"],
            robust["point"]["coupling_length"],
        ],
        "gap [um]": [nom["point"]["gap"], robust["point"]["gap"]],
        "overall yield [%]": [100 * nom_yield.yield_total, 100 * rob_yield.yield_total],
    }
)
for f in study.foms:
    summary[f"{f.name} yield [%]"] = [
        100 * nom_yield.yield_per_fom.get(f.name, 0),
        100 * rob_yield.yield_per_fom.get(f.name, 0),
    ]
summary.round(2)

fig, axes = plt.subplots(1, 2, figsize=(11, 3.6), sharey=True)
for ax, (label, yr_obj) in zip(
    axes, [("Nominal optimum", nom_yield), ("Robust optimum", rob_yield)]
):
    names = list(yr_obj.yield_per_fom.keys()) + ["ALL"]
    vals = [yr_obj.yield_per_fom[n] for n in yr_obj.yield_per_fom] + [
        yr_obj.yield_total
    ]
    colors = [C_DESIGN] * (len(names) - 1) + [C_PASS]
    ax.barh(names, vals, color=colors)
    ax.axvline(1.0, color="grey", lw=0.5)
    ax.set_xlim(0, 1.05)
    ax.set_xlabel("Pass rate")
    ax.set_title(f"{label}: {100 * yr_obj.yield_total:.1f}% overall")
    for i, v in enumerate(vals):
        ax.text(v, i, f" {100 * v:.1f}%", va="center", fontsize=9)
fig.suptitle("Side-by-side validation", fontsize=12)
fig.tight_layout()
plt.show()

import photonforge as pf
import siepic_forge as siepic

if "Basic" in str(pf.config.default_technology):
    pf.config.default_technology = siepic.ebeam()
_tech = pf.config.default_technology
assert "TE_1550_500" in _tech.ports

C0 = 299_792_458.0
freqs_prod = C0 / np.linspace(1.540e-6, 1.560e-6, 5)

robust_design = {"coupling_length": 4.54, "gap": 0.178}


def make_dc_fdtd(coupling_length, gap, min_steps_per_wvl=10):
    """Build the FDTD coupler with PhotonForge's s_bend_coupler.

    Important: PhotonForge's `coupling_distance` is the CENTER-TO-CENTER
    separation between the two coupling waveguides, NOT the edge-to-edge
    gap. The offline analytical study uses `gap` to mean the edge-to-edge
    separation (the standard photonics convention). This wrapper converts:

        coupling_distance = gap (edge-to-edge) + waveguide_width

    so the physical device matches the analytical-model assumption.
    """
    spec = pf.config.default_technology.ports["TE_1550_500"]
    try:
        wg_width, _ = spec.path_profile_for("Si")
    except Exception:
        wg_width = 0.5
    coupling_distance = gap + wg_width
    model = pf.Tidy3DModel(grid_spec=min_steps_per_wvl)
    return pf.parametric.s_bend_coupler(
        port_spec="TE_1550_500",
        coupling_length=coupling_length,
        coupling_distance=coupling_distance,
        s_bend_length=10,
        s_bend_offset=3,
        model=model,
    )


class _SAdapter:
    def __init__(self, S):
        self._S = S

    def __getitem__(self, k):
        try:
            return np.asarray(self._S[k])
        except (KeyError, TypeError):
            return np.asarray(self._S.elements[k])


def run_fdtd_point(L_c, gap, min_steps_per_wvl=10):
    comp = make_dc_fdtd(L_c, gap, min_steps_per_wvl=min_steps_per_wvl)
    smat = comp.s_matrix(freqs_prod)
    S = _SAdapter(smat)

    mid = len(smat.frequencies) // 2
    ports = ["P0@0", "P1@0", "P2@0", "P3@0"]
    p_in = sum(
        abs(smat.elements[("P0@0", po)][mid]) ** 2
        for po in ports
        if ("P0@0", po) in smat.elements
    )

    foms = {f.name: float(f.extract(S, freqs_prod)) for f in study.foms}
    return {
        "L_c": L_c,
        "gap": gap,
        "min_steps": min_steps_per_wvl,
        "smatrix": smat,
        "power_in": p_in,
        **foms,
    }

import gdstk
import os
import tempfile
from matplotlib.patches import Polygon as MplPolygon

robust_component = make_dc_fdtd(
    coupling_length=robust_design["coupling_length"],
    gap=robust_design["gap"],
    min_steps_per_wvl=10,
)

bbox = robust_component.bounds()
size_x = float(bbox[1][0] - bbox[0][0])
size_y = float(bbox[1][1] - bbox[0][1])

L_c_print = robust_design["coupling_length"]
gap_print = robust_design["gap"]
print(f"Robust design point:  L_c = {L_c_print:.3f} um, gap = {gap_print:.3f} um")
print("Bend topology:        s_bend_length = 10 um, s_bend_offset = 3 um")
print(f"Component bounds:     {size_x:.2f} um x {size_y:.2f} um")
print(f"Ports:                {list(robust_component.ports.keys())}")

_gds_tmp = os.path.join(tempfile.gettempdir(), "robust_dc_layout.gds")
robust_component.write_gds(_gds_tmp)

if not os.path.exists(_gds_tmp):
    fallback = "robust_dc_layout.gds"
    if os.path.exists(fallback):
        _gds_tmp = os.path.abspath(fallback)
    else:
        for f in os.listdir("."):
            if f.endswith(".gds") and "robust_dc_layout" in f:
                _gds_tmp = os.path.abspath(f)
                break
print(f"GDS written to: {_gds_tmp}")

_lib = gdstk.read_gds(_gds_tmp)
_top = _lib.top_level()[0]
_polys = _top.get_polygons(depth=None)

_layer_colors = {
    (1, 0): (C_DESIGN, 0.85, "Si waveguide (1/0)"),
    (1, 5): (C_NOMINAL, 0.45, "Si etch (1/5)"),
    (1, 10): (C_FAIL, 0.35, "DevRec (1/10)"),
    (68, 0): ("#888888", 0.35, "Floorplan (68/0)"),
}


def _draw_polys(ax):
    seen = set()
    for poly in _polys:
        layer_key = (poly.layer, poly.datatype)
        color, alpha, label = _layer_colors.get(
            layer_key, ("#cccccc", 0.30, f"layer {layer_key}")
        )
        legend_label = label if layer_key not in seen else None
        seen.add(layer_key)
        patch = MplPolygon(
            poly.points,
            closed=True,
            facecolor=color,
            edgecolor="#222222",
            linewidth=0.4,
            alpha=alpha,
            label=legend_label,
        )
        ax.add_patch(patch)
    return seen


fig = plt.figure(figsize=(13, 5.2))
gs = fig.add_gridspec(1, 2, width_ratios=[2.3, 1.0], wspace=0.25)
ax_full = fig.add_subplot(gs[0, 0])
ax_zoom = fig.add_subplot(gs[0, 1])

_draw_polys(ax_full)
for pname, port in robust_component.ports.items():
    cx_p, cy_p = float(port.center[0]), float(port.center[1])
    ax_full.plot(
        cx_p,
        cy_p,
        "o",
        color=C_PASS,
        ms=9,
        markeredgecolor="black",
        markeredgewidth=0.7,
        zorder=5,
    )
    ax_full.annotate(
        pname,
        (cx_p, cy_p),
        xytext=(8, 8),
        textcoords="offset points",
        fontsize=10,
        fontweight="bold",
        zorder=6,
    )
ax_full.set_xlim(bbox[0][0] - 1.5, bbox[1][0] + 1.5)
ax_full.set_ylim(bbox[0][1] - 1.5, bbox[1][1] + 1.5)
ax_full.set_aspect("equal")
ax_full.set_xlabel("x (um)")
ax_full.set_ylabel("y (um)")
ax_full.set_title(
    f"Robust directional coupler  (L_c = {L_c_print:.2f} um, gap = {gap_print:.3f} um)"
)
ax_full.grid(alpha=0.25)
ax_full.legend(loc="upper right", fontsize=8.5, framealpha=0.92)

cx = (bbox[0][0] + bbox[1][0]) / 2
cy = (bbox[0][1] + bbox[1][1]) / 2
zoom_half_x = robust_design["coupling_length"] * 0.55
zoom_half_y = 0.7

_draw_polys(ax_zoom)
ax_zoom.set_xlim(cx - zoom_half_x, cx + zoom_half_x)
ax_zoom.set_ylim(cy - zoom_half_y, cy + zoom_half_y)
ax_zoom.set_aspect("equal")
ax_zoom.set_xlabel("x (um)")
ax_zoom.set_ylabel("y (um)")
ax_zoom.set_title("Coupling region (zoomed in y)")
ax_zoom.grid(alpha=0.25)

gap = robust_design["gap"]
ax_zoom.axhline(cy + gap / 2, color=C_FAIL, lw=1.0, ls="--", alpha=0.85)
ax_zoom.axhline(cy - gap / 2, color=C_FAIL, lw=1.0, ls="--", alpha=0.85)
arrow_x = cx - zoom_half_x * 0.55
ax_zoom.annotate(
    "",
    xy=(arrow_x, cy + gap / 2),
    xytext=(arrow_x, cy - gap / 2),
    arrowprops=dict(arrowstyle="<->", color=C_FAIL, lw=1.6),
)
ax_zoom.text(
    arrow_x + 0.18,
    cy,
    f"gap = {gap * 1000:.0f} nm",
    fontsize=11,
    fontweight="bold",
    color=C_FAIL,
    ha="left",
    va="center",
    bbox=dict(boxstyle="round,pad=0.2", fc="white", ec=C_FAIL, lw=0.8, alpha=0.92),
)

fig.tight_layout()
plt.show()

try:
    os.remove(_gds_tmp)
except OSError:
    pass