tech = pf.basic_technology(slab_thickness=0.15)
pf.config.default_technology = tech

tech.layers

tech.extrusion_specs

# An empty mask covers the whole device bounds
air_extrusion = pf.ExtrusionSpec(
    pf.MaskSpec(), td.Medium(permittivity=1.0, name="air"), limits=(0, pf.Z_INF)
)

# Insert air volume first, so that all core extrusions override the air clad
tech.insert_extrusion_spec(0, air_extrusion)

# Print the updated extrusion specs
tech.extrusion_specs

# Setup parameters
wg_width = 6
wg_length = 30
fiber_angle = 10
fiber_translation = 3
waist_radius = 6.0
gaussian_port_height = 0.5

# Define the gaussian port input vector based on the fiber angle
sin_angle = np.sin(fiber_angle / 180 * np.pi)
cos_angle = np.cos(fiber_angle / 180 * np.pi)
input_vector = (-sin_angle, 0, -cos_angle)

# Create an empty Photonforge component
background = pf.Component()

# Add the waveguide layer
wg = pf.Rectangle((0, -wg_width / 2), (wg_length, wg_width / 2))
background.add("WG_CORE", wg)

# Add the waveguide port
input_wg_portspec = pf.PortSpec(
    "input wg",
    width=wg_width,
    limits=tech.ports["Strip"].limits,
    num_modes=1,
    target_neff=3.5,
    path_profiles={(wg_width, 0, tech.layers["WG_CORE"].layer)},
)
background.add_port(pf.Port((0, 0), 0, input_wg_portspec))

# Add the fiber (gaussian) port
gaussian_port_center = (wg_length / 2 + fiber_translation, 0, gaussian_port_height)
background.add_port(
    pf.GaussianPort(
        gaussian_port_center,
        input_vector,
        waist_radius,
        polarization_angle=90,
        field_tolerance=1e-2,
    )
)

background

# Add a Tidy3D simulation model to the component
background.add_model(pf.Tidy3DModel(), "Tidy3D")

wavelengths = np.linspace(1.5, 1.6, 21)
sims = background.models["Tidy3D"].get_simulations(background, pf.C_0 / wavelengths)
basesim = sims["P1@0"]  # This will be the base simulation that we work with

basesim.plot(y=0)
basesim.plot_eps(y=0)
plt.show()

# Update the simulation size so that it is 2D
size = (basesim.size[0], 0, basesim.size[2])

# Update the boundary spec for 2D simulations
boundary_spec = basesim.boundary_spec.updated_copy(
    y=td.Boundary(minus=td.Periodic(), plus=td.Periodic()),
)

# Update the monitor size
monitors = []
for monitor in basesim.monitors:
    # We will remove the field monitor since it's not needed during optimization
    if monitor.type == "ModeMonitor":
        mon_size = (monitor.size[0], td.inf, monitor.size[2])
        monitors.append(monitor.updated_copy(size=mon_size))

basesim = basesim.updated_copy(
    size=size,
    boundary_spec=boundary_spec,
    monitors=monitors,
)

print(basesim.size)

period = 0.63
gc_len = 15.5
ncells = int(round(gc_len / period))
fill_factor = 0.5

gc_widths = np.ones(shape=ncells) * period * fill_factor
gc_gaps = period - gc_widths

# Extract the layers and their bounds from the extrusion specs
extrusion_layers = [ex.mask_spec.layer for ex in tech.extrusion_specs]
extrusion_bounds = [ex.limits for ex in tech.extrusion_specs]

# Get the slab, core, and etch bounds
slabbounds = extrusion_bounds[extrusion_layers.index(tech.layers["SLAB"].layer)]
corebounds = extrusion_bounds[extrusion_layers.index(tech.layers["WG_CORE"].layer)]
etch_zbounds = (slabbounds[1], corebounds[1])

def gc_sim(
    gc_gaps: np.ndarray,
    gc_widths: np.ndarray,
    basesim: td.Simulation = basesim,
    gc_startx: float = 5,
):
    """Returns a 2D grating simulation for optimization.

    Params:
    -------
    gc_gaps : np.ndarray
        Widths of the partially-etched regions
    gc_widths : np.ndarray
        Widths of the unetched regions
    basesim : td.Simulation
        Base simulation to update
    gc_startx : float
        The position of where the grating begins.
    """

    # Add the grating teeth to the simulation
    grating_structures = []
    cx = gc_startx
    for w, g in zip(gc_widths, gc_gaps):
        cx += g / 2
        grating_structures.append(
            td.Structure(
                geometry=td.Box.from_bounds(
                    rmin=(cx - g / 2, -td.inf, etch_zbounds[0]),
                    rmax=(cx + g / 2, td.inf, etch_zbounds[1]),
                ),
                medium=td.Medium(permittivity=1.0),
            )
        )
        cx += g / 2 + w
    structures = list(basesim.structures) + grating_structures

    # Etch away the remaining waveguide layer
    structures.append(
        td.Structure(
            geometry=td.Box.from_bounds(
                rmin=(cx, -td.inf, corebounds[0]), rmax=(1e3, td.inf, corebounds[1])
            ),
            medium=td.Medium(permittivity=1.0),
        )
    )

    return basesim.updated_copy(structures=structures)

# test gc sim
testsim = gc_sim(
    gc_gaps,
    gc_widths,
)

testsim.plot(y=0)
plt.show()

sim_data = td.web.run(testsim, task_name="testgc", path="simdata/data.hdf5", verbose=False)

# Plot transmission
amps = sim_data["P0"].amps.sel(direction="-")
T = amps.abs**2
f, ax = plt.subplots()
ax.plot(td.C_0 / T.f, T)
ax.set_xlabel("wavelength (um)")
ax.set_ylabel("transmission")
ax.grid()

# sweep source position for best coupling
fiber_pos = np.linspace(5, 15, 11)

# Generate simulations with updated fiber source positions
sims = []
for fpos in fiber_pos:
    source = testsim.sources[0]
    source = source.updated_copy(center=(fpos, source.center[1], source.center[2]))
    sims.append(testsim.updated_copy(sources=[source]))

# Consolidate simulations into a dict for passing to td.web.run
gcsims = {f"fiber_pos{fpos}": sim for fpos, sim in zip(fiber_pos, sims)}

# Run the batch of simulations.
batchdata = td.web.run(gcsims, path="simdata", verbose=False)

# Get transmission at the center wavelength vs. source position
centerwl = 1.55
T_centerwl = []
for key, simdata in batchdata.items():
    amps = simdata["P0"].amps.sel(direction="-")
    T = amps.abs**2
    T_centerwl.append(T.interp(f=td.C_0 / centerwl).data[0])

f, ax = plt.subplots()
ax.plot(fiber_pos, T_centerwl, "-o")
ax.set_xlabel("fiber source position (um)")
ax.set_ylabel(f"transmission at {centerwl} wavelength")
ax.grid()
plt.show()

# update source position
best_fiber_pos = fiber_pos[np.argmax(T_centerwl)]
source = basesim.sources[0]
source = source.updated_copy(center=(best_fiber_pos, source.center[1], source.center[2]))
basesim = basesim.updated_copy(sources=[source])

def worstloss(simdata, wl_range):
    """Returns the worst transmission across measured frequencies from the monitor"""
    freqrange = td.C_0 / np.array(wl_range)
    amps = simdata["P0"].amps.sel(direction="-", f=slice(freqrange[0], freqrange[1])).data
    T = np.abs(amps) ** 2
    return smooth_min(T, tau=1e-2)

def design_to_array(design: dict[str, np.ndarray]) -> tuple[np.ndarray, tuple[int, int]]:
    """Converts design parameters into a single array"""
    gc_gaps = design["gc_gaps"]
    gc_widths = design["gc_widths"]
    array = np.hstack([gc_gaps, gc_widths])
    lengths = (
        len(gc_gaps),
        len(gc_widths),
    )
    return array, lengths


def array_to_design(array: np.ndarray, lengths: tuple[int, int]) -> dict[str, np.ndarray]:
    """Converts array back to design parameters"""
    len_gaps, len_widths = lengths
    idx = 0
    design = {}
    design["gc_gaps"] = array[idx : idx + len_gaps]
    idx += len_gaps
    design["gc_widths"] = array[idx : idx + len_widths]
    return design

def objective(params_array: np.ndarray, param_lengths, wl_range):
    """The objective function for optimization. Creates the simulation, and returns the loss."""
    # first convert the parameter array to a dictionary for passing to the simulation maker
    design = array_to_design(params_array, param_lengths)

    # then create and run the simulation
    sim = gc_sim(**design, basesim=basesim)
    simdata = td.web.run(sim, task_name="gcopt", path="simdata/gc.hdf5", verbose=False)

    # finally return the objective value
    return worstloss(simdata, wl_range=wl_range)

min_line = 0.15
min_space = 0.2


def minfeat_cliparray(design, min_line=min_line, min_space=min_space):
    """This function returns a single array of minimum line and space for clipping."""
    _, param_lengths = design_to_array(design)
    min_spaces = min_space * np.ones(shape=param_lengths[0])
    min_lines = min_line * np.ones(shape=param_lengths[1])
    return np.concatenate((min_spaces, min_lines))

# The wavelength range to optimize for
wl_range = (1.545, 1.555)

# Optimization parameters
learning_rate = 1e-2
num_steps = 40

# The starting grating design
design0 = {
    "gc_gaps": gc_gaps,
    "gc_widths": gc_widths,
}

# Initialize the optimizer with starting parameters
params_array0, param_lengths = design_to_array(design0)
params_array = np.array(params_array0).copy()
optimizer = adam(learning_rate=learning_rate)

# Create the minimum feature size array for clipping
minfeatarr = minfeat_cliparray(design=design0)

# Arrays that store the optimization history
objective_history = []
param_history = []
gradient_history = []

obj_fun = partial(objective, param_lengths=param_lengths, wl_range=wl_range)


def report_step(current_params, grad, state, step_index, objective_val):
    print(f"step = {step_index + 1}")
    print(f"\tJ = {objective_val:.4e}")
    print(f"\tgrad_norm = {np.linalg.norm(grad):.4e}")

    objective_history.append(objective_val)
    param_history.append(np.array(current_params))
    gradient_history.append(np.array(grad))

# maximize the objective directly using optimize(..., direction="max").
params_array, opt_state, history = optimize(
    obj_fun,
    params0=params_array,
    optimizer=optimizer,
    num_steps=num_steps,
    bounds=(minfeatarr, None),
    callback=report_step,
    direction="max",
)

# Plot the objective value versus iteration.
f, ax = plt.subplots()
ax.plot(objective_history, "-o")
ax.set_xlabel("iteration")
ax.set_ylabel("Objective value")
ax.grid()
ax.set_title("Objective vs. iteration")

# Select the final design parameters returned by optimize.
# The callback records pre-update iterates, so param_history
# does not include this final vector.
final_params = params_array
_, param_lengths = design_to_array(design0)
final_design = array_to_design(final_params, param_lengths)

# Compare the starting and final params.
f, ax = plt.subplots()
ax.plot(design0["gc_widths"], "o", label="initial design")
ax.plot(final_design["gc_widths"], "o", label="final design")
ax.set_ylabel("widths")
ax.grid()
ax.legend()
ax.set_title("Initial vs. final widths")

f, ax = plt.subplots()
ax.plot(design0["gc_gaps"], "o", label="initial design")
ax.plot(final_design["gc_gaps"], "o", label="final design")
ax.set_ylabel("gaps")
ax.grid()
ax.legend()
ax.set_title("Initial vs. final gaps")

plt.show()

sim_startdesign = gc_sim(**design0, basesim=basesim)
simdata_startdesign = td.web.run(
    sim_startdesign, task_name="gcopt", path="simdata/gcstart.hdf5", verbose=False
)
sim_finaldesign = gc_sim(**final_design, basesim=basesim)
simdata_finaldesign = td.web.run(
    sim_finaldesign, task_name="gcopt", path="simdata/gcfinal.hdf5", verbose=False
)

# Plot final transmission, compare to the initial design
amps_startdesign = simdata_startdesign["P0"].amps.sel(direction="-")
T_startdesign = amps_startdesign.abs**2
amps_finaldesign = simdata_finaldesign["P0"].amps.sel(direction="-")
T_finaldesign = amps_finaldesign.abs**2

f, ax = plt.subplots()
ax.plot(td.C_0 / T_startdesign.f, T_startdesign, label="initial design")
ax.plot(td.C_0 / T_finaldesign.f, T_finaldesign, label="final design")
ax.set_xlabel("wavelength (um)")
ax.set_ylabel("transmission")
ax.grid()
ax.legend()

plt.show()