# Standard python imports.
import pickle
from typing import List

# Import jax to be able to use automatic differentiation.
import jax
import jax.numpy as jnp
import matplotlib.pylab as plt
import numpy as np
import optax
import scipy as sp

# Import regular tidy3d.
import tidy3d as td

# Import the components we need from the adjoint plugin.
import tidy3d.plugins.adjoint as tda
import tidy3d.web as web
from jax import value_and_grad
from tidy3d.plugins.adjoint.utils.filter import ConicFilter
from tidy3d.plugins.adjoint.utils.penalty import ErosionDilationPenalty
from tidy3d.plugins.adjoint.web import run

# 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.03  # 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.1

# 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)

conic_filter = ConicFilter(radius=min_feature, design_region_dl=grid_size)


def tanh_projection(x, beta, eta=0.5):
    tanhbn = jnp.tanh(beta * eta)
    num = tanhbn + jnp.tanh(beta * (x - eta))
    den = tanhbn + jnp.tanh(beta * (1 - eta))
    return num / den


def filter_project(x, beta, eta=0.5):
    x = conic_filter.evaluate(x)
    return tanh_projection(x, beta=beta, eta=eta)


def pre_process(params, beta):
    params1 = filter_project(params, beta=beta)
    return params1


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

def include_constant_regions(eps, circ_center=[0, 0], circ_radius=1.0) -> jnp.ndarray:
    # Build the geometric mask.
    yv, xv = jnp.meshgrid(y_grid, x_grid)
    geo_mask = (
        jnp.where(
            jnp.abs((xv - circ_center[0]) ** 2 + (yv - circ_center[1]) ** 2)
            <= (2 * circ_radius) ** 2,
            1,
            0,
        )
        * eps_max
    )
    eps = jnp.maximum(geo_mask, eps)
    return eps

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

    if not unfold:
        # Creation of a JaxCustomMedium using the values of the design parameters.
        coords_yp = [0 + iy * grid_size for iy in range(ny_grid)]
        coords = dict(x=coords_x, y=coords_yp, z=[0], f=[freq])
        eps_jax = {
            f"eps_{dim}{dim}": tda.JaxDataArray(values=eps_val, coords=coords) for dim in "xyz"
        }
        eps_dataset = tda.JaxPermittivityDataset(**eps_jax)
        eps_medium = tda.JaxCustomMedium(eps_dataset=eps_dataset, interp_method="linear")
        box = tda.JaxBox(center=(cr_center_x, cr_w / 4, 0), size=(cr_l, cr_w / 2, wg_thick))
        structure = [tda.JaxStructure(geometry=box, medium=eps_medium)]

    else:
        # Creation of a CustomMedium using the values of the design parameters.
        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], f=[freq])
        eps_jax = {
            f"eps_{dim}{dim}": tda.JaxDataArray(
                values=jnp.concatenate((jnp.fliplr(jnp.copy(eps_val)), eps_val), axis=1),
                coords=coords,
            )
            for dim in "xyz"
        }
        eps_dataset = tda.JaxPermittivityDataset(**eps_jax)
        eps_medium = tda.JaxCustomMedium(eps_dataset=eps_dataset, interp_method="linear")
        box = tda.JaxBox(center=(cr_center_x, 0, 0), size=(cr_l, cr_w, wg_thick))
        structure = [tda.JaxStructure(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],
            colocate=False,
            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=[freq],
    name="field_xy",
)

def make_adjoint_sim(param, beta: float = 1.00, unfold=False) -> tda.JaxSimulation:
    # Builds the design region from the design parameters.
    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_jax = 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,
    )

    return tda.JaxSimulation(
        size=[size_x, size_y, size_z],
        center=[size_x / 2, 0, 0],
        grid_spec=td.GridSpec.auto(
            wavelength=wl_max,
            min_steps_per_wvl=15,
            override_structures=[adjoint_dr_mesh],
        ),
        symmetry=(0, -1, 0),
        structures=[substrate, waveguide],
        input_structures=structure_jax,
        sources=[dp_source],
        monitors=[field_monitor_xy],
        output_monitors=[mode_monitor_fom] + field_monitor_fom,
        run_time=run_time,
        subpixel=True,
    )

# Figure of Merit (FOM) calculation.
def fom(sim_data: tda.JaxSimulationData) -> float:
    """Return the coupling efficiency."""
    mode_amps = sim_data.output_data[0].amps.sel(direction="-", f=freq, mode_index=0)
    mode_power = jnp.sum(jnp.abs(mode_amps) ** 2)
    dip_power = 0
    for m in range(1, 7):
        field_mon = sim_data.output_data[m]
        dip_power += jnp.abs(field_mon.flux)
    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)
    ed_penalty = ErosionDilationPenalty(length_scale=min_feature, pixel_size=grid_size)
    return ed_penalty.evaluate(params_processed)


# 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)
    task_name = "inv_des"
    if step_num:
        task_name += f"_step_{step_num}"
    sim_data = 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.
obj_grad = value_and_grad(obj, has_aux=True)

# where to store history
history_fname = "./misc/qe_light_coupler.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:
    for i in range(iter_done, max_iter):
        print(f"Iteration = ({i + 1} / {max_iter})")

        # Compute gradient and current objective function value.
        beta_i = i // iter_steps + beta_min
        (value, data), gradient = obj_grad(params, beta=beta_i, step_num=(i + 1))
        sim_data_i, mode_power_i, dip_power_i, penalty_val_i = data
        # 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 = jnp.minimum(params, 1.0)
        params = jnp.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]

# 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.gds",
    z=0,
    permittivity_threshold=(eps_max + eps_min) / 2,
    frequency=freq,
)