wavelength = 1.55
freq = td.C_0 / wavelength

core_thickness = 0.22
width = 0.55
slot_widths = np.round(np.arange(0.04, 0.161, 0.02), 3)

n_core = 3.48
n_subs = 1.44
n_clad = 1.00
n_slot = n_clad

num_modes = 4
te_threshold = 0.5

sim_size = (5.0, 5.0, 1.0)
mode_plane_size = (5.0, 5.0, 0.0)

TE_CMAP = "jet"

clad = td.Medium(permittivity=n_clad**2)
substrate_medium = td.Medium(permittivity=n_subs**2)
core_medium = td.Medium(permittivity=n_core**2)
slot_medium = td.Medium(permittivity=n_slot**2)

def make_slot_waveguide_mode_sim(slot_width, mode_spec):
    rail_width = (width - slot_width) / 2

    structures = [
        td.Structure(
            geometry=td.Box(center=(0, -0.5, 0), size=(td.inf, 1.0, td.inf)),
            medium=substrate_medium,
            name="substrate",
        ),
        td.Structure(
            geometry=td.Box(
                center=(-(slot_width / 2 + rail_width / 2), core_thickness / 2, 0),
                size=(rail_width, core_thickness, td.inf),
            ),
            medium=core_medium,
            name="left_rail",
        ),
        td.Structure(
            geometry=td.Box(
                center=((slot_width / 2 + rail_width / 2), core_thickness / 2, 0),
                size=(rail_width, core_thickness, td.inf),
            ),
            medium=core_medium,
            name="right_rail",
        ),
        td.Structure(
            geometry=td.Box(
                center=(0, core_thickness / 2, 0),
                size=(slot_width, core_thickness, td.inf),
            ),
            medium=slot_medium,
            name="slot_region",
        ),
    ]

    return td.ModeSimulation(
        size=sim_size,
        medium=clad,
        structures=structures,
        boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()),
        grid_spec=td.GridSpec.auto(wavelength=wavelength, min_steps_per_wvl=40),
        mode_spec=mode_spec,
        freqs=[freq],
        plane=td.Box(center=(0, core_thickness / 2, 0), size=mode_plane_size),
    )


def make_simulations(mode_spec):
    simulations = {}

    for slot_width in slot_widths:
        name = f"slot={slot_width:.3f}"
        simulations[name] = make_slot_waveguide_mode_sim(slot_width, mode_spec)

    return simulations

def mode_dataframe(results):
    mode_tables = []

    for slot_width in slot_widths:
        name = f"slot={slot_width:.3f}"
        sim_data = results[name]

        mode_table = sim_data.modes.to_dataframe().reset_index()
        mode_table = mode_table.rename(
            columns={"n eff": "n_eff", "TE (Ex) fraction": "TE_fraction"}
        )
        mode_table["key"] = name
        mode_table["slot_width_um"] = slot_width

        mode_tables.append(
            mode_table[["key", "slot_width_um", "mode_index", "n_eff", "TE_fraction"]]
        )

    return (
        pd.concat(mode_tables, ignore_index=True)
        .sort_values(["slot_width_um", "mode_index"])
        .reset_index(drop=True)
    )


def add_te_fraction_line(ax, df, label, linewidth=2.0):
    df = df.sort_values("slot_width_um")
    x = df["slot_width_um"].to_numpy()
    y = df["n_eff"].to_numpy()
    c = df["TE_fraction"].to_numpy()

    ax.plot(x, y, color="black", linestyle="--", linewidth=linewidth, alpha=0.55)
    points = ax.scatter(x, y, c=c, cmap=TE_CMAP, vmin=0, vmax=1)

    ax.text(x[-1] + 0.002, y[-1], label, fontsize=9, va="center")
    return points


def plot_mode_paths(df, group_column, title, label_prefix="mode "):
    fig, ax = plt.subplots(figsize=(7.4, 4.8))

    mappable = None
    for group_value, group_df in df.groupby(group_column, sort=True):
        label = f"{label_prefix}{group_value}" if label_prefix else str(group_value)
        mappable = add_te_fraction_line(ax, group_df, label=label)

    ax.autoscale()
    ax.margins(x=0.12, y=0.08)
    ax.set_xlabel("Slot width (µm)")
    ax.set_ylabel(r"$n_\mathrm{eff}$")
    ax.set_title(title)
    ax.grid(True, alpha=0.22)

    cbar = fig.colorbar(mappable, ax=ax)
    cbar.set_label("TE fraction")

    plt.show()


def plot_mode_fields(results, modes_df, mode_indices, slot_values, title):
    fig, axes = plt.subplots(
        len(mode_indices),
        len(slot_values),
        figsize=(4.2 * len(slot_values), 3.4 * len(mode_indices)),
        tight_layout=True,
    )
    axes = np.atleast_2d(axes)

    for row_index, mode_index in enumerate(mode_indices):
        for col_index, slot_width in enumerate(slot_values):
            ax = axes[row_index, col_index]
            row = modes_df[
                (modes_df["slot_width_um"] == slot_width) & (modes_df["mode_index"] == mode_index)
            ].iloc[0]

            sim_data = results[row["key"]]
            sim_data.plot_field(
                "E",
                "abs",
                mode_index=int(row["mode_index"]),
                f=freq,
                ax=ax,
                robust=False,
            )

            ax.set_title(
                f"mode {mode_index}, slot={slot_width:.2f} µm\n"
                rf"$n_\mathrm{{eff}}$={row['n_eff']:.4f}, TE={row['TE_fraction']:.3f}",
                fontsize=9,
            )
            ax.set_xlim(-0.8, 0.8)
            ax.set_ylim(-0.5, 0.7)

    fig.suptitle(title, y=1.02, fontsize=12)
    plt.show()


def plot_branch_fields(results_by_branch, branches, branch_names, slot_values, title):
    fig, axes = plt.subplots(
        len(branch_names),
        len(slot_values),
        figsize=(4.2 * len(slot_values), 3.4 * len(branch_names)),
        tight_layout=True,
    )
    axes = np.atleast_2d(axes)

    for row_index, branch_name in enumerate(branch_names):
        branch = branches[branches["branch"] == branch_name].set_index("slot_width_um")
        results = results_by_branch[branch_name]

        for col_index, slot_width in enumerate(slot_values):
            ax = axes[row_index, col_index]
            row = branch.loc[slot_width]

            sim_data = results[row["key"]]
            mode_index = int(row["mode_index"])
            sim_data.plot_field("E", "abs", mode_index=mode_index, f=freq, ax=ax, robust=False)

            ax.set_title(
                f"{branch_name}, slot={slot_width:.2f} µm\n"
                rf"mode {mode_index}, $n_\mathrm{{eff}}$={row['n_eff']:.4f}",
                fontsize=9,
            )
            ax.set_xlim(-0.8, 0.8)
            ax.set_ylim(-0.5, 0.7)

    fig.suptitle(title, y=1.02, fontsize=12)
    plt.show()

def make_sort_spec(slot_width, filter_order="over"):
    slot_box = td.Box(
        center=(0, core_thickness / 2, 0),
        size=(slot_width, core_thickness, 0),
    )

    return td.ModeSortSpec(
        filter_key="TE_fraction",
        filter_reference=te_threshold,
        filter_order=filter_order,
        sort_key="fill_fraction_box",
        bounding_box=slot_box,
    )


sorted_results = {}

for slot_width in slot_widths:
    name = f"slot={slot_width:.3f}"
    sorted_results[name] = naive_results[name].sort_modes(make_sort_spec(slot_width))

sorted_modes_df = mode_dataframe(sorted_results)

plot_mode_paths(
    sorted_modes_df,
    group_column="mode_index",
    title="Sorted result: TE-like modes are placed first",
)

first_two_modes_df = sorted_modes_df[sorted_modes_df["mode_index"].isin([0, 1])]

plot_mode_paths(
    first_two_modes_df,
    group_column="mode_index",
    title="First two sorted modes: mode_index=1 is still not TM-like",
)

plot_mode_fields(
    sorted_results,
    sorted_modes_df,
    mode_indices=[0, 1],
    slot_values=[slot_widths[0], slot_widths[-1]],
    title="Sorted field check: mode_index=0 and mode_index=1 are both TE-like here",
)

tm_sorted_results = {
    f"slot={slot_width:.3f}": naive_results[f"slot={slot_width:.3f}"].sort_modes(
        make_sort_spec(slot_width, filter_order="under")
    )
    for slot_width in slot_widths
}

te_branch_df = sorted_modes_df[sorted_modes_df["mode_index"] == 0].assign(branch="TE-like")

tm_branch_df = mode_dataframe(tm_sorted_results)
tm_branch_df = tm_branch_df[tm_branch_df["mode_index"] == 0].assign(branch="TM-like")

branches_df = (
    pd.concat([te_branch_df, tm_branch_df], ignore_index=True)
    .sort_values(["branch", "slot_width_um"])
    .reset_index(drop=True)
)

plot_mode_paths(
    branches_df,
    group_column="branch",
    title="Physical branches selected from the same sorted run",
    label_prefix="",
)

plot_branch_fields(
    {"TE-like": sorted_results, "TM-like": tm_sorted_results},
    branches_df,
    branch_names=["TE-like", "TM-like"],
    slot_values=[slot_widths[0], slot_widths[-1]],
    title="Final field check: selected TE-like and TM-like branches",
)