nm = 1e-3

wvl0 = 1550 * nm
a = 400 * nm

num_freqs = 151

wvl_min, wvl_max = (1400 * nm, 1700 * nm)
freq_min = td.C_0 / wvl_max
freq_max = td.C_0 / wvl_min

freqs = np.linspace(freq_min, freq_max, num_freqs)
freq0 = td.C_0 / wvl0

fwidth = freq0 / 10
run_time = 300 / fwidth

sqrt3_div2 = np.sqrt(3) / 2.0

# silicon material (for slab)
n_si = 3.48
si = td.Medium(permittivity=n_si**2)
air = td.Medium()

# radius of holes
r0 = 90 * nm

# when we optimize, the radius will vary between these two values
# for now we just constrain all radius to r0, but change these values to add more degrees of freedom
r_range = rmin, rmax = (r0, r0)
rmid = r0

# how much centers can move from the center
drmax = a / 4

# number of holes in x and y
N_rows = 15
N_cols = 19
N_cols_static = 5

# total length of the PhC region
Lx_phc = N_cols * a + a / 2
Ly_phc = N_rows * sqrt3_div2 * a + a / 2

# buffer on each side of the design region
buffer = 1.0 * wvl0

# thickness of slab
t_slab = 220 * nm

# waveguide width
w_wg = 1.3 * a

# size of simulation
Lx = Lx_phc + 2 * buffer
Ly = Ly_phc + 2 * buffer
Lz = t_slab + 2 * buffer

# define grid resolution
steps_per_unit_cell = 14
dx = a / steps_per_unit_cell
dy = a * sqrt3_div2 / steps_per_unit_cell

grid_spec = td.GridSpec(
    grid_x=td.UniformGrid(dl=dx),
    grid_y=td.UniformGrid(dl=dy),
    grid_z=td.AutoGrid(min_steps_per_wvl=steps_per_unit_cell),
)

def make_holes(params) -> td.Structure:
    """Convenience function to make the phc holes given the design parameters."""

    hole_spacing_x = a
    hole_spacing_y = a * sqrt3_div2

    x_slab_length, y_slab_length = (
        hole_spacing_x * (N_cols + 0.5),
        hole_spacing_y * N_rows,
    )
    start_x, start_y = (
        -x_slab_length / 2 + hole_spacing_x / 2,
        -y_slab_length / 2 + hole_spacing_y / 2,
    )

    cylinders = []

    for i in range(0, N_cols):
        for j in range(0, N_rows):
            # depending on distance from central column, hole is either static
            i_dist = abs(i - N_cols // 2)
            if i_dist < ((N_cols_static) / 2):
                radius = rmid
                dx = dy = 0

            # or optimizable
            else:
                radius = params[0, i, j]
                dx = params[1, i, j]
                dy = params[2, i, j]

            x0 = dx + start_x + (i + (j % 2) * 0.5) * hole_spacing_x
            y0 = dy + start_y + j * hole_spacing_y
            if j != N_rows // 2:
                c = td.Cylinder(
                    axis=2,
                    radius=radius,
                    center=(x0, y0, 0),
                    length=td.inf,
                )
                cylinders.append(c)

    # use GeometryGroup since all same medium, for performance
    structure = td.Structure(
        geometry=td.GeometryGroup(geometries=cylinders),
        medium=air,
        background_medium=si,  # note: we need this for correct gradients when embedded in slab
    )

    return structure

source = td.ModeSource(
    center=(-Lx / 2 + wvl0 / 10, 0, 0),
    size=(0, Ly_phc / 2.0, td.inf),
    source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth),
    direction="+",
)

# this monitor is just for plotting the broadband flux response, not optimizing

slab = td.Structure(
    geometry=td.Box(
        center=(0, 0, 0),
        size=(Lx_phc, Ly_phc, t_slab),
    ),
    medium=si,
)

waveguide = td.Structure(
    geometry=td.Box(
        center=(0, 0, 0),
        size=(td.inf, w_wg, t_slab),
    ),
    medium=si,
)


mnt_flux = td.FluxMonitor(
    center=(-source.center[0], 0, 0),
    size=(0, td.inf, td.inf),
    freqs=freqs,
    name="flux",
)

# this monitor is for visualizing field patterns only
mnt_field = td.FieldMonitor(
    center=(0, 0, 0),
    size=(td.inf, td.inf, 0),
    freqs=[freq0],
    name="field",
)

# This monitor is used for the optimization.
mnt_flux_adjoint = td.FluxMonitor(
    center=(+Lx_phc / 2 - a, 0, 0),
    size=(0, td.inf, td.inf),
    freqs=[freq0],
    name="flux",
    enable_adjoint=True,
)

def make_sim(params, optimization_mode: bool = False) -> td.Simulation:
    """Function to generate a simulation with different monitors depending on whether optimizing."""

    # create the holes
    holes = make_holes(params)

    # decide which monitors to include
    if not optimization_mode:
        monitors = [mnt_flux, mnt_field]
    else:
        monitors = [mnt_flux_adjoint]

    return td.Simulation(
        center=[0, 0, 0],
        size=[Lx, Ly, Lz],
        grid_spec=grid_spec,
        structures=[slab, waveguide, holes],
        sources=[source],
        monitors=monitors,
        run_time=run_time,
        boundary_spec=td.BoundarySpec.pml(x=True, y=True, z=True),
        symmetry=(0, -1, 1),
    )

params0 = rmid * np.ones((3, N_cols, N_rows))
params0[1:] = 0.0

sim0 = make_sim(params0, optimization_mode=False)

sim_wg = sim0.updated_copy(structures=[waveguide])

_, (ax1, ax2) = plt.subplots(1, 2, tight_layout=True, figsize=(10, 4))
_ = sim0.plot(z=0.01, ax=ax1)
_ = sim_wg.plot(z=0.01, ax=ax2)

plt.show()

sim_data0 = web.Job(simulation=sim0, task_name="initial PhC").run()
sim_data_wg = web.Job(simulation=sim_wg, task_name="initial PhC norm").run()

_, (ax1, ax2) = plt.subplots(1, 2, tight_layout=True, figsize=(10, 4))
_ = sim_data_wg.plot_field("field", field_name="Ey", val="abs", ax=ax1)
_ = sim_data0.plot_field("field", field_name="Ey", val="abs", ax=ax2)

plt.show()

flux_wg = abs(sim_data_wg["flux"].flux)
flux0 = abs(sim_data0["flux"].flux) / flux_wg
flux0.plot(x="f")
plt.ylabel("flux (no optimization)")
plt.plot([freq0, freq0], [0, 1], linestyle="--", label="optimization freq")
plt.legend()
plt.show()

flux_wg = abs(sim_data_wg["flux"].flux)
flux0 = abs(sim_data0["flux"].flux) / flux_wg
flux0.plot(x="f")
plt.ylabel("flux (no optimization)")
plt.plot([freq0, freq0], [0, 1], linestyle="--", label="optimization freq")
plt.legend()
plt.show()

flux0_freq0 = flux_wg.interp(f=freq0).item()
print(flux0_freq0)

def objective(params: anp.ndarray) -> float:
    """Maximize flux in -y, minimize flux in +x."""

    sim = make_sim(params, optimization_mode=True)
    sim_data = web.run(sim, task_name="phc_adjoint", verbose=False)

    flux_measure_freq0 = anp.sum(anp.abs(sim_data["flux"].flux.data))

    return flux_measure_freq0 / flux_wg.interp(f=freq0).item()

val_grad = autograd.value_and_grad(objective)

from autograd.tracer import getval
from tidy3d.plugins.autograd import adam, optimize

# hyperparameters
num_steps = 10
learning_rate = a / 40
# note: the step size needs to be quite low because of the direct modification of geometric parameter

# initialize adam optimizer with starting parameters
params = np.array(params0).copy()
optimizer = adam(learning_rate=learning_rate)

# store history
objective_history = []
param_history = []
data_history = []

# define lower and upper bounds for the optimization
lower_bounds = np.zeros_like(params0)
upper_bounds = np.zeros_like(params0)
lower_bounds[0] = rmin
upper_bounds[0] = rmax
lower_bounds[1:] = -drmax
upper_bounds[1:] = drmax

%%time


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}")
    param_history.append(current_params.copy())


params, opt_state, history = optimize(
    objective,
    params0=params,
    optimizer=optimizer,
    num_steps=num_steps,
    bounds=(lower_bounds, upper_bounds),
    callback=report_step,
    direction="max",
)
param_history.append(np.array(params))
objective_history = history["objective_fn_val"]

plt.plot(objective_history)
plt.xlabel("iteration")
plt.ylabel("objective function")
plt.show()

params_final = param_history[-1]

sim_final = make_sim(params_final, optimization_mode=False)
_ = sim_final.plot(z=0.01)
plt.show()

sim_data_final = web.run(sim_final, task_name="phc")

_, (ax1, ax2) = plt.subplots(1, 2, tight_layout=True, figsize=(10, 4))
# vmax = 30

ax1 = sim_data0.plot_field("field", field_name="Ey", val="abs", ax=ax1)
ax2 = sim_data_final.plot_field("field", field_name="Ey", val="abs", ax=ax2)
ax1.set_title("original")
ax2.set_title("optimized")
plt.show()

flux_final = abs(sim_data_final["flux"].flux)

flux0.plot(x="f", label="original")
flux_final.plot(x="f", label="optimized")

plt.ylim([0, 1])
plt.title("transmission")
plt.plot([freq0, freq0], [0, 1], linestyle="--", label="optimization freq")

plt.legend()
plt.show()