import re
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from IPython.display import display, Image

import tidy3d as td
import tidy3d.web as web

# Keep repeated non-fatal geometry warnings from cluttering the submission notebook.
# Remove this line while debugging geometry if you want to see all Tidy3D warnings.
td.config.logging_level = "ERROR"

print("Tidy3D version:", td.__version__)
print("NumPy version:", np.__version__)

DATA_DIR = Path("data_mmi_splitter")
DATA_DIR.mkdir(exist_ok=True)

SWEEP_CSV = DATA_DIR / "mmi_length_sweep_results.csv"
CONV_CSV = DATA_DIR / "mmi_convergence_check.csv"
FIELD_PNG = DATA_DIR / "best_mmi_field.png"
FIELD_HDF5 = DATA_DIR / "best_mmi_field.hdf5"

# Units: Tidy3D uses micrometers for length.
um = 1.0

# Wavelength and frequency setup.
lambda0 = 1.55 * um
freq0 = td.C_0 / lambda0
fwidth = 0.12 * freq0
freqs = np.array([freq0])

# 2D effective-index model.
# These are typical effective-index-like values, not a specific foundry stack.
n_core = 2.80
n_clad = 1.44

core = td.Medium(permittivity=n_core**2)
clad = td.Medium(permittivity=n_clad**2)

# Waveguide and MMI geometry.
wg_width = 0.45 * um
mmi_width = 3.20 * um
output_sep = 1.05 * um

input_len = 2.0 * um
output_len = 2.0 * um

# Simulation domain padding.
pad_x = 1.0 * um
pad_y = 1.2 * um

# Temporal and numerical settings.
min_steps_per_wvl_default = 30
run_time = 1.8e-12

# Compact MMI length sweep.
mmi_lengths = np.linspace(3.8, 6.2, 9) * um

print("Number of MMI lengths:", len(mmi_lengths))
print("MMI length sweep [um]:", np.round(mmi_lengths, 3))

preview_sim = make_mmi_sim(
    mmi_lengths[len(mmi_lengths) // 2], include_field_monitor=True
)

fig, ax = plt.subplots(figsize=(9, 4))
preview_sim.plot(z=0, ax=ax)
ax.set_title("2D effective-index 1x2 MMI splitter geometry")
plt.show()

if SWEEP_CSV.exists():
    print(f"Loading cached sweep results from: {SWEEP_CSV}")
    sweep_df = pd.read_csv(SWEEP_CSV)
else:
    sweep_sims = {
        f"L_{float(L):.3f}um": make_mmi_sim(L, include_field_monitor=False)
        for L in mmi_lengths
    }

    sweep_batch = web.Batch(simulations=sweep_sims, verbose=True)
    sweep_batch_results = sweep_batch.run(path_dir=str(DATA_DIR / "sweep"))

    sweep_df = summarize_mmi_batch(
        batch_results=sweep_batch_results,
        output_csv=SWEEP_CSV,
    )
    print(f"Saved results to: {SWEEP_CSV}")

display(sweep_df)

best_row = sweep_df.sort_values("split_score", ascending=False).iloc[0]
best_length = float(best_row["mmi_length_um"]) * um

print("\nBest MMI length:")
display(best_row)
print(f"Best length = {best_length:.3f} um")

fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(sweep_df["mmi_length_um"], sweep_df["P_top"], marker="o", label="Top output")
ax.plot(
    sweep_df["mmi_length_um"], sweep_df["P_bottom"], marker="s", label="Bottom output"
)
ax.plot(sweep_df["mmi_length_um"], sweep_df["P_total"], marker="^", label="Total")
ax.set_xlabel("MMI length (um)")
ax.set_ylabel("Modal power")
ax.set_title("Output modal powers versus MMI length")
ax.grid(True, alpha=0.3)
ax.legend()
plt.show()

fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(sweep_df["mmi_length_um"], sweep_df["imbalance"], marker="o")
ax.set_xlabel("MMI length (um)")
ax.set_ylabel("Splitting imbalance")
ax.set_title("50:50 splitting error versus MMI length")
ax.grid(True, alpha=0.3)
plt.show()

if FIELD_PNG.exists():
    print(f"Using cached field plot: {FIELD_PNG}")
    display(Image(filename=str(FIELD_PNG)))
else:
    best_sim = make_mmi_sim(
        best_length,
        min_steps_per_wvl=min_steps_per_wvl_default,
        include_field_monitor=True,
    )

    best_data = web.run(
        best_sim,
        task_name=f"best_mmi_L_{best_length:.3f}um",
        path=str(FIELD_HDF5),
        verbose=True,
    )

    field_data = best_data["field_xy"]

    # Prefer Ez for a clean scalar plot; fall back to another component if needed.
    component_name = None
    for candidate in ["Ez", "Ey", "Ex"]:
        if hasattr(field_data, candidate):
            component_name = candidate
            break

    if component_name is None:
        raise RuntimeError("No E-field component was found in the field monitor data.")

    field = (
        getattr(field_data, component_name)
        .sel(f=freq0, method="nearest")
        .real.squeeze()
    )

    fig, ax = plt.subplots()
    field.plot(ax=ax)
    ax.set_title(
        f"Real({component_name}) field for best MMI design, L = {best_length:.3f} um"
    )
    ax.set_aspect("equal")
    plt.tight_layout()
    fig.savefig(FIELD_PNG, dpi=200)
    plt.show()

print(f"Field plot path: {FIELD_PNG}")

if CONV_CSV.exists():
    print(f"Loading cached convergence results from: {CONV_CSV}")
    conv_df = pd.read_csv(CONV_CSV)
else:
    grid_values = [24, 30, 36]
    conv_sims = {
        f"grid_{grid}": make_mmi_sim(
            best_length,
            min_steps_per_wvl=grid,
            include_field_monitor=False,
        )
        for grid in grid_values
    }

    conv_batch = web.Batch(simulations=conv_sims, verbose=True)
    conv_batch_results = conv_batch.run(path_dir=str(DATA_DIR / "convergence"))

    rows = []
    for task_name in conv_batch_results.keys():
        sim_data = conv_batch_results[task_name]
        grid = int(str(task_name).replace("grid_", ""))

        p_top = extract_forward_mode_power(sim_data, "top_mode")
        p_bottom = extract_forward_mode_power(sim_data, "bottom_mode")
        metrics = score_splitter(p_top, p_bottom)

        rows.append(
            {
                "min_steps_per_wvl": grid,
                "mmi_length_um": best_length,
                **metrics,
            }
        )

    conv_df = pd.DataFrame(rows).sort_values("min_steps_per_wvl").reset_index(drop=True)
    conv_df.to_csv(CONV_CSV, index=False)
    print(f"Saved convergence results to: {CONV_CSV}")

display(conv_df)

fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(
    conv_df["min_steps_per_wvl"], conv_df["P_total"], marker="o", label="Total power"
)
ax.plot(
    conv_df["min_steps_per_wvl"], conv_df["imbalance"], marker="s", label="Imbalance"
)
ax.set_xlabel("Minimum grid steps per wavelength")
ax.set_title("Convergence check for best MMI design")
ax.grid(True, alpha=0.3)
ax.legend()
plt.show()