import xarray as xr
import numpy as np
from numpy.polynomial import Polynomial

import tidy3d as td
from tidy3d.plugins.dispersion import AdvancedFastFitterParam, FastDispersionFitter
from tidy3d.plugins.resonance import ResonanceFinder
from tidy3d.plugins.mode import ModeSolver
from tidy3d.components.simulation import AbstractYeeGridSimulation
from tidy3d.components.structure import Structure
from tidy3d.components.grid.grid import Grid, Coords
import matplotlib.cm as cm
import matplotlib.pyplot as plt
from scipy.constants import c

from tqdm import tqdm

td.config.logging_level = "ERROR"

def get_fdtd_sim(
    simulation,
    ports,
    freqs,
    layer_z_centers,
    inputs,
    outputs,
    mode_volume=None,
    farfield=None,
    flux=None,
    output_freqs=None,
    custom_sources=None,
    sim_kwargs=None,
):
    """Returns the FDTD simulation object for a given set of inputs and outputs.

    Args:
        simulation: base td.Simulation (structures + domain, no sources/monitors)
        ports: dict of port info, e.g.
            {'o1': {'center': (x,y,z), 'size': (sx,sy,sz),
                    'direction': '+'/'-', 'mode_spec': td.ModeSpec}}
        freqs: array of frequencies used to compute default bandwidth
        layer_z_centers: dict {'diamond': z_value, ...} for field monitor positioning
        inputs: dict{port_name: {modes, amps, phases, freqs, fwidths}}
        outputs: dict{port_name: {types: list[str]}}
        mode_volume: dict with keys 'layer', 'thickness', 'downsample', etc.
        farfield: dict with keys 'downsample', 'surface', 'res', etc.
        flux: dict with keys 'downsample', 'apodization', etc.
        output_freqs: list of floats
        custom_sources: list of custom td sources
        sim_kwargs: dict of kwargs passed to simulation update
    """
    sources = []
    if custom_sources is not None:
        if not isinstance(custom_sources, list):
            custom_sources = [custom_sources]
        sources.extend(custom_sources)
    monitors = []
    if sim_kwargs is None:
        sim_kwargs = {}

    # Validate port names
    port_names = set(list(inputs.keys()) + list(outputs.keys()))
    missing_ports = port_names - set(ports.keys())
    if missing_ports:
        raise ValueError(f"Ports {missing_ports} not found in ports dict")
    if output_freqs is None:
        output_freqs = np.mean(freqs)

    for port_name, port_info in ports.items():
        p_center = port_info["center"]
        p_size = port_info["size"]
        p_direction = port_info["direction"]
        p_mode_spec = port_info["mode_spec"]

        if port_name in inputs:
            inp = inputs[port_name]
            if "freqs" not in inp:
                inp["freqs"] = [None] * len(inp["modes"])
            if "fwidths" not in inp:
                inp["fwidths"] = [None] * len(inp["modes"])
            if "port_offset" not in inp:
                inp["port_offset"] = [(0, 0, 0)] * len(inp["modes"])
            for (
                input_mode,
                input_amp,
                input_phase,
                input_freq,
                input_fwidth,
                input_port_offset,
            ) in zip(
                inp["modes"],
                inp["amps"],
                inp["phases"],
                inp["freqs"],
                inp["fwidths"],
                inp["port_offset"],
            ):
                if isinstance(input_mode, str) and ("dipole" in input_mode):
                    freq0 = input_freq if input_freq is not None else np.mean(freqs)
                    if input_fwidth is None:
                        fdiff = max(freqs) - min(freqs)
                        fwidth = max(fdiff, max(output_freqs) / 10)
                    else:
                        fwidth = input_fwidth
                    dipole_source = td.PointDipole(
                        center=(
                            p_center[0] + input_port_offset[0],
                            p_center[1] + input_port_offset[1],
                            p_center[2] + input_port_offset[2],
                        ),
                        source_time=td.GaussianPulse(
                            freq0=freq0,
                            fwidth=fwidth,
                            phase=input_phase,
                            amplitude=input_amp,
                        ),
                        name=port_name
                        + f"@{input_mode}"
                        + (f"@{input_freq}" if input_freq is not None else ""),
                        polarization=input_mode.split("_")[1],
                    )
                    sources.append(dipole_source)
                else:
                    freq0 = input_freq if input_freq is not None else np.mean(freqs)
                    if input_fwidth is None:
                        fdiff = max(freqs) - min(freqs)
                        fwidth = max(fdiff, freq0 * 0.1)
                    else:
                        fwidth = input_fwidth
                    source_time = td.GaussianPulse(
                        freq0=freq0,
                        fwidth=fwidth,
                        amplitude=input_amp,
                        phase=input_phase,
                    )
                    mode_source = td.ModeSource(
                        center=p_center,
                        size=p_size,
                        source_time=source_time,
                        mode_spec=p_mode_spec,
                        mode_index=input_mode,
                        direction=p_direction,
                        name=port_name,
                    )
                    sources.append(mode_source)

        if port_name in outputs:
            if len(outputs[port_name]["types"]) == 0:
                outputs[port_name]["types"] = ["mode"]
            for output_type in outputs[port_name]["types"]:
                if output_type == "mode":
                    output_monitor = td.ModeMonitor(
                        center=p_center,
                        size=p_size,
                        freqs=output_freqs
                        if isinstance(output_freqs, (list, np.ndarray))
                        else [output_freqs],
                        mode_spec=p_mode_spec,
                        name=port_name,
                    )
                elif "time" in output_type:
                    fdiff = max(freqs) - min(freqs)
                    fwidth = max(fdiff, max(output_freqs) / 10)
                    t_start = 2 / fwidth
                    output_monitor = td.FieldTimeMonitor(
                        fields=[output_type.split("_")[1]],
                        center=p_center,
                        size=(0, 0, 0),
                        start=t_start,
                        name=port_name + "_" + output_type,
                    )
                monitors.append(output_monitor)

    sim_center = simulation.center
    sim_size = sim_kwargs.get("size", simulation.size)
    if mode_volume is not None:
        if "downsample" not in mode_volume or mode_volume["downsample"] is None:
            mode_volume["downsample"] = (4, 4, 2)
        if "apodization" in mode_volume:
            apod_spec = td.ApodizationSpec(
                start=mode_volume["apodization"]["start"],
                width=mode_volume["apodization"]["width"],
            )
        else:
            apod_spec = td.ApodizationSpec()
        if "colocate" in mode_volume:
            colocate = mode_volume["colocate"]
        else:
            colocate = True
        monitors.append(
            td.FieldMonitor(
                name="field",
                center=[0, 0, layer_z_centers[mode_volume["layer"]]],
                size=[td.inf, td.inf, mode_volume["thickness"]],
                interval_space=mode_volume["downsample"],
                freqs=output_freqs
                if isinstance(output_freqs, (list, np.ndarray))
                else [output_freqs],
                apodization=apod_spec,
                colocate=colocate,
            )
        )
    if farfield is not None:
        if "downsample" not in farfield or farfield["downsample"] is None:
            farfield["downsample"] = (4, 4, 4)
        if "apodization" in farfield:
            apod_spec = td.ApodizationSpec(
                start=farfield["apodization"]["start"],
                width=farfield["apodization"]["width"],
            )
        else:
            apod_spec = td.ApodizationSpec()
        distance_from_pml = 0.1
        if "surface" not in farfield or farfield["surface"] == "spherical":
            thetas = farfield.get("thetas", np.linspace(0, np.pi, farfield["res"]))
            phis = farfield.get("phis", np.linspace(0, 2 * np.pi, 2 * farfield["res"]))
            exclude_surfaces = farfield.get("exclude_surfaces", None)
            monitors.append(
                td.FieldProjectionAngleMonitor(
                    name="farfield_monitor",
                    center=sim_center,
                    size=(
                        sim_size[0] - 2 * distance_from_pml,
                        sim_size[1] - 2 * distance_from_pml,
                        sim_size[2] - 2 * distance_from_pml,
                    ),
                    interval_space=farfield["downsample"],
                    freqs=output_freqs
                    if isinstance(output_freqs, (list, np.ndarray))
                    else [output_freqs],
                    theta=thetas,
                    phi=phis,
                    apodization=apod_spec,
                    exclude_surfaces=exclude_surfaces,
                )
            )
        elif farfield["surface"] == "cartesian":
            proj_distance = farfield.get("proj_distance", 1e6)
            res = farfield.get("res", 100)
            x = farfield.get(
                "x", np.linspace(-proj_distance / 2, proj_distance / 2, res)
            )
            y = farfield.get(
                "y", np.linspace(-proj_distance / 2, proj_distance / 2, res)
            )
            exclude_surfaces = farfield.get("exclude_surfaces", None)
            monitors.append(
                td.FieldProjectionCartesianMonitor(
                    name="farfield_monitor_cartesian",
                    center=sim_center,
                    size=(
                        sim_size[0] - 2 * distance_from_pml,
                        sim_size[1] - 2 * distance_from_pml,
                        sim_size[2] - 2 * distance_from_pml,
                    ),
                    interval_space=farfield["downsample"],
                    freqs=output_freqs
                    if isinstance(output_freqs, (list, np.ndarray))
                    else [output_freqs],
                    x=x,
                    y=y,
                    proj_axis=2,
                    proj_distance=proj_distance,
                    apodization=apod_spec,
                    exclude_surfaces=exclude_surfaces,
                )
            )
    if flux is not None:
        if "downsample" not in flux or flux["downsample"] is None:
            flux["downsample"] = (4, 4, 4)
        if "apodization" in flux:
            apod_spec = td.ApodizationSpec(
                start=flux["apodization"]["start"], width=flux["apodization"]["width"]
            )
        else:
            apod_spec = td.ApodizationSpec()
        distance_from_pml = 0.1
        surfaces = {
            "left": "x-",
            "right": "x+",
            "back": "y-",
            "front": "y+",
            "bottom": "z-",
            "top": "z+",
        }
        for name, direction in surfaces.items():
            excludes = list({k: v for k, v in surfaces.items() if k != name}.values())
            monitors.append(
                td.FluxMonitor(
                    name="flux_" + name,
                    center=sim_center,
                    size=(
                        sim_size[0] - 2 * distance_from_pml,
                        sim_size[1] - 2 * distance_from_pml,
                        sim_size[2] - 2 * distance_from_pml,
                    ),
                    interval_space=flux["downsample"],
                    freqs=output_freqs
                    if isinstance(output_freqs, (list, np.ndarray))
                    else [output_freqs],
                    apodization=apod_spec,
                    exclude_surfaces=excludes,
                )
            )

    sim = simulation.updated_copy(sources=sources, monitors=monitors, **sim_kwargs)
    return sim

class NeffMixer:
    def __init__(self, freqs, neffs, num_modes, layer_widths, layer_height, meta_data):
        """Initializes a NeffMixer object.
        inputs:
            freqs (Nx1 array): list of frequencies
            neffs (WxNxM array): list of effective indices
            num_modes (int): M, the number of modes
            layer_widths (Wx1 array): list of widths for the chosen layer
            layer_height (float): height of the chosen layer
            meta_data (dict(str, any)): dictionary of meta data
        """
        self.freqs = freqs
        self.neffs = neffs
        self.num_modes = num_modes
        self.layer_widths = layer_widths
        self.layer_height = layer_height
        self.meta_data = meta_data

    def plot(self, mode_num=0):
        freqs_grid, wg_widths_grid = np.meshgrid(self.freqs, self.layer_widths)
        plt.figure()
        plt.pcolormesh(
            freqs_grid,
            wg_widths_grid,
            self.neffs[:, :, mode_num],
            shading="auto",
            cmap="bwr",
        )
        plt.colorbar(label="neff")
        plt.xlabel("Frequency (THz)")
        plt.ylabel("Waveguide Width (um)")
        plt.title(f"Mode {mode_num} Effective Index")
        plt.show()

    def net_momentum_matrix(
        self,
        fixed_freqs,
        mode_factors,
        mode_nums=0,
        unit_ks=1,
        fit_tol=1e-3,
        poly_order=3,
        plot=False,
        num_freqs=100,
        center_colorbar=False,
    ):
        """Returns the fwm net momentum matrix for mixing that targets the given output frequency.
        inputs:
            fixed_freqs ((K-2)x1 array): the frequencies of the fixed modes, where K is the number of modes involved in the mixing
            mode_factors (Kx1 array): list of mode factors, where the last two are the mode factors of the independent and dependent modes, respectively
            mode_nums (Kx1 array): list of mode numbers
            unit_ks (Kx1 array): indicates the unit vector magnitude in the direction of net momentum
            fit_tol (float): tolerance for the fit
            plot (bool): whether to plot the net momentum matrix
        outputs:
            f_A (Vx1 array): the frequencies of the independent mode, where V is the number of valid frequencies
            f_B (Vx1 array): the frequencies of the dependent mode, where V is the number of valid frequencies
            net_momentum_matrix (WxV array): the net momentum matrix, where W is the number of widths and L is the number of valid frequencies
        """
        if np.isscalar(mode_nums):
            mode_nums = np.array([mode_nums] * len(mode_factors))
        if np.isscalar(unit_ks):
            unit_ks = np.array([unit_ks] * len(mode_factors))

        if len(fixed_freqs) != len(mode_factors) - 2:
            raise ValueError(
                "The number of fixed frequencies must be equal to the number of modes minus 2"
            )
        if np.max(mode_nums) > self.num_modes:
            raise ValueError(
                f"The requested mode number ({np.max(mode_nums)}) is too high for the number of modes in the mixer ({self.num_modes})"
            )
        # Independent mode
        min_freq = np.min(self.freqs)
        max_freq = np.max(self.freqs)
        f_A = np.linspace(min_freq, max_freq, num_freqs)
        # Dependent mode
        f_B = -mode_factors[-1] * (
            np.sum(mode_factors[:-2] * fixed_freqs) + mode_factors[-2] * f_A
        )
        if plot:
            plt.figure()
            plt.plot(f_A / 1e12, f_B / 1e12)
            plt.vlines(
                [min_freq / 1e12, max_freq / 1e12],
                min_freq / 1e12,
                max_freq / 1e12,
                color="k",
            )
            plt.hlines(
                [min_freq / 1e12, max_freq / 1e12],
                min_freq / 1e12,
                max_freq / 1e12,
                color="k",
            )
            plt.xlabel("Pump A Frequency (THz)")
            plt.ylabel("Pump B Frequency (THz)")
            plt.title("Pump Frequencies")
            plt.legend(["Pumps", "Mixer Boundary"])
            plt.show()
        # Crop the frequencies to the range of the mixer
        f_A = f_A[(f_B < np.max(self.freqs)) & (f_B > np.min(self.freqs))]
        f_B = f_B[(f_B < np.max(self.freqs)) & (f_B > np.min(self.freqs))]

        if len(f_A) == 0:
            raise ValueError("No valid frequencies found for the mixer")

        net_momentum_matrix = np.zeros((len(self.layer_widths), len(f_A)))
        min_n_matrix = np.zeros((len(self.layer_widths), len(f_A)))
        for w_i, width in enumerate(self.layer_widths):
            k_fits = {}
            n_fits = {}
            for m_n in mode_nums:
                if m_n not in k_fits:
                    n_fit, n_err = Polynomial.fit(
                        self.freqs, self.neffs[w_i, :, m_n], poly_order, full=True
                    )
                    n_fits[m_n] = n_fit
                    k_fits[m_n] = lambda f: 2 * np.pi * f / td.C_0 * n_fit(f)
                    if np.abs(n_err[0]) > fit_tol:
                        print(
                            f"Warning: n_err {np.abs(n_err[0][0]):.3e} is too large for mode {m_n}"
                        )
            for m_n, n_fit in n_fits.items():
                if w_i == 0 or w_i == len(self.layer_widths) - 1:
                    plt.plot(self.freqs, n_fit(self.freqs), label=f"Mode {m_n}")
                    plt.scatter(self.freqs, self.neffs[w_i, :, m_n])
            if w_i == 0 or w_i == len(self.layer_widths) - 1:
                plt.legend()
                plt.show()
            # Add the fixed modes to the net momentum matrix
            fixed_iterable = zip(
                mode_nums[:-2], fixed_freqs, mode_factors[:-2], unit_ks[:-2]
            )
            fixed_ks = np.array(
                [
                    k_fits[mode_num](fixed_freq) * mode_factor * unit_k
                    for mode_num, fixed_freq, mode_factor, unit_k in fixed_iterable
                ]
            )
            net_momentum_matrix[w_i, :] = np.sum(fixed_ks)
            # Add the independent mode to the net momentum matrix
            independent_k = k_fits[mode_nums[-2]](f_A) * mode_factors[-2] * unit_ks[-2]
            net_momentum_matrix[w_i, :] += independent_k
            # Add the dependent mode to the net momentum matrix
            dependent_k = k_fits[mode_nums[-1]](f_B) * mode_factors[-1] * unit_ks[-1]
            net_momentum_matrix[w_i, :] += dependent_k
            # Record the smallest effective index involved in the mixing
            fixed_ns = np.array(
                [
                    n_fits[mode_num](fixed_freq)
                    for mode_num, fixed_freq in zip(mode_nums[:-2], fixed_freqs)
                ]
            )
            independent_n = n_fits[mode_nums[-2]](f_A)
            dependent_n = n_fits[mode_nums[-1]](f_B)
            fixed_ns = np.tile(fixed_ns, (len(independent_n), 1))
            all_ns = np.concatenate(
                (fixed_ns, independent_n[:, np.newaxis], dependent_n[:, np.newaxis]),
                axis=1,
            )
            min_n_matrix[w_i, :] = np.min(all_ns, axis=1)

        if plot:
            f_A_grid, widths_grid = np.meshgrid(f_A, self.layer_widths)
            plt.figure()
            if center_colorbar:
                vmin = -np.max(np.abs(net_momentum_matrix))
                vmax = np.max(np.abs(net_momentum_matrix))
            else:
                vmin = None
                vmax = None
            plt.pcolormesh(
                f_A_grid / 1e12,
                widths_grid,
                net_momentum_matrix,
                cmap="bwr",
                shading="auto",
                vmin=vmin,
                vmax=vmax,
            )
            plt.title("Net Momentum Matrix")
            plt.xlabel("Pump A Frequency (THz)")
            plt.ylabel("Width (um)")
            plt.colorbar()
            plt.show()
            plt.figure()
            plt.pcolormesh(
                f_A_grid / 1e12, widths_grid, min_n_matrix, cmap="hot", shading="auto"
            )
            plt.title("Minimum Effective Index Matrix")
            plt.xlabel("Pump A Frequency (THz)")
            plt.ylabel("Width (um)")
            plt.colorbar()
            plt.show()

        return f_A, f_B, net_momentum_matrix, min_n_matrix

    def scale_neffs(self, new_height, plot=False, plot_mode_num=0):
        """
        Scale the frequencies and the widths to a new size, using the height of the layer at the given index
        inputs:
            new_height (float): the new height
        """
        scale_factor = new_height / self.layer_height

        self.layer_widths = self.layer_widths * scale_factor
        self.layer_height = self.layer_height * scale_factor
        self.freqs = self.freqs / scale_factor
        if plot:
            self.plot(plot_mode_num)

    def get_resonances(
        self, length, mode_num=0, plot=False, poly_order=3, fit_tol=2e-3
    ):
        """Returns the resonances for the given length
        inputs:
            length (float): the length of the resonator
        outputs:
            resonances (dict): Dictionary mapping width index to array of resonance frequencies and modes
                resonances[w_i]['resonance_freqs'] (array): the resonance frequencies
                resonances[w_i]['resonance_modes'] (array): the resonance modes
        """

        # Calculate accumulated phase for each width and frequency

        # Ensure self.freqs is a 1D array of frequencies (Hz)
        freqs = np.array(self.freqs)
        widths = np.array(self.layer_widths)
        n_effs = np.array(self.neffs)  # shape: (num_widths, num_freqs, num_modes)

        # Only use the selected mode
        n_effs_mode = n_effs[:, :, mode_num]  # shape: (num_widths, num_freqs)

        # Calculate wavelength for each frequency
        wavelengths = c / freqs  # shape: (num_freqs,)

        # Calculate accumulated phase for each width and frequency
        # shape: (num_widths, num_freqs)
        accumulated_phase = (
            2 * np.pi * n_effs_mode * length / wavelengths[np.newaxis, :]
        )

        # Mode number as a float
        resonance_number = accumulated_phase / (
            2 * np.pi
        )  # shape: (num_widths, num_freqs)

        resonances = {}

        if plot:
            plt.figure(figsize=(5, 4.5), dpi=200)
            colors = cm.viridis(np.linspace(0, 1, len(widths)))

        for w_i, width in enumerate(widths):
            # Find integer mode numbers within the range
            min_mode = np.ceil(np.min(resonance_number[w_i, :]))
            max_mode = np.floor(np.max(resonance_number[w_i, :]))
            resonance_modes = np.arange(min_mode, max_mode + 1)
            # Interpolate frequency at these integer mode numbers
            interp_func, interp_err = Polynomial.fit(
                resonance_number[w_i, :], freqs / 1e12, poly_order, full=True
            )
            if np.abs(interp_err[0]) > fit_tol:
                print(
                    f"Warning: interp_err {np.abs(interp_err[0][0]):.3e} is too large for width {width:.2f} um"
                )
            resonance_freqs = interp_func(resonance_modes) * 1e12
            resonances[w_i] = {}
            resonances[w_i]["resonance_freqs"] = resonance_freqs
            resonances[w_i]["resonance_modes"] = resonance_modes

            if plot:
                resonance_wavelengths = c / resonance_freqs
                resonance_neffs = resonance_modes / length * resonance_wavelengths
                plt.scatter(
                    resonance_freqs / 1e12,
                    resonance_neffs,
                    color=colors[w_i],
                    label=f"w = {width:.2f} $\\mu$m",
                    s=1,
                )
                # plt.plot(freqs/1e12, resonance_number[w_i, :], color=colors[w_i])
                # plt.scatter(resonance_freqs / 1e12, resonance_modes, color=colors[w_i], label=f'Width={width:.2f} um', s=1)

        if plot:
            plt.xlabel("Resonance Frequency (THz)", fontsize=15)
            plt.ylabel("Effective Index", fontsize=15)
            # plt.title(f'Resonances vs Mode Number (length={length*1e6:.2f} um)', fontsize=15)
            plt.legend(title="Waveguide Width")
            plt.tight_layout()
            plt.show()

        return resonances

def get_waveguide_modes(
    layer_widths,
    layer_heights,
    freqs,
    num_modes,
    layer_mats,
    bend_radius,
    material_library,
    box_mat="air",
    cladding="air",
    target_neff=None,
    plot_geom=True,
    return_fields=True,
    return_n2=False,
    web_modes=False,
):
    """Returns the modes for the given layer widths, heights, and frequencies
    inputs:
        layer_widths (Hx1 array): a list of swept widths for each layer
        layer_heights (Hx1 array): heights of the layers
        freqs (Nx1 array): list of frequencies
        num_modes (int): number of modes to solve for
        layer_mats (Hx1 array): materials of the layers
        bend_radius (float): radius of the bend
        material_library (dict): dictionary of materials
        box_mat (str): material of the box
        cladding (str): material of the cladding
        target_neff (float): target effective index
        plot_geom (bool): whether to plot the mode solver geometry
        return_fields (bool): whether to return the fields
        return_n2 (bool): whether to return the n2 of the mode
        web_modes (bool): whether to run the modes in the web app
    outputs:
        mode_data (ModeData): the mode data
        mode_solver (ModeSolver): the mode solver
        n2 (np.complex64): the n2 of the mode, if return_n2 is True
    """
    # central frequency
    fmin = freqs[0]
    fmax = freqs[-1]
    freq0 = (fmin + fmax) / 2

    # automatic grid specification
    run_time = 1e-12
    grid_spec = td.GridSpec.auto(min_steps_per_wvl=25, wavelength=td.C_0 / freq0)

    bend_axis = 1
    structures = []
    total_height = np.sum(layer_heights)
    weighted_width = np.sum(layer_widths * layer_heights) / total_height

    Lx, Lz = 2 + weighted_width, 2 + total_height
    z_min = -total_height / 2

    for i, (width, height, mat) in enumerate(
        zip(layer_widths, layer_heights, layer_mats)
    ):
        z_center = z_min + height / 2
        layer = td.Structure(
            geometry=td.Box(
                size=(width, td.inf, height), center=(bend_radius, 0, z_center)
            ),
            medium=material_library[mat],
        )
        structures.append(layer)
        z_min += height

    # Add the box
    box_thickness = (Lz) / 2 + 2
    box_z = -total_height / 2 - box_thickness / 2
    box = td.Structure(
        geometry=td.Box(
            size=(Lx + 1, td.inf, box_thickness), center=(bend_radius, 0, box_z)
        ),
        medium=material_library[box_mat],
    )
    structures.append(box)

    sim = td.Simulation(
        size=(Lx, 1, Lz),
        center=(bend_radius if bend_radius else 0, 0, 0),
        grid_spec=grid_spec,
        structures=structures,
        run_time=run_time,
        medium=material_library[cladding],
        boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()),
    )
    mode_spec = td.ModeSpec(
        num_modes=num_modes,
        target_neff=target_neff,
        bend_radius=bend_radius if bend_radius else None,
        bend_axis=bend_axis,
        track_freq="highest",
        filter_pol="te",
    )
    plane = td.Box(center=(bend_radius, 0, 0), size=(Lx, 0, Lz))
    mode_solver = ModeSolver(
        simulation=sim,
        plane=plane,
        mode_spec=mode_spec,
        freqs=freqs,
        fields=["Ex", "Ey", "Ez"] if return_fields else [],
    )
    if plot_geom:
        mode_solver.plot()

    if web_modes:
        mode_data = td.web.run(mode_solver, "bend_modes")
    else:
        mode_data = mode_solver.solve()
    if not return_n2:
        return mode_data, mode_solver
    else:
        freq = np.mean(freqs)
        coords = td.Coords(
            x=mode_solver.grid_snapped.centers.x,
            y=[-1, 1],
            z=mode_solver.grid_snapped.centers.z,
        )
        grid = td.Grid(boundaries=coords)
        n2 = n2_on_grid(sim=sim, grid=grid, coord_key="centers", freq=freq)
        return mode_data, mode_solver, n2


def get_n2(structure: Structure, frequency: float, coords: Coords):
    """Returns the n2 of the structure at the given frequency and coordinates"""

    n2_vals = {
        "diamond": 0.082e-18,
        "nitride": 0.250e-18,
        "air": 0.0,
        "oxide": 0.0,
        "reflector": 0.0,
        None: 0.0,
    }
    return n2_vals[structure.medium.name]


def n2_on_grid(
    sim: AbstractYeeGridSimulation,
    grid: Grid,
    coord_key: str = "centers",
    freq: float = None,
    n2_fn: callable = get_n2,
) -> xr.DataArray:
    """Get array of permittivity at a given freq on a given grid.

    Parameters
    ----------
    sim : :class:`.AbstractYeeGridSimulation`
        Simulation object to get the permittivity from.
    grid : :class:`.Grid`
        Grid specifying where to measure the permittivity.
    coord_key : str = 'centers'
        Specifies at what part of the grid to return the permittivity at.
        Accepted values are ``{'centers', 'boundaries', 'Ex', 'Ey', 'Ez', 'Exy', 'Exz', 'Eyx',
        'Eyz', 'Ezx', Ezy'}``. The field values (eg. ``'Ex'``) correspond to the corresponding field
        locations on the yee lattice. If field values are selected, the corresponding diagonal
        (eg. ``eps_xx`` in case of ``'Ex'``) or off-diagonal (eg. ``eps_xy`` in case of ``'Exy'``) epsilon
        component from the epsilon tensor is returned. Otherwise, the average of the main
        values is returned.
    freq : float = None
        The frequency to evaluate the mediums at.
        If not specified, evaluates at infinite frequency.
    Returns
    -------
    xarray.DataArray
        Datastructure containing the relative permittivity values and location coordinates.
        For details on xarray DataArray objects,
        refer to `xarray's Documentation <https://tinyurl.com/2zrzsp7b>`_.
    """

    def make_n2_data(coords: Coords):
        """returns epsilon data on grid of points defined by coords"""
        arrays = (np.array(coords.x), np.array(coords.y), np.array(coords.z))
        n2_background = n2_fn(
            structure=sim.scene.background_structure, frequency=freq, coords=coords
        )
        shape = tuple(len(array) for array in arrays)
        n2_array = n2_background * np.ones(shape, dtype=np.float64)
        # replace 2d materials with volumetric equivalents
        for structure in sim.volumetric_structures:
            # Indexing subset within the bounds of the structure

            inds = structure.geometry._inds_inside_bounds(*arrays)

            # Get permittivity on meshgrid over the reduced coordinates
            coords_reduced = tuple(arr[ind] for arr, ind in zip(arrays, inds))
            if any(coords.size == 0 for coords in coords_reduced):
                continue

            red_coords = Coords(**dict(zip("xyz", coords_reduced)))
            n2_structure = n2_fn(structure=structure, frequency=freq, coords=red_coords)

            # Update permittivity array at selected indexes within the geometry
            is_inside = structure.geometry.inside_meshgrid(*coords_reduced)
            n2_array[inds][is_inside] = (n2_structure * is_inside)[is_inside]

        coords = dict(zip("xyz", arrays))
        return xr.DataArray(n2_array, coords=coords, dims=("x", "y", "z"))

    # combine all data into dictionary
    if coord_key[0] == "E":
        # off-diagonal components are sampled at respective locations (eg. `eps_xy` at `Ex`)
        coords = grid[coord_key[0:2]]
    else:
        coords = grid[coord_key]
    return make_n2_data(coords)


def n2_on_coords(
    sim: AbstractYeeGridSimulation,
    coords: Coords,
    freq: float = None,
    n2_fn: callable = get_n2,
) -> xr.DataArray:
    """Get array of permittivity at a given freq on a given grid.

    Parameters
    ----------
    sim : :class:`.AbstractYeeGridSimulation`
        Simulation object to get the permittivity from.
    coords : :class:`.Coords`
        Grid specifying where to measure the permittivity.
    coord_key : str = 'centers'
        Specifies at what part of the grid to return the permittivity at.
        Accepted values are ``{'centers', 'boundaries', 'Ex', 'Ey', 'Ez', 'Exy', 'Exz', 'Eyx',
        'Eyz', 'Ezx', Ezy'}``. The field values (eg. ``'Ex'``) correspond to the corresponding field
        locations on the yee lattice. If field values are selected, the corresponding diagonal
        (eg. ``eps_xx`` in case of ``'Ex'``) or off-diagonal (eg. ``eps_xy`` in case of ``'Exy'``) epsilon
        component from the epsilon tensor is returned. Otherwise, the average of the main
        values is returned.
    freq : float = None
        The frequency to evaluate the mediums at.
        If not specified, evaluates at infinite frequency.
    Returns
    -------
    xarray.DataArray
        Datastructure containing the relative permittivity values and location coordinates.
        For details on xarray DataArray objects,
        refer to `xarray's Documentation <https://tinyurl.com/2zrzsp7b>`_.
    """

    def make_n2_data(coords: Coords):
        """returns epsilon data on grid of points defined by coords"""
        arrays = (np.array(coords.x), np.array(coords.y), np.array(coords.z))
        n2_background = n2_fn(
            structure=sim.scene.background_structure, frequency=freq, coords=coords
        )
        shape = tuple(len(array) for array in arrays)
        n2_array = n2_background * np.ones(shape, dtype=np.float64)
        # replace 2d materials with volumetric equivalents
        for structure in sim.volumetric_structures:
            # Indexing subset within the bounds of the structure
            if n2_fn(structure=structure, frequency=freq, coords=coords) != 0:
                inds = structure.geometry._inds_inside_bounds(*arrays)

                # Get permittivity on meshgrid over the reduced coordinates
                coords_reduced = tuple(arr[ind] for arr, ind in zip(arrays, inds))
                if any(coords.size == 0 for coords in coords_reduced):
                    continue

                red_coords = Coords(**dict(zip("xyz", coords_reduced)))
                n2_structure = n2_fn(
                    structure=structure, frequency=freq, coords=red_coords
                )

                # Update permittivity array at selected indexes within the geometry
                is_inside = structure.geometry.inside_meshgrid(*coords_reduced)
                n2_array[inds][is_inside] = (n2_structure * is_inside)[is_inside]

        coords = dict(zip("xyz", arrays))
        return xr.DataArray(n2_array, coords=coords, dims=("x", "y", "z"))

    return make_n2_data(coords)


def fpm_spectrum(
    resonances, nonlinear_env, net_momentum_env, w_i, plot=False, mlim_a=None
):
    """Compute pump resonance pairs and output frequencies for four-wave mixing.

    Inputs:
        resonances (dict): output of get_resonances(). Use `resonances[w_i]`.
        nonlinear_env (dict): contains keys:
            - 'fixed_freqs': array-like of fixed mode frequencies (Hz) with length K-2
            - 'mode_factors': array-like of length K (integers, e.g. [+1, +1, -1, -1])
            - 'unit_ks': array-like of length K (direction indicators in {-1, 0, +1})
        net_momentum_env (dict): included for API symmetry; not used directly here.
        w_i (int): width index selecting which resonance set to use.
        plot (bool): if True, plot m_pump_a and m_pump_b vs output frequency.
        mlim_a (tuple): if not None, restrict the m_pump_a to the given limits.
    Returns:
        m_pump_a (list[int]): resonance numbers for pump A.
        m_pump_b (list[int]): resonance numbers for pump B.
        output_freqs (np.ndarray): computed output frequencies (Hz) for each pair.
    """
    # Extract resonance mapping for the specified width index
    width_res = resonances[w_i]
    res_modes = np.asarray(width_res["resonance_modes"]).astype(int)
    res_freqs = np.asarray(width_res["resonance_freqs"])  # Hz

    # Build lookup from mode number to frequency
    mode_to_freq = {int(m): f for m, f in zip(res_modes, res_freqs)}

    fixed_freqs = np.asarray(net_momentum_env["fixed_freqs"])  # length K-2
    mode_factors = np.asarray(nonlinear_env["mode_factors"])
    unit_ks = np.asarray(nonlinear_env["unit_ks"])

    if len(mode_factors) != len(unit_ks):
        raise ValueError("mode_factors and unit_ks must have the same length")
    if len(fixed_freqs) != len(mode_factors) - 2:
        raise ValueError("fixed_freqs length must be len(mode_factors) - 2")

    num_fixed = len(fixed_freqs)

    # For each fixed frequency, find the nearest resonance number by frequency proximity
    fixed_res_numbers = []
    for f in fixed_freqs:
        idx = int(np.argmin(np.abs(res_freqs - f)))
        fixed_res_numbers.append(int(res_modes[idx]))
    fixed_res_numbers = np.asarray(fixed_res_numbers, dtype=int)
    # Compute leftover resonance number required to satisfy momentum conservation in mode-number space
    leftover_resonance_number = int(
        np.sum(fixed_res_numbers * mode_factors[:num_fixed] * unit_ks[:num_fixed])
    )
    # Define coefficients for pumps A and B in the resonance-number equation:
    # m_a * (mode_factors[-2] * unit_ks[-2]) + m_b * (mode_factors[-1] * unit_ks[-1]) = leftover_resonance_number
    coeff_a = int(mode_factors[-2] * unit_ks[-2])
    coeff_b = int(mode_factors[-1] * unit_ks[-1])

    if coeff_a == 0 and coeff_b == 0:
        # No pump contributions possible to satisfy leftover; no valid pairs
        return [], [], np.array([])

    available_modes = set(int(m) for m in res_modes.tolist())
    m_pump_a = []
    m_pump_b = []

    # Search all feasible pump pairs within available resonance numbers
    # Solve for m_b given m_a when possible, or vice versa.
    if coeff_b != 0:
        for m_a in available_modes:
            rhs = leftover_resonance_number + coeff_a * m_a
            m_b = -rhs // coeff_b
            if m_b in available_modes:
                m_pump_a.append(int(m_a))
                m_pump_b.append(int(m_b))
    else:
        # coeff_b == 0, require coeff_a * m_a == leftover
        if coeff_a != 0 and (leftover_resonance_number % coeff_a == 0):
            m_a_solution = leftover_resonance_number // coeff_a
            if m_a_solution in available_modes:
                # Any m_b in available_modes is acceptable; pair each with the single m_a
                for m_b in available_modes:
                    m_pump_a.append(int(m_a_solution))
                    m_pump_b.append(int(m_b))
    # Compute output frequency for each valid pump pair using the energy relation
    # f_out = sum_j mode_factor[j] * abs(unit_k[j]) * f_j (fixed + pumps)
    # Fixed modes contribution (constant shift)
    fixed_contrib = float(
        np.sum(mode_factors[:num_fixed] * np.abs(unit_ks[:num_fixed]) * fixed_freqs)
    )

    output_freqs = []
    f_as = []
    f_bs = []
    for m_a, m_b in zip(m_pump_a, m_pump_b):
        # Lookup pump frequencies from resonance numbers
        f_a = mode_to_freq[m_a]
        f_b = mode_to_freq[m_b]
        f_as.append(f_a)
        f_bs.append(f_b)
        pump_contrib = (mode_factors[-2] * np.abs(unit_ks[-2]) * f_a) + (
            mode_factors[-1] * np.abs(unit_ks[-1]) * f_b
        )
        output_freqs.append(fixed_contrib + pump_contrib)
    f_as = np.asarray(f_as)
    f_bs = np.asarray(f_bs)
    output_freqs = np.asarray(output_freqs)

    # Crop the data to the limits of mlim_a
    if mlim_a is not None:
        idx_min = np.argmin(np.abs(np.array(m_pump_a) - mlim_a[0]))
        idx_max = np.argmin(np.abs(np.array(m_pump_a) - mlim_a[1]))
        m_pump_a = m_pump_a[idx_min : idx_max + 1]
        m_pump_b = m_pump_b[idx_min : idx_max + 1]
        output_freqs = output_freqs[idx_min : idx_max + 1]

    if plot and len(output_freqs) > 1:
        fig, ax0 = plt.subplots(1, 1, figsize=(5, 5), dpi=200)
        ax1 = ax0.twiny()
        # ax1.scatter(m_pump_b, output_freqs / 1e12, s=10, color='tab:orange')
        ax1.spines["bottom"].set_position(
            ("outward", 40)
        )  # 40 points below the original
        # if mlim_a is None:
        ax0.set_xlim((m_pump_a[0], m_pump_a[-1]))
        ax1.set_xlim((m_pump_b[0], m_pump_b[-1]))  # Reverse the x-axis direction
        ax0.scatter(
            m_pump_a, c / output_freqs * 1e6, s=400 / len(m_pump_a), color="tab:blue"
        )
        if len(m_pump_a) < 22:
            ax0.set_xticks(m_pump_a)
            ax1.set_xticks(m_pump_b)
        ax0.set_xlabel(r"Mode Number $m_a$", fontsize=15)
        ax1.set_xlabel(r"Mode Number $m_b$", fontsize=15, labelpad=10, loc="center")
        ax1.xaxis.set_label_position("top")
        ax1.xaxis.set_ticks_position("top")
        ax0.set_ylabel("Output Wavelength ($\\mu$m)", fontsize=15)
        # ax0.set_title(f'Output Wavelengths for m([Signal, Idler]) = {fixed_res_numbers * mode_factors[:num_fixed] * unit_ks[:num_fixed]}', fontsize=15)
        plt.tight_layout()
        plt.show()

    return m_pump_a, m_pump_b, f_as, f_bs, output_freqs

def get_neff_mixer(
    swept_layer_widths,
    layer_heights,
    freqs,
    num_modes,
    mode_args,
    chosen_layer_i=0,
    plotted_modes=0,
    plot_fields=False,
    plot_freq=None,
):
    """Return a NeffMixer object for the given layer widths, heights, and frequencies
    inputs:
        swept_layer_widths (HxWx1 array): a list of swept widths for each layer
        layer_heights (Hx1 array): heights of the layers
        freqs (Nx1 array): list of frequencies in increasing order
        num_modes (int): number of modes to solve for
        mode_args (dict): dictionary of mode arguments for get_waveguide_modes
        chosen_layer_i (int): index of the layer to use for the mixer
        plotted_modes (int): number of modes to plot
        plot_fields (bool): whether to plot the fields
    outputs:
        neff_mixer (NeffMixer): the NeffMixer object
    """
    neffs = np.zeros((np.shape(swept_layer_widths)[1], len(freqs), num_modes))
    f_min = freqs[0]
    if plot_freq is None:
        plot_freq = f_min
    if plot_fields:
        print(f"Plotting fields at {plot_freq / 1e12:.2f} THz")
        fig, axs = plt.subplots(
            np.shape(swept_layer_widths)[1],
            2 * plotted_modes,
            figsize=(9 * plotted_modes, 4 * np.shape(swept_layer_widths)[1]),
        )

    for sweep_i in tqdm(range(np.shape(swept_layer_widths)[1])):
        layer_widths = swept_layer_widths[:, sweep_i]
        mode_data, mode_solver = get_waveguide_modes(
            layer_widths, layer_heights, freqs, num_modes, **mode_args
        )
        if plot_fields:
            for i in range(plotted_modes):
                mode_solver.plot_field(
                    "Ex", f=plot_freq, mode_index=i, ax=axs[sweep_i, 2 * i]
                )
                mode_solver.plot_field(
                    "Ez", f=plot_freq, mode_index=i, ax=axs[sweep_i, 2 * i + 1]
                )
                axs[sweep_i, i].set_title(
                    f"Width: {layer_widths[chosen_layer_i]:.2f} um, m={i + 1}"
                )
        neff = mode_data.n_eff
        neff = np.array(neff)
        neffs[sweep_i, :] = neff

    meta_data = {}
    meta_data["swept_layer_widths"] = swept_layer_widths
    meta_data["layer_heights"] = layer_heights
    meta_data["chosen_layer_i"] = chosen_layer_i
    meta_data["layer_mats"] = mode_args["layer_mats"]
    meta_data["bend_radius"] = mode_args["bend_radius"]
    meta_data["box_mat"] = mode_args["box_mat"]
    meta_data["cladding"] = mode_args["cladding"]
    meta_data["target_neff"] = mode_args["target_neff"]

    mixer = NeffMixer(
        freqs,
        neffs,
        num_modes,
        swept_layer_widths[chosen_layer_i, :],
        layer_heights[chosen_layer_i],
        meta_data,
    )
    for i in range(plotted_modes):
        mixer.plot(i)
    return mixer

idler_wavelength = 1.31
bend_radius = 6

######################################################### Mixer parameters #########################################################
layer_mats = ["diamond"]
n_idler = [mats[mat].nk_model(td.C_0 / idler_wavelength)[0] for mat in layer_mats]
layer_heights = [
    0.55
]  # [idler_wavelength/n_idler[i]/len(layer_mats) for i in range(len(layer_mats))]

min_width = 0.65
max_width = 1.05
num_widths = 5
swept_layer_widths = [
    np.linspace(min_width, max_width, num_widths) for _ in range(len(layer_mats))
]
wl_min = 0.6
wl_max = 2.4
f_min = td.C_0 / wl_max
f_max = td.C_0 / wl_min
num_freqs = 26
freqs = np.linspace(f_min, f_max, num_freqs)
freq_plot_i = np.argmin(np.abs(freqs - 146e12))

env = {}
mixer_env = {
    "layer_heights": np.array(layer_heights),
    "swept_layer_widths": np.array(swept_layer_widths),
    "chosen_layer_i": np.array(0),
    "num_modes": np.array(5),
    "freqs": np.array(freqs),
    "plotted_modes": np.array(2),
    "plot_fields": True,
    "plot_freq": freqs[freq_plot_i],
}
env["mixer"] = mixer_env

mode_env = {
    "layer_mats": layer_mats,
    "bend_radius": bend_radius,
    "cladding": "air",
    "box_mat": "oxide",
    "target_neff": None,
    "web_modes": False,
    "plot_geom": False,
    "material_library": mats,
    "return_fields": True,
}
env["mode"] = mode_env
######################################################### Nonlinear Bahar Parameters#########################################################
nonlinear_env = {
    #### Order of the signals is (signal, idler, pump A, pump B)
    "mode_nums": [0, 0, 0, 0],
    "mode_factors": [1, -1, 1, -1],
    "unit_ks": [1, 0, -1, 1],
}
env["nonlinear"] = nonlinear_env

net_momentum_env = {
    "fixed_freqs": td.C_0 / np.array([0.619, idler_wavelength]),
    "fit_tol": 1e-3,
    "poly_order": 15,
    "num_freqs": 1000,
}
env["net_momentum"] = net_momentum_env

fpm_env = {
    "zero_momentum_tolerance": 0.1,
    "pump_freq_nearby": None,
}
env["fpm"] = fpm_env

mn_env = {
    "mlim_a": (0, 50),
    "w_i": 1,
    "poly_order": 20,
}
env["mnumber"] = mn_env

result = {}

def build_layer_structures(
    mats,
    diamond_thickness,
    box_thickness,
    lower_nitride_thickness=0.22,
    upper_nitride_thickness=0.0,
    lower_nitride_top_depth=2.22,
    upper_nitride_bottom_depth=2,
    reflector_diamond_gap=0.87,
    reflector_thickness=0.5,
    ring_radius=6,
    ring_width=0.85,
    window_size=100,
    add_lower_nitride=False,
):
    """Build all td.Structure objects for the diamond ring on nitride/oxide stack.

    Returns:
        structures: list of td.Structure
        layer_z_centers: dict mapping layer name to z center coordinate
    """
    # Compute z positions (same logic as original get_nitride_window_layer_stack)
    zmin_lower_nitride = box_thickness
    zmin_upper_nitride = (
        zmin_lower_nitride
        + lower_nitride_thickness
        + lower_nitride_top_depth
        - upper_nitride_bottom_depth
    )
    zmin_diamond = max(zmin_upper_nitride, zmin_upper_nitride + upper_nitride_thickness)
    oxide_thickness = box_thickness + lower_nitride_thickness + lower_nitride_top_depth

    structures = []

    # 1. Oxide substrate (full slab)
    structures.append(
        td.Structure(
            geometry=td.Box(
                center=(0, 0, oxide_thickness / 2),
                size=(td.inf, td.inf, oxide_thickness),
            ),
            medium=mats["oxide"],
            name="oxide",
        )
    )

    # 2. Lower nitride film (structure only added if add_lower_nitride=True;
    #    lower_nitride_thickness still affects z positions of upper layers)
    if lower_nitride_thickness > 0 and add_lower_nitride:
        structures.append(
            td.Structure(
                geometry=td.Box(
                    center=(0, 0, zmin_lower_nitride + lower_nitride_thickness / 2),
                    size=(td.inf, td.inf, lower_nitride_thickness),
                ),
                medium=mats["nitride"],
                name="lower_nitride",
            )
        )

    # 3. Upper nitride film (may be zero thickness)
    if upper_nitride_thickness > 0:
        structures.append(
            td.Structure(
                geometry=td.Box(
                    center=(0, 0, zmin_upper_nitride + upper_nitride_thickness / 2),
                    size=(td.inf, td.inf, upper_nitride_thickness),
                ),
                medium=mats["nitride"],
                name="upper_nitride",
            )
        )

    # 4. Window (air etch region — centered rectangle)
    half_w = window_size / 2
    window_verts = [
        (-half_w, -half_w),
        (half_w, -half_w),
        (half_w, half_w),
        (-half_w, half_w),
    ]
    structures.append(
        td.Structure(
            geometry=td.PolySlab(
                vertices=window_verts,
                axis=2,
                slab_bounds=(
                    zmin_upper_nitride,
                    zmin_upper_nitride + upper_nitride_bottom_depth,
                ),
            ),
            medium=td.Medium(),  # air
            name="window",
        )
    )

    # 5. Reflector
    if reflector_thickness > 0:
        zmin_refl = zmin_diamond - reflector_diamond_gap - reflector_thickness
        structures.append(
            td.Structure(
                geometry=td.Box(
                    center=(0, 0, zmin_refl + reflector_thickness / 2),
                    size=(td.inf, td.inf, reflector_thickness),
                ),
                medium=mats["reflector"],
                name="reflector",
            )
        )

    # 6. Diamond ring — outer cylinder minus inner cylinder
    z_center_diamond = zmin_diamond + diamond_thickness / 2
    ring_geo = td.ClipOperation(
        operation="difference",
        geometry_a=td.Cylinder(
            center=(0, 0, z_center_diamond),
            radius=ring_radius + ring_width / 2,
            length=diamond_thickness,
            axis=2,
        ),
        geometry_b=td.Cylinder(
            center=(0, 0, z_center_diamond),
            radius=ring_radius - ring_width / 2,
            length=diamond_thickness,
            axis=2,
        ),
    )
    structures.append(
        td.Structure(geometry=ring_geo, medium=mats["diamond"], name="diamond_ring")
    )

    layer_z_centers = {
        "diamond": z_center_diamond,
        "lower_nitride": zmin_lower_nitride + lower_nitride_thickness / 2,
        "upper_nitride": zmin_upper_nitride + upper_nitride_thickness / 2,
        "oxide": oxide_thickness / 2,
    }

    return structures, layer_z_centers


# Build the layer structures
layer_kwargs = {
    "diamond_thickness": layer_heights[0],
    "box_thickness": 10,
    "upper_nitride_thickness": 0.0,
    "lower_nitride_thickness": 0.22,  # used for z-position calculation only; no structure added (add_lower_nitride=False)
    "upper_nitride_bottom_depth": 2,
    "lower_nitride_top_depth": 2.22,
    "reflector_diamond_gap": 0.87,
    "reflector_thickness": 0.5,
}

# Ring and simulation parameters
ring_radius = env["mode"]["bend_radius"]
ring_width = env["mixer"]["swept_layer_widths"][0][env["mnumber"]["w_i"]]

# Build structures using cylinder-based ring
structures, layer_z_centers = build_layer_structures(
    mats=mats, ring_radius=ring_radius, ring_width=ring_width, **layer_kwargs
)

# Port at (0, -radius) pointing in +x direction (same as original orientation=0)
port_z_center = layer_z_centers["diamond"]
port_size_mult = (4, 2.5)
port_info = {
    "o1": {
        "center": (0, -ring_radius, port_z_center),
        "size": (0, ring_width * port_size_mult[0], ring_width * port_size_mult[1]),
        "direction": "-",
        "mode_spec": td.ModeSpec(
            num_modes=10,
            filter_pol="te",
            bend_radius=-ring_radius,
            bend_axis=1,
            num_pml=(5, 5),
        ),
    }
}

# Frequency range for modeler (used for default bandwidth calculations)
modeler_bandwidth = 0.02  # um
modeler_num_freqs = 3
modeler_wavelength = 1.55  # original get_component_modeler default
# Fixed modeler freqs (original always used wavelength=1.55)
modeler_freqs = td.C_0 / np.linspace(
    modeler_wavelength - modeler_bandwidth / 2,
    modeler_wavelength + modeler_bandwidth / 2,
    modeler_num_freqs,
)

# Base simulation
base_sim = td.Simulation(
    size=(2 * ring_radius + 4, 2 * ring_radius + 4, 4),
    center=(0, 0, port_z_center),
    structures=structures,
    grid_spec=td.GridSpec.auto(wavelength=idler_wavelength, min_steps_per_wvl=20),
    boundary_spec=td.BoundarySpec(
        x=td.Boundary(minus=td.PML(), plus=td.PML()),
        y=td.Boundary(minus=td.PML(), plus=td.PML()),
        z=td.Boundary(minus=td.PML(), plus=td.PML()),
    ),
    run_time=2e-12,
    symmetry=(0, 0, 0),
)

env["port_info"] = port_info
env["layer_z_centers"] = layer_z_centers
env["base_sim"] = base_sim
env["modeler_freqs"] = modeler_freqs

env["custom_fdtd"] = {}
env["chi3_reslock"] = {
    "downsample": (2, 2, 2),
    "mode_volume_layer": "diamond",
    "ff_res": 60,
    "amp_a": 1000,
    "amp_b": 1000,
    "Q_threshold": 1e3,
}