# Standard python imports.
import pickle
from typing import List

import gdstk

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

# 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.web import run

plt.rcParams["font.size"] = "12"

# Geometric parameters.
y_width = 1.7  # Y-branch maximum width (um).
y_length = 1.7  # Y-branch maximum length (um).
w_thick = 0.22  # Waveguide thickness (um).
w_width = 0.5  # Waveguide width (um).
w_length = 1.0  # Input output waveguide length (um).
w_gap = 0.3  # Gap between the output waveguides (um).
bend_length = 3  # Output waveguide bend length (um).
bend_offset = 0.5  # Offset between output bends (um).

# Material.
nSi = 3.48  # Silicon refractive index.

# Inverse design set up parameters.
grid_size = 0.016  # Simulation grid size on design region (um).
ls_grid_size = 0.004  # Discretization size of the level set function (um).
ls_down_sample = (
    20  # The spacing between the level set control knots is given by ls_grid_size*ls_down_sample.
)
fom_name_1 = "fom_field1"  # Name of the monitor used to compute the objective function.
min_feature_size = 0.14  # Minimum fabrication feature size (um).
gap_par = 1.0  # Parameter to minimum gap fabrication constraint.
curve_par = 1.5  # Parameter of minimum curvature fabrication constraint.

# Optimizer parameters.
iterations = 100  # Maximum number of iterations in optimization.
learning_rate = 0.03

# Simulation wavelength.
wl = 1.55  # Central simulation wavelength (um).
bw = 0.06  # Simulation bandwidth (um).
n_wl = 61  # Number of wavelength points within the bandwidth.

# Minimum and maximum values for the permittivities.
eps_max = nSi**2
eps_min = 1.0

# Material definition.
mat_si = td.Medium(permittivity=eps_max)  # Waveguide material.

# 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 = 5e-13

# Computational domain size.
pml_spacing = 0.6 * wl
size_x = 2 * w_length + y_length + bend_length
size_y = w_gap + 2 * (bend_offset + w_width + pml_spacing)
size_z = w_thick + 2 * pml_spacing
eff_inf = 10

# Source and monitor positions.
mon_w = 3 * w_width
mon_h = 5 * w_thick

# Separation between the level set control knots.
rho_size = ls_down_sample * ls_grid_size

# Number of points on the parameter grid (rho) and level set grid (phi)
nx_rho = int(y_length / rho_size) + 1
ny_rho = int(y_width / rho_size / 2) + 1
nx_phi = int(y_length / ls_grid_size) + 1
ny_phi = int(y_width / ls_grid_size / 2) + 1
npar = nx_rho * ny_rho
ny_rho *= 2
ny_phi *= 2

# Design region size
dr_size_x = (nx_phi - 1) * ls_grid_size
dr_size_y = (ny_phi - 1) * ls_grid_size
dr_center_x = -size_x / 2 + w_length + dr_size_x / 2

# xy coordinates of the parameter and level set grids.
x_rho = np.linspace(dr_center_x - dr_size_x / 2, dr_center_x + dr_size_x / 2, nx_rho)
x_phi = np.linspace(dr_center_x - dr_size_x / 2, dr_center_x + dr_size_x / 2, nx_phi)
y_rho = np.linspace(-dr_size_y / 2, dr_size_y / 2, ny_rho)
y_phi = np.linspace(-dr_size_y / 2, dr_size_y / 2, ny_phi)

class LevelSetInterp:
    """This class implements the level set surface using Gaussian radial basis functions."""

    def __init__(
        self,
        x0: jnp.ndarray = None,
        y0: jnp.ndarray = None,
        z0: jnp.ndarray = None,
        sigma: float = None,
    ):
        # Input data.
        x, y = jnp.meshgrid(y0, x0)
        xy0 = jnp.column_stack((x.reshape(-1), y.reshape(-1)))
        self.xy0 = xy0
        self.z0 = z0
        self.sig = sigma
        # Builds the level set interpolation model.
        gauss_kernel = self.gaussian(self.xy0, self.xy0)
        self.model = jnp.dot(jnp.linalg.inv(gauss_kernel), self.z0)

    def gaussian(self, xyi, xyj):
        dist = jnp.sqrt(
            (xyi[:, 1].reshape(-1, 1) - xyj[:, 1].reshape(1, -1)) ** 2
            + (xyi[:, 0].reshape(-1, 1) - xyj[:, 0].reshape(1, -1)) ** 2
        )
        return jnp.exp(-(dist**2) / (2 * self.sig**2))

    def get_ls(self, x1, y1):
        xx, yy = jnp.meshgrid(y1, x1)
        xy1 = jnp.column_stack((xx.reshape(-1), yy.reshape(-1)))
        ls = self.gaussian(self.xy0, xy1).T @ self.model
        return ls


# Function to plot the level set surface.
def plot_level_set(x0, y0, rho, x1, y1, phi):
    y, x = np.meshgrid(y0, x0)
    yy, xx = np.meshgrid(y1, x1)

    fig = plt.figure(figsize=(12, 6), tight_layout=True)
    ax1 = fig.add_subplot(1, 2, 1, projection="3d")
    ax1.view_init(elev=45, azim=-45, roll=0)
    ax1.plot_surface(xx, yy, phi, cmap="RdBu", alpha=0.8)
    ax1.contourf(
        xx, yy, phi, levels=[np.amin(phi), 0], zdir="z", offset=0, colors=["k", "w"], alpha=0.5
    )
    ax1.contour3D(xx, yy, phi, 1, cmap="binary", linewidths=[2])
    ax1.scatter(x, y, rho, color="black", linewidth=1.0)
    ax1.set_title("Level set surface")
    ax1.set_xlabel(r"x ($\mu m$)")
    ax1.set_ylabel(r"y ($\mu m$)")
    ax1.xaxis.pane.fill = False
    ax1.yaxis.pane.fill = False
    ax1.zaxis.pane.fill = False
    ax1.xaxis.pane.set_edgecolor("w")
    ax1.yaxis.pane.set_edgecolor("w")
    ax1.zaxis.pane.set_edgecolor("w")

    ax2 = fig.add_subplot(1, 2, 2)
    ax2.contourf(xx, yy, phi, levels=[0, np.amax(phi)], colors=[[0, 0, 0]])
    ax2.set_title("Zero level set contour")
    ax2.set_xlabel(r"x ($\mu m$)")
    ax2.set_ylabel(r"y ($\mu m$)")
    ax2.set_aspect("equal")
    plt.show()

def mirror_param(design_param):
    param = jnp.array(design_param).reshape((nx_rho, int(ny_rho / 2)))
    return jnp.concatenate((jnp.fliplr(jnp.copy(param)), param), axis=1).flatten()


def get_eps(design_param, sharpness=10, plot_levelset=False) -> np.ndarray:
    """Returns the permittivities defined by the zero level set isocontour"""
    phi_model = LevelSetInterp(x0=x_rho, y0=y_rho, z0=design_param, sigma=rho_size)
    phi = phi_model.get_ls(x1=x_phi, y1=y_phi)

    # Calculates the permittivities from the level set surface.
    eps_phi = 0.5 * (jnp.tanh(sharpness * phi) + 1)
    eps = eps_min + (eps_max - eps_min) * eps_phi
    eps = jnp.maximum(eps, eps_min)
    eps = jnp.minimum(eps, eps_max)

    # Reshapes the design parameters into a 2D matrix.
    eps = jnp.reshape(eps, (nx_phi, ny_phi))

    # Plots the level set surface.
    if plot_levelset:
        rho = np.reshape(design_param, (nx_rho, ny_rho))
        phi = np.reshape(phi, (nx_phi, ny_phi))
        plot_level_set(x0=x_rho, y0=y_rho, rho=rho, x1=x_phi, y1=y_phi, phi=phi)

    return eps

def update_design(eps, unfold=False) -> List[tda.JaxStructure]:
    # Reflects the structure about the x-axis.
    eps_val = jnp.array(eps).reshape((nx_phi, ny_phi, 1, 1))
    coords_x = [(dr_center_x - dr_size_x / 2) + ix * ls_grid_size for ix in range(nx_phi)]

    if not unfold:
        # Creation of a JaxCustomMedium using the values of the design parameters.
        coords_yp = [0 + iy * ls_grid_size for iy in range(int(ny_phi / 2))]
        coords = dict(x=coords_x, y=coords_yp, z=[0], f=[freq])
        eps_jax = {
            f"eps_{dim}{dim}": tda.JaxDataArray(
                values=eps_val[:, int(ny_phi / 2) :, :, :], 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=(dr_center_x, dr_size_y / 4, 0), size=(dr_size_x, dr_size_y / 2, w_thick)
        )
        structure = [tda.JaxStructure(geometry=box, medium=eps_medium)]

    else:
        # Creation of a CustomMedium using the values of the design parameters.
        coords_y = [-dr_size_y / 2 + iy * ls_grid_size for iy in range(ny_phi)]
        coords = dict(x=coords_x, y=coords_y, 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=(dr_center_x, 0, 0), size=(dr_size_x, dr_size_y, w_thick))
        structure = [tda.JaxStructure(geometry=box, medium=eps_medium)]
    return structure

# Figure of Merit (FOM) calculation.
def fom_eps(eps_ref: jnp.ndarray, eps: jnp.ndarray) -> float:
    """Calculate the L2 norm between eps_ref and eps."""
    return jnp.mean(jnp.abs(eps_ref - eps) ** 2)


# Objective function to be passed to the optimization algorithm.
def obj_eps(design_param, eps_ref) -> float:
    param = mirror_param(design_param)
    eps = get_eps(param)
    return fom_eps(eps_ref, eps)


# Function to calculate the objective function value and its
# gradient with respect to the design parameters.
obj_grad_eps = value_and_grad(obj_eps)

# Input waveguide.
wg_input = td.Structure(
    geometry=td.Box.from_bounds(
        rmin=(-eff_inf, -w_width / 2, -w_thick / 2),
        rmax=(-size_x / 2 + w_length + grid_size, w_width / 2, w_thick / 2),
    ),
    medium=mat_si,
)

# Output bends.
x_start = (
    -size_x / 2 + w_length + dr_size_x - grid_size
)  # x-coordinate of the starting point of the waveguide bends.
x = np.linspace(x_start, x_start + bend_length, 100)  # x-coordinates of the top edge vertices.
y = (
    (x - x_start) * bend_offset / bend_length
    - bend_offset * np.sin(2 * np.pi * (x - x_start) / bend_length) / (np.pi * 2)
    + (w_gap + w_width) / 2
)  # y coordinates of the top edge vertices

# adding the last point to include the straight waveguide at the output
x = np.append(x, eff_inf)
y = np.append(y, y[-1])

# add path to the cell
cell = gdstk.Cell("bend")
cell.add(gdstk.FlexPath(x + 1j * y, w_width, layer=1, datatype=0))  # Top waveguide bend.
cell.add(gdstk.FlexPath(x - 1j * y, w_width, layer=1, datatype=0))  # Bottom waveguide bend.

# Define top waveguide bend structure.
wg_bend_top = td.Structure(
    geometry=td.PolySlab.from_gds(
        cell,
        gds_layer=1,
        axis=2,
        slab_bounds=(-w_thick / 2, w_thick / 2),
    )[1],
    medium=mat_si,
)

# Define bottom waveguide bend structure.
wg_bend_bot = td.Structure(
    geometry=td.PolySlab.from_gds(
        cell,
        gds_layer=1,
        axis=2,
        slab_bounds=(-w_thick / 2, w_thick / 2),
    )[0],
    medium=mat_si,
)

# Input mode source.
mode_spec = td.ModeSpec(num_modes=1, target_neff=nSi)
source = td.ModeSource(
    center=(-size_x / 2 + 0.15 * wl, 0, 0),
    size=(0, mon_w, mon_h),
    source_time=td.GaussianPulse(freq0=freq, fwidth=freqw),
    direction="+",
    mode_spec=mode_spec,
    mode_index=0,
)

# Monitor where we will compute the objective function from.
fom_monitor_1 = td.ModeMonitor(
    center=[size_x / 2 - 0.25 * wl, (w_gap + w_width) / 2 + bend_offset, 0],
    size=[0, mon_w, mon_h],
    freqs=[freq],
    mode_spec=mode_spec,
    name=fom_name_1,
)

### Monitors used only to visualize the initial and final y-branch results.
# Field monitors to visualize the final fields.
field_xy = td.FieldMonitor(
    size=(td.inf, td.inf, 0),
    freqs=[freq],
    name="field_xy",
)

# Monitor where we will compute the objective function from.
fom_final_1 = td.ModeMonitor(
    center=[size_x / 2 - 0.25 * wl, (w_gap + w_width) / 2 + bend_offset, 0],
    size=[0, mon_w, mon_h],
    freqs=freqs,
    mode_spec=mode_spec,
    name="out_1",
)

def make_adjoint_sim(design_param, unfold=False) -> tda.JaxSimulation:
    # Builds the design region from the design parameters.
    eps = get_eps(design_param)
    design_structure = update_design(eps, unfold=unfold)

    # Creates a uniform mesh for the design region.
    adjoint_dr_mesh = td.MeshOverrideStructure(
        geometry=td.Box(center=(dr_center_x, 0, 0), size=(dr_size_x, dr_size_y, w_thick)),
        dl=[grid_size, grid_size, grid_size],
        enforce=True,
    )

    return tda.JaxSimulation(
        size=[size_x, size_y, size_z],
        center=[0, 0, 0],
        grid_spec=td.GridSpec.auto(
            wavelength=wl_max,
            min_steps_per_wvl=15,
            override_structures=[adjoint_dr_mesh],
        ),
        symmetry=(0, -1, 1),
        structures=[wg_input, wg_bend_top, wg_bend_bot],
        input_structures=design_structure,
        sources=[source],
        monitors=[],
        output_monitors=[fom_monitor_1],
        run_time=run_time,
        subpixel=True,
    )

# Auxiliary function to calculate first and second order partial derivatives.
def ls_derivatives(phi, d_size):
    SC = 1e-12
    phi_1 = jnp.gradient(phi, d_size)
    phi_x = phi_1[0] + SC
    phi_y = phi_1[1] + SC
    phi_2x = jnp.gradient(phi_x, d_size)
    phi_2y = jnp.gradient(phi_y, d_size)
    phi_xx = phi_2x[0]
    phi_xy = phi_2x[1]
    phi_yy = phi_2y[1]
    return phi_x, phi_y, phi_xx, phi_xy, phi_yy


# Minimum gap size fabrication constraint integrand calculation.
# The "beta" parameter relax the constraint near the zero plane.
def fab_penalty_ls_gap(params, beta=1, min_feature_size=min_feature_size, grid_size=ls_grid_size):
    # Get the level set surface.
    phi_model = LevelSetInterp(x0=x_rho, y0=y_rho, z0=params, sigma=rho_size)
    phi = phi_model.get_ls(x1=x_phi, y1=y_phi)
    phi = jnp.reshape(phi, (nx_phi, ny_phi))

    # Calculates their derivatives.
    phi_x, phi_y, phi_xx, phi_xy, phi_yy = ls_derivatives(phi, grid_size)

    # Calculates the gap penalty over the level set grid.
    pi_d = np.pi / (1.3 * min_feature_size)
    phi_v = jnp.maximum(jnp.power(phi_x**2 + phi_y**2, 0.5), jnp.power(1e-32, 1 / 4))
    phi_vv = (phi_x**2 * phi_xx + 2 * phi_x * phi_y * phi_xy + phi_y**2 * phi_yy) / phi_v**2
    return (
        jnp.maximum((jnp.abs(phi_vv) / (pi_d * jnp.abs(phi) + beta * phi_v) - pi_d), 0)
        * grid_size**2
    )


# Minimum radius of curvature fabrication constraint integrand calculation.
# The "alpha" parameter controls its relative weight to the gap penalty.
# The "sharpness" parameter controls the smoothness of the surface near the zero-contour.
def fab_penalty_ls_curve(
    params, alpha=1, sharpness=1, min_feature_size=min_feature_size, grid_size=ls_grid_size
):
    # Get the permittivity surface and calculates their derivatives.
    eps = get_eps(params, sharpness=sharpness)
    eps_x, eps_y, eps_xx, eps_xy, eps_yy = ls_derivatives(eps, grid_size)

    # Calculates the curvature penalty over the permittivity grid.
    pi_d = np.pi / (1.1 * min_feature_size)
    eps_v = jnp.maximum(jnp.sqrt(eps_x**2 + eps_y**2), jnp.power(1e-32, 1 / 6))
    k = (eps_x**2 * eps_yy - 2 * eps_x * eps_y * eps_xy + eps_y**2 * eps_xx) / eps_v**3
    curve_const = jnp.abs(k * jnp.arctan(eps_v / eps)) - pi_d
    return alpha * jnp.maximum(curve_const, 0) * grid_size**2


# Gap and curvature fabrication constraints calculation.
# Penalty values are normalized by "norm_gap" and "norm_curve".
def fab_penalty_ls(
    params,
    beta=gap_par,
    alpha=curve_par,
    sharpness=4,
    min_feature_size=min_feature_size,
    grid_size=ls_grid_size,
    norm_gap=1,
    norm_curve=1,
):
    # Get the gap penalty fabrication constraint value.
    gap_penalty_int = fab_penalty_ls_gap(
        params=params, beta=beta, min_feature_size=min_feature_size, grid_size=grid_size
    )
    gap_penalty = jnp.nansum(gap_penalty_int) / norm_gap

    # Get the curvature penalty fabrication constraint value.
    curve_penalty_int = fab_penalty_ls_curve(
        params=params,
        alpha=alpha,
        sharpness=sharpness,
        min_feature_size=min_feature_size,
        grid_size=grid_size,
    )
    curve_penalty = jnp.nansum(curve_penalty_int) / norm_curve

    return gap_penalty, curve_penalty

# Figure of Merit (FOM) calculation.
def fom(sim_data: tda.JaxSimulationData) -> float:
    """Return the power at the mode index of interest."""
    output_amps1 = sim_data.output_data[0].amps
    amp1 = output_amps1.sel(direction="+", f=freq, mode_index=0)
    eta1 = jnp.sum(jnp.abs(amp1)) ** 2
    return jnp.abs(0.5 - eta1), eta1


# Objective function to be passed to the optimization algorithm.
def obj(
    design_param, fab_const: float = 0.0, norm_gap=1.0, norm_curve=1.0, verbose: bool = False
) -> float:
    param = mirror_param(design_param)
    sim = make_adjoint_sim(param)
    sim_data = run(sim, task_name="inv_des_ybranch", verbose=verbose)
    fom_val, eta1 = fom(sim_data)
    fab_gap, fab_curve = fab_penalty_ls(param, norm_gap=norm_gap, norm_curve=norm_curve)
    J = fom_val + fab_const * (fab_gap + fab_curve)
    return J, [sim_data, eta1, fab_gap, fab_curve]


# 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/y_branch_fab.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

td.config.logging_level = "ERROR"

iter_done = len(history_dict["values"])

if iter_done < iterations:
    for i in range(iter_done, iterations):
        # Compute gradient and current objective function value.
        (value, data), gradient = obj_grad(
            params, fab_const=0.05, norm_gap=init_fab_gap, norm_curve=init_fab_curve
        )
        sim_data_i, eta1, penalty_gap, penalty_curve = data

        # outputs
        print(f"Step = {i + 1}")
        print(f"\tobj_val = {value:.4e}")
        print(f"\tgrad_norm = {np.linalg.norm(gradient):.4e}")
        print(f"\teta1 = {eta1:.3f}")
        print(f"\tpenalty gap = {penalty_gap:.3f}")
        print(f"\tpenalty curve = {penalty_curve:.3f}")

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

        # Save history.
        history_dict["values"].append(value)
        history_dict["eta1"].append(eta1)
        history_dict["penalty_gap"].append(penalty_gap)
        history_dict["penalty_curve"].append(penalty_curve)
        history_dict["params"].append(params)
        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"])
eta1_vals = np.array(history_dict["eta1"])
pen_gap_vals = np.array(history_dict["penalty_gap"])
pen_curve_vals = np.array(history_dict["penalty_curve"])
final_par = history_dict["params"][-1]

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