import tidy3d as td
import tidy3d.web as web
from tidy3d.plugins.mode import ModeSolver

# Custom anisotropic hBN medium for 585 nm wavelength

# Initialize wavelength and wavenumber points
wavelengths = np.linspace(0.25, 1, 1000)  # Wavelength range in um
wavenumber = 1e4 / wavelengths

# Extraordinary Axis

# Define fitting parameters from paper
A = 110000  # sigma P0 para
B = 90000  # #sigma 0 para

# Generate data points
n_data = np.array((1 + A**2 / (B**2 - (wavenumber**2))) ** (1 / 2))

# Fit
advanced_param = AdvancedFastFitterParam(weights=(1, 1), num_iters=100)
fitter = FastDispersionFitter(wvl_um=wavelengths, n_data=n_data, wvl_range=(0.25, 1))
hbn_zz, rms_error = fitter.fit(
    max_num_poles=3, tolerance_rms=5e-2, eps_inf=1.0, advanced_param=advanced_param
)


# Plot
fitter.plot(hbn_zz)
plt.xlabel("wavelength ($\mu m$)")
plt.ylabel("n")
plt.title("Extraordinary Refractive Index Data")
plt.legend()
plt.show()

# Ordinary Axis

# Define fitting parameters from paper
C = 124100  # sigma P0 perp
D = 75100  # #sigma 0 perp
E = 53200  # sigma P1 perp
F = 49200  # sigma 1 perp
G = 1560  # gamma 1 perp

# Generate data points
n_data = np.array(
    (
        1
        + (C**2 / (D**2 - (wavenumber**2)))
        + (
            (E**2 * (F**2 - wavenumber**2))
            / ((F**2 - wavenumber**2) ** 2 + (G**2 * wavenumber**2))
        )
    )
    ** (1 / 2)
)


advanced_param = AdvancedFastFitterParam(weights=(1, 1), num_iters=100)
fitter = FastDispersionFitter(wvl_um=wavelengths, n_data=n_data, wvl_range=(0.25, 1))
hbn_xx, rms_error = fitter.fit(
    max_num_poles=3, tolerance_rms=5e-2, eps_inf=1.0, advanced_param=advanced_param
)

fitter.plot(hbn_xx)
plt.xlabel("wavelength ($\mu m$)")
plt.ylabel("n")
plt.title("Ordinary Refractive Index Data")
plt.legend()
plt.show()

hBN_ani = td.AnisotropicMedium(name="hBN_ani", xx=hbn_xx, yy=hbn_xx, zz=hbn_zz)

# Define the sweep
waveguide_thickness_sweep = np.linspace(0.08, 0.14, 13)
nwg_thickness_sweep = np.linspace(0.04, 0.12, 9)

# Define the path to save data
path = "data/20260122 thickness sweep/"

# Build the simulation batch
sims_thick = {
    f"wg_thickness={wgt * 1e3:.3f}nm;nwg_thickness={nwgt * 1e3:.3f}nm": Nanowaveguide_sim_thick(
        wgt, nwgt
    )
    for wgt in waveguide_thickness_sweep
    for nwgt in nwg_thickness_sweep
}

# Run the simulation batch
batch_thick = web.Batch(simulations=sims_thick)
est_fc = batch_thick.estimate_cost()
batch_thick_results = batch_thick.run(path_dir=path)
real_fc = batch_thick.real_cost()

# Analyze the results

# Uncomment below to load from file
# batch_thick = web.Batch.from_file("data/20260122 thickness sweep/batch.hdf5")
# batch_thick_results = batch_thick.load("data/20260122 thickness sweep/")

# Define arrays
Nwg_mode_thick = np.zeros([len(nwg_thickness_sweep), len(waveguide_thickness_sweep)])
Wg_mode_thick = np.zeros([len(nwg_thickness_sweep), len(waveguide_thickness_sweep)])

# Populate data arrays with coupling coefficients averaged across simulated frequencies and normalized by the dipole source power
for i in range(len(waveguide_thickness_sweep)):
    for j in range(len(nwg_thickness_sweep)):
        wgt = waveguide_thickness_sweep[i]
        nwgt = nwg_thickness_sweep[j]
        dipole_power = np.mean(
            batch_thick_results[
                f"wg_thickness={wgt * 1e3:.3f}nm;nwg_thickness={nwgt * 1e3:.3f}nm"
            ]["flux_dip"].flux
        )
        norm_coup_nwg = (
            np.mean(
                np.abs(
                    batch_thick_results[
                        f"wg_thickness={wgt * 1e3:.3f}nm;nwg_thickness={nwgt * 1e3:.3f}nm"
                    ]["Nwg_mode"].amps.sel(direction="+")
                )
                ** 2
            )
            / dipole_power
        )
        norm_coup_wg = (
            np.mean(
                np.abs(
                    batch_thick_results[
                        f"wg_thickness={wgt * 1e3:.3f}nm;nwg_thickness={nwgt * 1e3:.3f}nm"
                    ]["Wg_mode"].amps.sel(direction="+")
                )
                ** 2
            )
            / dipole_power
        )
        Nwg_mode_thick[j][i] = norm_coup_nwg
        Wg_mode_thick[j][i] = norm_coup_wg


# Coupling through the taper between light in the nanowaveguide mode and the waveguide mode found using the ratio of the total coupling in the two modes
Nwg_Wg_coup_thick = Wg_mode_thick / Nwg_mode_thick

# Redefine ticks in nm for plotting
wg_tick = np.linspace(80, 140, 7)
nwg_tick = np.linspace(40, 120, 9)

# Initialize figure
fig, axs = plt.subplots(3, 1)
fig.tight_layout()
fig.set_size_inches(6, 18)
fig.set_dpi(300)
c_map = cm.get_cmap("plasma")

# Plot data arrays
im1 = axs[0].imshow(
    Nwg_mode_thick,
    interpolation="none",
    cmap=c_map,
    aspect="auto",
    extent=[77.5, 142.5, 125, 35],
)
im2 = axs[1].imshow(
    Nwg_Wg_coup_thick,
    interpolation="none",
    cmap=c_map,
    aspect="auto",
    extent=[77.5, 142.5, 125, 35],
)
im3 = axs[2].imshow(
    Wg_mode_thick,
    interpolation="none",
    cmap=c_map,
    aspect="auto",
    extent=[77.5, 142.5, 125, 35],
)

# Figure formating for emitter to nanowaveguide coupling analysis
axs[0].set_title("Emitter to Nanowaveguide - Coupling Coefficient", fontsize=18)
axs[0].set_ylabel("hBN  Thickness (nm)", fontsize=18)
axs[0].set_xlabel("Si3N4 Thickness (nm)", fontsize=18)
axs[0].set_xticks(wg_tick)
axs[0].set_yticks(nwg_tick)
axs[0].tick_params(which="major", direction="in", length=8)
axs[0].tick_params(which="minor", direction="in", length=4)
axs[0].tick_params(axis="both", labelsize=18, pad=5)
axs[0].set_ylabel("hBN  Thickness (nm)", fontsize=18)
axs[0].set_xlabel("Si3N4 Thickness (nm)", fontsize=18)

cbar1 = plt.colorbar(im1)
cbar1.set_label("Coupling Coefficient", fontsize=18)
cbar1.ax.tick_params(labelsize=18)

# Figure formating for emitter to nanowaveguide coupling analysis
axs[1].set_title("Nanowaveguide to Waveguide - Coupling Coefficient", fontsize=18)
axs[1].set_xticks(wg_tick)
axs[1].set_yticks(nwg_tick)
axs[1].tick_params(which="major", direction="in", length=8)
axs[1].tick_params(which="minor", direction="in", length=4)
axs[1].tick_params(axis="both", labelsize=18, pad=5)
axs[1].set_ylabel("hBN  Thickness (nm)", fontsize=18)
axs[1].set_xlabel("Si3N4 Thickness (nm)", fontsize=18)

cbar2 = plt.colorbar(im2)
cbar2.set_label("Coupling Coefficient", fontsize=18)
cbar2.ax.tick_params(labelsize=18)

# Figure formating for emitter to nanowaveguide coupling analysis
axs[2].set_title("Total Waveguide - Coupling Coefficient", fontsize=18)
axs[2].set_xticks(wg_tick)
axs[2].set_yticks(nwg_tick)
axs[2].tick_params(which="major", direction="in", length=8)
axs[2].tick_params(which="minor", direction="in", length=4)
axs[2].tick_params(axis="both", labelsize=18, pad=5)
axs[2].set_ylabel("hBN  Thickness (nm)", fontsize=18)
axs[2].set_xlabel("Si3N4 Thickness (nm)", fontsize=18)

cbar3 = plt.colorbar(im3)
cbar3.set_label("Coupling Coefficient", fontsize=18)
cbar3.ax.tick_params(labelsize=18)

plt.show()

max_flat_index = np.argmax(Wg_mode_thick)
nwg_ind, wg_ind = np.unravel_index(max_flat_index, Wg_mode_thick.shape)

max_wg = waveguide_thickness_sweep[wg_ind]
max_nwg = nwg_thickness_sweep[nwg_ind]
max_coup = Wg_mode_thick[nwg_ind][wg_ind]

print(
    f"Maximum total coupling of {max_coup:.3f} with {round(max_wg * 1e3)} nm thick waveguide and {round(max_nwg * 1e3)} nm thick nanowaveguide"
)  # Output: 89

waveguide_thickness = max_wg
# for these simulations we choose to overide the optimal waveguide thickness, for a slightly thicker waveguide because it benefits racetrack resonator coupling for the full device investigate in the companion paper
waveguide_thickness = 0.12  # um
nwg_thickness = max_nwg

# Define the sweep
taper_length_sweep = np.linspace(1, 50, 50)
taper_slope_sweep = np.linspace(1e-3, 30e-3, 30)

# Define the path to save data
path = "data/20260126 taper sweep/"

# Build the simulation batch. We include an if statement to filter out invalid cases where the taper length is too long for larger taper slopes resulting in invalid polygons.
sims_tap = {
    f"taper_length={tl:.3f}um;taper_slope={ts * 1e3:.3f}": Nanowaveguide_sim_tap(tl, ts)
    for tl in taper_length_sweep
    for ts in taper_slope_sweep
    if (nwg_width / 2 - ts * tl) > 0
}
# Run the simulation batch
batch_tap = web.Batch(simulations=sims_tap)
est_fc = batch_tap.estimate_cost()
batch_tap_results = batch_tap.run(path_dir=path)
real_fc = batch_tap.real_cost()

# Analyze the results

# Uncomment below to load from file
# batch_tap = web.Batch.from_file("data/20260126 taper sweep/batch.hdf5")
# batch_tap_results = batch_tap.load("data/20260126 taper sweep/")

# Define arrays
Nwg_mode_tap = np.zeros([len(taper_slope_sweep), len(taper_length_sweep)])
Wg_mode_tap = np.zeros([len(taper_slope_sweep), len(taper_length_sweep)])

# Populate data arrays with coupling coefficients averaged across simulated frequencies and normalized by the dipole source power
for i in range(len(taper_length_sweep)):
    for j in range(len(taper_slope_sweep)):
        tl = taper_length_sweep[i]
        ts = taper_slope_sweep[j]
        if (nwg_width / 2 - ts * tl) > 0:
            dipole_power = np.mean(
                batch_tap_results[
                    f"taper_length={tl:.3f}um;taper_slope={ts * 1e3:.3f}"
                ]["flux_dip"].flux
            )
            norm_coup_nwg = (
                np.mean(
                    np.abs(
                        batch_tap_results[
                            f"taper_length={tl:.3f}um;taper_slope={ts * 1e3:.3f}"
                        ]["Nwg_mode"].amps.sel(direction="+")
                    )
                    ** 2
                )
                / dipole_power
            )
            norm_coup_wg = (
                np.mean(
                    np.abs(
                        batch_tap_results[
                            f"taper_length={tl:.3f}um;taper_slope={ts * 1e3:.3f}"
                        ]["Wg_mode"].amps.sel(direction="+")
                    )
                    ** 2
                )
                / dipole_power
            )
            Nwg_mode_tap[j][i] = norm_coup_nwg
            Wg_mode_tap[j][i] = norm_coup_wg

# Redefine ticks in nm for plotting
len_tick = [1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50]
slope_tick = np.linspace(2e-3, 30e-3, 15) * 1e3

# Initialize figure
fig, axs = plt.subplots(3, 1)
fig.tight_layout()
fig.set_size_inches(6, 18)
fig.set_dpi(300)
c_map = cm.get_cmap("plasma")

# Set invalid taper length and slope combinations to NaN in array and white on plot
c_map.set_bad("white")
Nwg_mode_tap_copy = Nwg_mode_tap.copy()
Nwg_mode_tap[Nwg_mode_tap == 0] = np.nan
Wg_mode_tap_copy = Wg_mode_tap.copy()
Wg_mode_tap[Wg_mode_tap == 0] = np.nan

# Coupling through the taper between light in the nanowaveguide mode and the waveguide mode found using the ratio of the total coupling in the two modes
Nwg_Wg_coup_tap = Wg_mode_tap / Nwg_mode_tap

# Plot data arrays
im1 = axs[0].imshow(
    Nwg_mode_tap,
    interpolation="none",
    cmap=c_map,
    aspect="auto",
    extent=[0.5, 50.5, 30.5, 0.5],
)
im2 = axs[1].imshow(
    Nwg_Wg_coup_tap,
    interpolation="none",
    cmap=c_map,
    aspect="auto",
    extent=[0.5, 50.5, 30.5, 0.5],
)
im3 = axs[2].imshow(
    Wg_mode_tap,
    interpolation="none",
    cmap=c_map,
    aspect="auto",
    extent=[0.5, 50.5, 30.5, 0.5],
)

# Figure formating for emitter to nanowaveguide coupling analysis
axs[0].set_title("Emitter to Nanowaveguide - Coupling Coefficient", fontsize=18)
axs[0].set_ylabel("hBN  Thickness (nm)", fontsize=18)
axs[0].set_xlabel("Si3N4 Thickness (nm)", fontsize=18)
axs[0].set_xticks(len_tick)
axs[0].set_yticks(slope_tick)
axs[0].tick_params(which="major", direction="in", length=8)
axs[0].tick_params(which="minor", direction="in", length=4)
axs[0].tick_params(axis="both", labelsize=18, pad=5)
axs[0].set_ylabel("Taper Slope (nm/$\mu m$)", fontsize=18)
axs[0].set_xlabel("Taper Length ($\mu m$)", fontsize=18)

cbar1 = plt.colorbar(im1)
cbar1.set_label("Coupling Coefficient", fontsize=18)
cbar1.ax.tick_params(labelsize=18)

# Figure formating for emitter to nanowaveguide coupling analysis
axs[1].set_title("Nanowaveguide to Waveguide - Coupling Coefficient", fontsize=18)
axs[1].set_xticks(len_tick)
axs[1].set_yticks(slope_tick)
axs[1].tick_params(which="major", direction="in", length=8)
axs[1].tick_params(which="minor", direction="in", length=4)
axs[1].tick_params(axis="both", labelsize=18, pad=5)
axs[1].set_ylabel("Taper Slope (nm/$\mu m$)", fontsize=18)
axs[1].set_xlabel("Taper Length ($\mu m$)", fontsize=18)

cbar2 = plt.colorbar(im2)
cbar2.set_label("Coupling Coefficient", fontsize=18)
cbar2.ax.tick_params(labelsize=18)

# Figure formating for emitter to nanowaveguide coupling analysis
axs[2].set_title("Total Waveguide - Coupling Coefficient", fontsize=18)
axs[2].set_xticks(len_tick)
axs[2].set_yticks(slope_tick)
axs[2].tick_params(which="major", direction="in", length=8)
axs[2].tick_params(which="minor", direction="in", length=4)
axs[2].tick_params(axis="both", labelsize=18, pad=5)
axs[2].set_ylabel("Taper Slope (nm/$\mu m$)", fontsize=18)
axs[2].set_xlabel("Taper Length ($\mu m$)", fontsize=18)

cbar3 = plt.colorbar(im3)
cbar3.set_label("Coupling Coefficient", fontsize=18)
cbar3.ax.tick_params(labelsize=18)

plt.show()

max_flat_index = np.nanargmax(Wg_mode_tap)
ts_ind, tl_ind = np.unravel_index(max_flat_index, Wg_mode_tap.shape)

max_tl = taper_length_sweep[tl_ind]
max_ts = taper_slope_sweep[ts_ind]
max_coup = Wg_mode_tap[ts_ind][tl_ind]

print(
    f"Maximum total coupling of {max_coup:.3f} with {max_tl} um taper length and {max_ts * 1e3:.1f} nm/um taper slope"
)  # Output: 89

taper_length = max_tl
taper_slope = max_ts

path = "data/20260126 field profile/"

sim = Nanowaveguide_sim_hires()

name = (
    "Nanowaveguide_TL_"
    + str(int(taper_length))
    + "um_TS_"
    + str(round(taper_slope * 1e3))
    + "_wgT_"
    + str(round(waveguide_thickness * 1e3))
    + "nm_nwgT_"
    + str(round(nwg_thickness * 1e3))
    + "nm_hires"
)
job = web.Job(simulation=sim, task_name=name)
FC = web.estimate_cost(job.task_id)
path = path + name + ".hdf5"
sim_data = job.run(path=path)

time.sleep(2)
real_fc = web.real_cost(job.task_id)

# Taper Field Profile
# Uncomment to load from file
# sim_data = td.SimulationData.from_file('data/20260122 field profile//Nanowaveguide_TL_9um_TS_15_wgT_120nm_nwgT_100nm_hires.hdf5')

Ez = sim_data["field"].Ez
Ex = sim_data["field"].Ex
Ey = sim_data["field"].Ey

E_tot = (abs(Ez) ** 2 + abs(Ex) ** 2 + abs(Ey) ** 2) ** (1 / 2)
E_tot = E_tot / np.max(E_tot)

fig, ax1 = plt.subplots()
fig.set_size_inches(12, 7)
fig.tight_layout()
fig.set_dpi(300)

im = E_tot.plot(x="x", y="z", cmap="jet", ax=ax1)

ax1.set_title(None)
ax1.set_ylabel("Z ($\mu m$)", fontsize=18)
ax1.set_xlabel("X ($\mu m$)", fontsize=18)

ax1.set_ylim(-0.3, 0.3)
ax1.tick_params(which="major", direction="in", length=8)
ax1.tick_params(which="minor", direction="in", length=4)
ax1.tick_params(axis="both", labelsize=18, pad=5)

cbar = im.colorbar
cbar.set_label("|E|", fontsize=18)
cbar.ax.tick_params(labelsize=18)

plt.show()

# Coupling coeffiecient summary

dipole_power = np.mean(sim_data["flux_dip"].flux)

norm_coup_nwg = (
    np.mean(np.abs(sim_data["Nwg_mode"].amps.sel(direction="+")) ** 2) / dipole_power
)

norm_coup_wg = (
    np.mean(np.abs(sim_data["Wg_mode"].amps.sel(direction="+")) ** 2) / dipole_power
)

print(
    f"For a nanowaveguide with {taper_length:.0f} um taper length, {taper_slope * 1e3:.0f} nm/um taper slope, {round(waveguide_thickness * 1e3)} nm thick waveguide and {round(nwg_thickness * 1e3)} nm thick nanowaveguide: \nCoupling from emitter to nanowaveguide: {norm_coup_nwg:.3f} \nCoupling from nanowaveguide to waveguide: {norm_coup_wg / norm_coup_nwg:.3f} \nMaximum total coupling from emitter to waveguide {norm_coup_wg:.3f}"
)

ideal_total_CE = norm_coup_wg

# Axes parameters
wg_tick = np.linspace(80, 140, 7)
nwg_tick = np.linspace(40, 120, 9)
len_tick = [1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50]
slope_tick = np.linspace(1e-3, 15e-3, 15) * 1e3

# Initialize figure
c_map = cm.get_cmap("plasma")
c_map.set_bad("white")
fig, axs = plt.subplots(2, 1)
fig.tight_layout()
fig.set_dpi(300)
fig.set_size_inches(6, 12)

# Find coupling loss with respect to ideal parameters
# try:
#     ideal_total_CE = ideal_total_CE.item()
# except:
#     None
max_wg_thick = np.max(np.max(Wg_mode_thick))
max_wg_tap = np.nanmax(np.nanmax(Wg_mode_tap))
Wg_mode_thick_fab = np.subtract(Wg_mode_thick, max_wg_thick)
Wg_mode_tap_fab = np.subtract(Wg_mode_tap, max_wg_tap)

Wg_mode_thick_fab[Wg_mode_thick_fab < -0.02] = np.nan
Wg_mode_tap_fab[Wg_mode_tap_fab < -0.02] = np.nan


# Plot and format
im1 = axs[0].imshow(
    Wg_mode_thick_fab,
    interpolation="none",
    cmap=c_map,
    aspect="auto",
    extent=[77.5, 142.5, 125, 35],
)
im2 = axs[1].imshow(
    Wg_mode_tap_fab,
    interpolation="none",
    cmap=c_map,
    aspect="auto",
    extent=[0.5, 50.5, 15.5, 0.5],
)

fig.set_size_inches(6, 12)
fig.set_dpi(300)

axs[0].set_title("Total Coupling Fabrication Robustness", fontsize=18)
axs[0].set_xlabel("Si3N4 Thickness (nm)", fontsize=18)
axs[0].set_ylabel("hBN  Thickness (nm)", fontsize=18)
axs[0].set_xticks(wg_tick)
axs[0].set_yticks(nwg_tick)
axs[0].tick_params(which="major", direction="in", length=8)
axs[0].tick_params(which="minor", direction="in", length=4)
axs[0].tick_params(axis="both", labelsize=18, pad=5)

axs[1].set_title("Total Coupling Fabrication Robustness", fontsize=18)
axs[1].set_ylabel("Taper Slope (nm/$\mu m$)", fontsize=18)
axs[1].set_xlabel("Taper Length ($\mu m$)", fontsize=18)
axs[1].set_xticks(len_tick)
axs[1].set_yticks(slope_tick)
axs[1].tick_params(which="major", direction="in", length=8)
axs[1].tick_params(which="minor", direction="in", length=4)
axs[1].tick_params(axis="both", labelsize=18, pad=5)

cbar1 = plt.colorbar(im1)
cbar1.set_label("Coupling Loss", fontsize=18)
cbar1.ax.tick_params(labelsize=18)

cbar2 = plt.colorbar(im2)
cbar2.set_label("Coupling Loss", fontsize=18)
cbar2.ax.tick_params(labelsize=18)

plt.show()

# Start with dipole height simulation

# Define the sweep
dipole_height_sweep = np.arange(0.001, 0.1, 0.005)

# Define the path to save data
path = "data/20260126 dipole height sweep/"

# Build the simulation batch.
sims_dh = {
    f"dipole_height={dh * 1e3:.0f}nm": Nanowaveguide_sim_fab(
        dh, 0, 0, [np.pi / 2, np.pi / 2]
    )
    for dh in dipole_height_sweep
}
# Run the simulation batch
batch_dh = web.Batch(simulations=sims_dh)
est_fc = batch_dh.estimate_cost()
batch_dh_results = batch_dh.run(path_dir=path)
real_fc = batch_dh.real_cost()

# Next: deviation in emitter x position

# Define the sweep
dipole_dX_sweep = np.arange(0, 2, 0.1)

# Define the path to save data
path = "data/20260126 dipole X deviation sweep/"

# Build the simulation batch.
sims_dX = {
    f"dipole_dX={dX * 1e3:.0f}nm": Nanowaveguide_sim_fab(
        0, dX, 0, [np.pi / 2, np.pi / 2]
    )
    for dX in dipole_dX_sweep
}
# Run the simulation batch
batch_dX = web.Batch(simulations=sims_dX)
est_fc = batch_dX.estimate_cost()
batch_dX_results = batch_dX.run(path_dir=path)
real_fc = batch_dX.real_cost()

# Next: deviation in emitter y position

# Define the sweep
dipole_dY_sweep = np.arange(0, 0.199, 0.01)

# Define the path to save data
path = "data/20260126 dipole Y deviation sweep/"

# Build the simulation batch.
sims_dY = {
    f"dipole_dY={dY * 1e3:.0f}nm": Nanowaveguide_sim_fab(
        0, 0, dY, [np.pi / 2, np.pi / 2]
    )
    for dY in dipole_dY_sweep
}
# Run the simulation batch
batch_dY = web.Batch(simulations=sims_dY)
est_fc = batch_dY.estimate_cost()
batch_dY_results = batch_dY.run(path_dir=path)
real_fc = batch_dY.real_cost()


# Polarization Sweeps
dipole_theta_sweep = np.linspace(0, 1, 11) * np.pi / 2
path = "data/20260126 dipole theta sweep/"
sims_theta = {
    f"dipole_theta={t:.03}rad": Nanowaveguide_sim_fab(0, 0, 0, [t, np.pi / 2])
    for t in dipole_theta_sweep
}
batch_theta = web.Batch(simulations=sims_theta)
est_fc = batch_theta.estimate_cost()
batch_theta_results = batch_theta.run(path_dir=path)
real_fc = batch_theta.real_cost()


dipole_phi_sweep = np.linspace(0, 1, 11) * np.pi / 2
path = "data/20260126 dipole phi sweep/"
sims_phi = {
    f"dipole_phi={p:.03}rad": Nanowaveguide_sim_fab(0, 0, 0, [np.pi / 2, p])
    for p in dipole_phi_sweep
}
batch_phi = web.Batch(simulations=sims_phi)
est_fc = batch_phi.estimate_cost()
batch_phi_results = batch_phi.run(path_dir=path)
real_fc = batch_phi.real_cost()