wavelength = 1550 * nm
freq = td.C_0 / wavelength
fwidth = freq / 10
run_time = 50 / fwidth

# device parameters
thick_substrate = 2.0 * um
thick_Si = 220 * nm
thick_SiN = 220 * nm
thick_etch = 220 * nm
space_SiSiN = 100 * nm
width_waveguide = 1000 * nm

# size of design region
size_design_x = 8 * um
size_design_y = 8 * um

# spacings between things (technically infinite, but finite for sim)
length_waveguide = 1 * wavelength
thick_Si_substrate = 550 * nm
space_cladding_top = 1 * wavelength
space_design_edge = 1 * wavelength

# compute sim size
size_sim_x = length_waveguide + size_design_x + space_design_edge
size_sim_y = space_design_edge + size_design_y + space_design_edge
size_sim_z = (
    thick_Si_substrate
    + thick_substrate
    + thick_Si
    + space_SiSiN
    + thick_SiN
    + space_cladding_top
)
size_sim = (size_sim_x, size_sim_y, size_sim_z)

# references to various locations
sim_top_z = size_sim_z / 2.0
sim_bot_z = -size_sim_z / 2.0
center_SiN_z = sim_top_z - space_cladding_top - thick_SiN / 2.0
center_Si_z = sim_top_z - space_cladding_top - thick_SiN - space_SiSiN - thick_Si / 2.0
center_Si_etch_z = (
    sim_top_z - space_cladding_top - thick_SiN - space_SiSiN - thick_etch / 2.0
)
center_source_z = sim_top_z - space_cladding_top / 2.0
center_Si_substrate_z = sim_bot_z + thick_Si_substrate / 2

center_waveguide_x = -size_sim_x / 2 + length_waveguide / 2.0
center_design_x = size_sim_x / 2.0 - space_design_edge - size_design_x / 2.0

# dispersionless materials for now. Note: can also do `td.Medium.from_nk(n=n, k=k, freq=freq)`

n_SiO2 = 1.44
n_SiN = 2.0
n_Si = 3.5

eps_Si = n_Si**2
eps_SiN = n_SiN**2
eps_SiO2 = n_SiO2**2

medium_Si = td.Medium(permittivity=eps_Si)
medium_SiN = td.Medium(permittivity=eps_SiN)
medium_SiO2 = td.Medium(permittivity=eps_SiO2)

spot_size = np.sqrt(size_design_x**2 + size_design_y**2)
tilt_angle_deg = 0.0

waveguide = td.Structure(
    geometry=td.Box(
        center=(-size_sim_x + length_waveguide + size_design_x / 2.0, 0, center_Si_z),
        size=(size_sim_x, width_waveguide, thick_Si),
    ),
    medium=medium_SiN,
)

design_region = td.Structure(
    geometry=td.Box(
        center=(center_design_x, 0, center_Si_z),
        size=(size_design_x, size_design_y, thick_Si),
    ),
    medium=medium_SiN,
)

slab_SiN = td.Structure(
    geometry=td.Box(
        center=(center_design_x, 0, center_SiN_z),
        size=(size_design_x, size_design_y, thick_SiN),
    ),
    medium=medium_SiN,
)

substrate_Si = td.Structure(
    geometry=td.Box(
        center=(0, 0, -size_sim_z + thick_Si_substrate),
        size=(td.inf, td.inf, size_sim_z),
    ),
    medium=medium_Si,
)

gaussian_beam = td.GaussianBeam(
    center=(center_design_x, 0, center_source_z),
    size=(size_design_x, size_design_y, 0),
    source_time=td.GaussianPulse(freq0=freq, fwidth=fwidth),
    pol_angle=np.pi / 2,
    angle_theta=tilt_angle_deg * np.pi / 180.0,
    direction="-",
    num_freqs=7,
    waist_radius=spot_size / 2,
)

mode_monitor = td.ModeMonitor(
    size=(0, width_waveguide * 6, thick_Si * 6),
    center=(center_waveguide_x, 0, center_Si_z),
    name="mode",
    freqs=[freq],
    mode_spec=td.ModeSpec(),
)

field_monitor_xy = td.FieldMonitor(
    center=(0, 0, center_Si_z),
    size=(td.inf, td.inf, 0),
    freqs=[freq],
    name="field_xy",
)

field_monitor_xz = td.FieldMonitor(
    center=(0, 0, 0),
    size=(td.inf, 0, td.inf),
    freqs=[freq],
    name="field_xz",
)

sim_no_etch = td.Simulation(
    size=size_sim,
    run_time=run_time,
    medium=medium_SiO2,
    structures=[waveguide, design_region, slab_SiN, substrate_Si],
    grid_spec=td.GridSpec.auto(wavelength=wavelength),
    sources=[gaussian_beam],
    monitors=[mode_monitor, field_monitor_xy, field_monitor_xz],
)

from tidy3d.plugins.autograd import make_filter_and_project, rescale

# resolution of the design region pixels
pixel_size_Si = 15 * nm
pixel_size_SiN = 15 * nm

# radius of the circular filter (um) (higher = larger features)
# here the feature size could be defined to 150 nm
radius_Si = 100 * nm
radius_SiN = 100 * nm

# projection strengths (higher = more binarized)
beta_Si = 5
beta_SiN = 5
beta_penalty = 5

nx_Si = int(size_design_x // pixel_size_Si)
ny_Si = int(size_design_y // pixel_size_Si)
nx_SiN = int(size_design_x // pixel_size_SiN)
ny_SiN = int(size_design_y // pixel_size_SiN)

etch_Si_geometry = geometry = td.Box(
    size=(size_design_x, size_design_y, thick_etch),
    center=(center_design_x, 0, center_Si_etch_z),
)


def get_density(
    params: np.ndarray, radius: float, beta: float, pixel_size: float
) -> np.ndarray:
    """Generic function to get the etch density as a function of design parameters, using filter and projection."""
    filter_project = make_filter_and_project(radius, pixel_size, beta=beta)
    return filter_project(params)


def get_density_Si(params: np.ndarray) -> td.Structure:
    """Get the density of the Si etch as a function of its design parameters."""
    return get_density(
        params=params, radius=radius_Si, beta=beta_Si, pixel_size=pixel_size_Si
    )


def get_density_SiN(params: np.ndarray) -> td.Structure:
    """Get the density of the SiN etch as a function of its design parameters."""
    return get_density(
        params=params, radius=radius_SiN, beta=beta_SiN, pixel_size=pixel_size_SiN
    )


def get_permittivity(density: np.ndarray, eps_max: float) -> np.ndarray:
    """Get the permittivity array between SiO2 and the max, given an array of material densities between 0 and 1."""
    return rescale(density, eps_SiO2, eps_max)


def make_etch_structure(
    density: np.ndarray, eps_max: float, geometry: td.Box
) -> td.Structure:
    """Make a `td.Structure` containing a `td.CustomMedium` corresponding to this density array, given a geometry."""

    permittivity_data = get_permittivity(density, eps_max=eps_max)

    rmin, rmax = geometry.bounds
    coords = {}
    for key, pt_min, pt_max, num_pts in zip("xyz", rmin, rmax, density.shape):
        coord_edges = np.linspace(pt_min, pt_max, num_pts + 1)
        coord_centers = (coord_edges[1:] + coord_edges[:-1]) / 2.0
        coords[key] = coord_centers

    custom_medium = td.CustomMedium(
        permittivity=td.ScalarFieldDataArray(
            permittivity_data,
            coords=coords,
        )
    )

    return td.Structure(geometry=geometry, medium=custom_medium)


# change the geometry
def get_etch_Si(params: np.ndarray) -> td.Structure:
    """Get the etch region for Si, using the Si parameters."""
    density = get_density_Si(params)
    return make_etch_structure(
        density=density, eps_max=eps_SiN, geometry=etch_Si_geometry
    )


def get_etch_SiN(params: np.ndarray) -> td.Structure:
    """Get the etch region for SiN, using the SiN parameters."""
    density = get_density_SiN(params)
    return make_etch_structure(
        density=density, eps_max=eps_SiN, geometry=slab_SiN.geometry
    )

def make_sim_with_etch(params_Si: np.ndarray, params_SiN: np.ndarray) -> td.Simulation:
    """Make a simulation as a function of the design parameters for Si and SiN etch regions."""
    etch_Si = get_etch_Si(params_Si)
    etch_SiN = get_etch_SiN(params_SiN)

    # add uniform mesh override structures to simulation
    design_region_mesh_Si = td.MeshOverrideStructure(
        geometry=design_region.geometry,
        dl=[pixel_size_Si] * 3,
        enforce=True,
    )

    design_region_mesh_SiN = td.MeshOverrideStructure(
        geometry=slab_SiN.geometry,
        dl=[pixel_size_SiN] * 3,
        enforce=True,
    )

    grid_spec = sim_no_etch.grid_spec.updated_copy(
        override_structures=list(sim_no_etch.grid_spec.override_structures)
        + [design_region_mesh_Si, design_region_mesh_SiN]
    )

    return sim_no_etch.updated_copy(
        structures=list(sim_no_etch.structures) + [etch_Si, etch_SiN],
        grid_spec=grid_spec,
    )

from tidy3d.plugins.autograd import make_erosion_dilation_penalty

penalty_fn_Si = make_erosion_dilation_penalty(
    radius_Si, pixel_size_Si, beta=beta_penalty
)
penalty_fn_SiN = make_erosion_dilation_penalty(
    radius_SiN, pixel_size_SiN, beta=beta_penalty
)


def penalty_Si(params: np.ndarray) -> float:
    """Define the erosion dilation invariance penalty for Si density parameters."""
    density = get_density_Si(params)
    return penalty_fn_Si(density)


def penalty_SiN(params: np.ndarray) -> float:
    """Define the erosion dilation invariance penalty for SiN density parameters."""
    density = get_density_SiN(params)
    return penalty_fn_SiN(density)


def coupling_efficiency(sim_data: td.SimulationData) -> float:
    """Coupling efficiency into the waveguide as a function of the simulation output data."""
    output_amps = sim_data["mode"].amps
    amp = output_amps.sel(direction="-", f=freq, mode_index=0).values
    return np.sum(np.abs(amp) ** 2)

def objective(
    params_Si: np.ndarray,
    params_SiN,
    verbose: bool = False,
    include_field_mnts: bool = False,
) -> float:
    """Combined objective function over the full set of parameters."""

    # coupling efficiency calculation, through a differentiable simulation
    sim = make_sim_with_etch(params_Si, params_SiN)
    if not include_field_mnts:
        sim = sim.updated_copy(monitors=[mode_monitor])
    data = web.run(sim, task_name="coupler", verbose=verbose)
    efficiency = coupling_efficiency(data)

    # fabrication penalty calculation for both design regions
    penalty_val_Si = penalty_Si(params_Si)
    penalty_val_SiN = penalty_SiN(params_SiN)

    # total objective (to maximize)
    coupling_objective = 2 * efficiency
    penalty_objective = 0.5 * (penalty_val_Si + penalty_val_SiN)
    return coupling_objective - penalty_objective

val_grad_fn = autograd.value_and_grad(objective, argnum=(0, 1))

params_final = params_history[-3]
objective_value_final = objective(*params_final)
objective_values_history.append(objective_value_final)

# optionally save the coupler history if we want to load it later.
# import pickle
# with open("coupler.pkl", "wb") as file:
#    pickle.dump(params_history, file)