# Standard python imports.
import pickle
from typing import List

# Import autograd for automatic differentiation.
import autograd as ag
import autograd.numpy as anp
import matplotlib.pylab as plt
import numpy as np
import optax
import scipy as sp

# Import regular tidy3d.
import tidy3d as td
import tidy3d.web as web
from tidy3d.plugins.autograd import (
    make_erosion_dilation_penalty,
    make_filter_and_project,
    rescale,
    value_and_grad,
)

# Geometric parameters.
cr_w = 1.0  # Coupling region width (um).
cr_l = 3.0  # Coupling region length (um).
wg_thick = 0.19  # Collection waveguide thickness (um).
wg_width = 0.35  # Collection waveguide width (um).
wg_length = 1.0  # Collection waveguide length (um).

# Material.
n_wg = 3.50  # Structure refractive index.
n_sub = 1.44  # Substrate refractive index.

# Fabrication constraints.
min_feature = 0.06  # Minimum feature size.
non_etch_r = 0.06  # Non-etched circular region radius (um).

# Inverse design set up parameters.
grid_size = 0.015  # Simulation grid size on design region (um).
max_iter = 100  # Maximum number of iterations.
iter_steps = 5  # Beta is increased at each iter_steps.
beta_min = 1.0  # Minimum value for the tanh projection parameter.
learning_rate = 0.02

# Simulation wavelength.
wl = 0.94  # Central simulation wavelength (um).
bw = 0.04  # Simulation bandwidth (um).
n_wl = 41  # Number of wavelength points within the bandwidth.

# Minimum and maximum values of the permittivity.
eps_max = n_wg**2
eps_min = 1.0

# Material definition.
mat_wg = td.Medium(permittivity=eps_max)
mat_sub = td.Medium(permittivity=n_sub**2)

# Wavelengths and frequencies.
wl_max = wl + bw / 2
wl_min = wl - bw / 2
wl_range = np.linspace(wl_min, wl_max, n_wl)
freq = td.C_0 / wl
freqs = td.C_0 / wl_range
freqw = 0.5 * (freqs[0] - freqs[-1])
run_time = 3e-12

# Computational domain size.
pml_spacing = 0.6 * wl
size_x = wg_length + cr_l + pml_spacing
size_y = cr_w + 2 * pml_spacing
size_z = wg_thick + 2 * pml_spacing
eff_inf = 10

# Source position and monitor planes.
cr_center_x = wg_length + cr_l / 2
qe_pos = td.Box(center=(cr_center_x - 0.5, 0, 0), size=(0, 0, 0))
qe_field_plan = td.Box.surfaces(center=(cr_center_x, 0, 0), size=(cr_l, cr_w, 2 * wg_thick))
wg_mode_plan = td.Box(center=(wl / 4, 0, 0), size=(0, 4 * wg_width, 5 * wg_thick))

# Number of points on design grid.
nx_grid = int(cr_l / grid_size)
ny_grid = int(cr_w / grid_size / 2)

# xy coordinates of design grid.
x_grid = np.linspace(cr_center_x - cr_l / 2, cr_center_x + cr_l / 2, nx_grid)
y_grid = np.linspace(0, cr_w / 2, ny_grid)

def pre_process(params, beta):
    filter_project = make_filter_and_project(
        radius=min_feature, dl=grid_size, beta=beta, eta=0.5, filter_type="conic"
    )
    params1 = filter_project(params, beta)
    return params1


def get_eps(params, beta: float = 1.00):
    """Returns the permittivities after filter and projection transformations"""
    params1 = pre_process(params, beta=beta)
    eps = eps_min + (eps_max - eps_min) * params1
    eps = anp.maximum(eps, eps_min)
    eps = anp.minimum(eps, eps_max)
    return eps

def include_constant_regions(eps, circ_center=[0, 0], circ_radius=1.0):
    # Build the geometric mask.
    yv, xv = anp.meshgrid(y_grid, x_grid)

    # Shouldn't this be --> |x-x0|^2 + |y-y0|^2 <= r*2
    geo_mask = (
        anp.where(
            anp.abs((xv - circ_center[0]) ** 2 + (yv - circ_center[1]) ** 2) <= (circ_radius**2),
            1,
            0,
        )
        * eps_max
    )
    eps = anp.maximum(geo_mask, eps)
    return eps

def update_design(eps, unfold=False) -> List[td.Structure]:
    # Definition of the coordinates x,y along the design region.
    coords_x = [(cr_center_x - cr_l / 2) + ix * grid_size for ix in range(nx_grid)]
    eps_val = anp.array(eps).reshape((nx_grid, ny_grid, 1))

    if not unfold:
        coords_yp = [0 + iy * grid_size for iy in range(ny_grid)]
        coords = dict(x=coords_x, y=coords_yp, z=[0])
        eps1 = td.SpatialDataArray(eps_val, coords)
        eps_medium = td.CustomMedium(permittivity=eps1)
        box = td.Box(center=(cr_center_x, cr_w / 4, 0), size=(cr_l, cr_w / 2, wg_thick))
        structure = [td.Structure(geometry=box, medium=eps_medium)]

    # VJP for one of anp.copy(), anp.concatenate(), or anp.fliplr() not defined,
    # so the optimization should only be run with `unfold=False` for now
    else:
        coords_y = [-cr_w / 2 + iy * grid_size for iy in range(2 * ny_grid)]
        coords = dict(x=coords_x, y=coords_y, z=[0])
        eps1 = td.SpatialDataArray(
            anp.concatenate((anp.fliplr(anp.copy(eps_val)), eps_val), axis=1), coords
        )
        eps_medium = td.CustomMedium(permittivity=eps1)
        box = td.Box(center=(cr_center_x, 0, 0), size=(cr_l, cr_w, wg_thick))
        structure = [td.Structure(geometry=box, medium=eps_medium)]
    return structure

# Input/output waveguide.
waveguide = td.Structure(
    geometry=td.Box.from_bounds(
        rmin=(-eff_inf, -wg_width / 2, -wg_thick / 2),
        rmax=(wg_length, wg_width / 2, wg_thick / 2),
    ),
    medium=mat_wg,
)

# Substrate layer.
substrate = td.Structure(
    geometry=td.Box.from_bounds(
        rmin=(-eff_inf, -eff_inf, -eff_inf), rmax=(eff_inf, eff_inf, -wg_thick / 2)
    ),
    medium=mat_sub,
)

# Point dipole source located at the center of TiO2 thin film.
dp_source = td.PointDipole(
    center=qe_pos.center,
    source_time=td.GaussianPulse(freq0=freq, fwidth=freqw),
    polarization="Ey",
)

# Mode monitor to compute the FOM.
mode_spec = td.ModeSpec(num_modes=1, target_neff=n_wg)
mode_monitor_fom = td.ModeMonitor(
    center=wg_mode_plan.center,
    size=wg_mode_plan.size,
    freqs=[freq],
    mode_spec=mode_spec,
    name="mode_monitor_fom",
)

# Field monitor to compute the FOM.
field_monitor_fom = []
for i, plane in enumerate(qe_field_plan):
    field_monitor_fom.append(
        td.FieldMonitor(
            center=plane.center,
            size=plane.size,
            freqs=[freq],
            name=f"field_monitor_fom_{i}",
        )
    )

# Mode monitor to compute spectral response.
mode_spec = td.ModeSpec(num_modes=1, target_neff=n_wg)
mode_monitor = td.ModeMonitor(
    center=wg_mode_plan.center,
    size=wg_mode_plan.size,
    freqs=freqs,
    mode_spec=mode_spec,
    name="mode_monitor",
)

# Field monitor to compute spectral response.
field_monitor = []
for i, plane in enumerate(qe_field_plan):
    field_monitor.append(
        td.FieldMonitor(
            center=plane.center, size=plane.size, freqs=freqs, name=f"field_monitor_{i}"
        )
    )

# Field monitor to visualize the fields.
field_monitor_xy = td.FieldMonitor(
    center=(size_x / 2, 0, 0),
    size=(size_x, size_y, 0),
    freqs=freqs,
    name="field_xy",
)

def make_adjoint_sim(param, beta: float = 1.00, unfold=False):
    eps = get_eps(param, beta)
    eps = include_constant_regions(
        eps, circ_center=[qe_pos.center[0], qe_pos.center[1]], circ_radius=non_etch_r
    )
    structure = update_design(eps, unfold=unfold)

    # Creates a uniform mesh for the design region.
    adjoint_dr_mesh = td.MeshOverrideStructure(
        geometry=td.Box(center=(cr_center_x, 0, 0), size=(cr_w, cr_l, wg_thick)),
        dl=[grid_size, grid_size, grid_size],
        enforce=True,
    )

    grid_spec = td.GridSpec.auto(
        wavelength=wl_max,
        min_steps_per_wvl=15,
        override_structures=[adjoint_dr_mesh],
    )

    return td.Simulation(
        size=[size_x, size_y, size_z],
        center=[size_x / 2, 0, 0],
        grid_spec=grid_spec,
        symmetry=(0, -1, 0),
        structures=[substrate, waveguide] + structure,
        sources=[dp_source],
        monitors=[field_monitor_xy, mode_monitor_fom] + field_monitor_fom,
        run_time=run_time,
        subpixel=True,
    )

# Figure of Merit (FOM) calculation.
def fom(sim_data: td.SimulationData) -> float:
    """Return the coupling efficiency."""
    # best to use autograd-wrapped numpy functions for differentiation
    mode_amps = sim_data["mode_monitor_fom"].amps.sel(direction="-", f=freq, mode_index=0).data
    mode_power = anp.sum(anp.abs(mode_amps) ** 2)

    # unlike Jax version, should avoid in-place operators (e.g, `+=`), use numpy when possible
    field_mon_list = [sim_data[f"field_monitor_fom_{i}"] for i in range(0, 6)]
    dip_power = anp.sum([anp.abs(mon.flux.data) for mon in field_mon_list])

    return mode_power, dip_power


def penalty(params, beta) -> float:
    """Penalize changes in structure after erosion and dilation to enforce larger feature sizes."""
    params_processed = pre_process(params, beta=beta)
    erode_dilate_penalty = make_erosion_dilation_penalty(radius=min_feature, dl=grid_size)
    ed_penalty = erode_dilate_penalty(params_processed)
    return ed_penalty


# Objective function to be passed to the optimization algorithm.
def obj(param, beta: float = 1.0, step_num: int = None, verbose: bool = False) -> float:
    sim = make_adjoint_sim(param, beta, unfold=False)  # non-differentiable if `unfold=True`
    task_name = "inv_des"
    if step_num:
        task_name += f"_step_{step_num}"
    sim_data = web.run(sim, task_name=task_name, verbose=verbose)
    mode_power, dip_power = fom(sim_data)
    fom_val = mode_power / dip_power
    penalty_weight = 0.1
    penalty_val = penalty(param, beta)
    J = fom_val - penalty_weight * penalty_val
    return J, [sim_data, mode_power, dip_power, penalty_val]


# Function to calculate the objective function value and its gradient with respect to the design parameters.
# Use tidy3d's wrapped ag.value_and_grad() for it's auxiliary data functionality
obj_grad = value_and_grad(obj, has_aux=True)

# where to store history
history_fname = "misc/qe_light_coupler_autograd.pkl"


def save_history(history_dict: dict) -> None:
    """Convenience function to save the history to file."""
    with open(history_fname, "wb") as file:
        pickle.dump(history_dict, file)


def load_history() -> dict:
    """Convenience method to load the history from file."""
    with open(history_fname, "rb") as file:
        history_dict = pickle.load(file)
    return history_dict

iter_done = len(history_dict["values"])
if iter_done < max_iter:
    # small # of iters for quick testing
    for i in range(iter_done, max_iter):
        print(f"Iteration = ({i + 1} / {max_iter})")
        plt.subplots(1, 1, figsize=(3, 2))
        plt.imshow(np.flipud(1 - params.T), cmap="gray", vmin=0, vmax=1)
        plt.axis("off")
        plt.show()

        # Compute gradient and current objective function value.
        beta_i = i // iter_steps + beta_min
        (value, gradient), data = obj_grad(params, beta=beta_i, step_num=(i + 1))
        sim_data_i, mode_power_i, dip_power_i, penalty_val_i = [data[0]] + [
            dat._value for dat in data[1:]
        ]
        # Outputs.
        print(f"\tbeta = {beta_i}")
        print(f"\tJ = {value:.4e}")
        print(f"\tgrad_norm = {np.linalg.norm(gradient):.4e}")
        print(f"\tpenalty = {penalty_val_i:.3f}")
        print(f"\tmode power = {mode_power_i:.3f}")
        print(f"\tdip power = {dip_power_i:.3f}")
        print(f"\tcoupling efficiency = {mode_power_i / dip_power_i:.3f}")

        # Compute and apply updates to the optimizer based on gradient (-1 sign to maximize obj_fn).
        updates, opt_state = optimizer.update(-gradient, opt_state, params)
        params = optax.apply_updates(params, updates)

        # Cap parameters between 0 and 1.
        params = anp.minimum(params, 1.0)
        params = anp.maximum(params, 0.0)

        # Save history.
        history_dict["values"].append(value)
        history_dict["coupl_eff"].append(mode_power_i / dip_power_i)
        history_dict["penalty"].append(penalty_val_i)
        history_dict["params"].append(params)
        history_dict["beta"].append(beta_i)
        history_dict["gradients"].append(gradient)
        history_dict["opt_states"].append(opt_state)
        # history_dict["data"].append(sim_data_i)  # Uncomment to store data, can create large files.
        save_history(history_dict)

obj_vals = np.array(history_dict["values"])
ce_vals = np.array(history_dict["coupl_eff"])
pen_vals = np.array(history_dict["penalty"])
final_par_density = history_dict["params"][-1]
final_beta = history_dict["beta"][-1]

# just to inspect design at different iterations
def unfold_params(params):
    params = np.concatenate((np.fliplr(np.copy(params)), params), axis=1)
    return params


params1 = history_dict["params"][32]
params1_full = pre_process(params1, beta=final_beta)
params1_full = include_constant_regions(
    params1_full, circ_center=[qe_pos.center[0], qe_pos.center[1]], circ_radius=non_etch_r
)
params1_full = unfold_params(params1_full)

params2 = history_dict["params"][-1]
params2_full = pre_process(params2, beta=final_beta)
params2_full = include_constant_regions(
    params2_full, circ_center=[qe_pos.center[0], qe_pos.center[1]], circ_radius=non_etch_r
)
params2_full = unfold_params(params2_full)

# make the misc/ directory to store the GDS file if it doesn't exist already
import os

if not os.path.exists("./misc/"):
    os.mkdir("./misc/")

sim_final.to_gds_file(
    fname="./misc/inv_des_light_extractor_autograd.gds",
    z=0,
    permittivity_threshold=(eps_max + eps_min) / 2,
    frequency=freq,
)