# Standard python and external package imports
import pickle
from pathlib import Path

import gdstk
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
import numpy as np

# tidy3D imports
import tidy3d as td
import tidy3d.plugins.design as tdd
import torch
import torch.nn as nn
import torch.optim as optim
from tidy3d import web
from torch.utils.data import DataLoader

# Configure notebook options for reusing simulation data if it exists
use_previous_simulation_data = False

output_dir = Path("./misc/dc_surrogate/")
mc_data_dir = output_dir / "data"
mc_data_dir.mkdir(parents=True, exist_ok=True)

# Wavelength / frequency
lambda0 = 1.550  # all length scales in microns
freq0 = td.constants.C_0 / lambda0
fwidth = freq0 / 10
ldas = np.linspace(1.5, 1.6, 101)  # wavelength range
freqs = td.C_0 / ldas  # frequency range

# Permittivity of waveguide and substrate
wg_n = 3.48
sub_n = 1.45
mat_wg = td.Medium(permittivity=wg_n**2)
mat_sub = td.Medium(permittivity=sub_n**2)

# Waveguide dimensions

# Waveguide height
wg_height = 0.22
# Waveguide width
wg_width = 0.45
# Waveguide separation in the beginning/end
wg_spacing_in = 8
# Reference plane where the cross section of the device is defined
reference_plane = "bottom"
# Angle of the sidewall deviating from the vertical ones, positive values for the base larger than the top
sidewall_angle = np.pi / 20.0
# Total device length along propagation direction
device_length = 100
# Length of the bend region
bend_length = 16
# space between waveguide and PML
pml_spacing = 1
# resolution control: minimum number of grid cells per wavelength in each material
grid_cells_per_wvl = 16

# Coupler optimization parameter ranges to consider
wg_spacing_coup_range = (0.05, 0.25)
coup_length_range = (2, 16)
top_waveguide_transmission_goal = 0.25
bottom_waveguide_transmission_goal = 0.75

# Optimization parameters
num_monte_carlo_sims = 100

filename_mc_data = mc_data_dir / f"result_{num_monte_carlo_sims}.pkl"
if not filename_mc_data.exists() and use_previous_simulation_data:
    print(
        "Warning, MC simulation data does not exist to be loaded. Reconfiguring use_previous_simulation_data flag to False and recomputing."
    )
    use_previous_simulation_data = False

def tanh_interp(max_arg):
    """Interpolator for tanh with adjustable extension"""
    scale = 1 / np.tanh(max_arg)
    return lambda u: 0.5 * (1 + scale * np.tanh(max_arg * (u * 2 - 1)))


def make_coupler(
    length,
    wg_spacing_in,
    wg_width,
    wg_spacing_coup,
    coup_length,
    bend_length,
):
    """Make an integrated coupler using the gdstk RobustPath object."""
    # bend interpolator
    interp = tanh_interp(3)
    delta = wg_width + wg_spacing_coup - wg_spacing_in
    offset = lambda u: wg_spacing_in + interp(u) * delta

    coup = gdstk.RobustPath(
        (-0.5 * length, 0),
        (wg_width, wg_width),
        wg_spacing_in,
        simple_path=True,
        layer=1,
        datatype=[0, 1],
    )
    coup.segment((-0.5 * coup_length - bend_length, 0))
    coup.segment(
        (-0.5 * coup_length, 0),
        offset=[lambda u: -0.5 * offset(u), lambda u: 0.5 * offset(u)],
    )
    coup.segment((0.5 * coup_length, 0))
    coup.segment(
        (0.5 * coup_length + bend_length, 0),
        offset=[lambda u: -0.5 * offset(1 - u), lambda u: 0.5 * offset(1 - u)],
    )
    coup.segment((0.5 * length, 0))
    return coup

def make_sim(coup_length, wg_spacing_coup, add_field_monitor=False):
    """Make a simulation with a given length of the coupling region and
    distance between the waveguides in that region. If ``add_field_monitor``
    is True, a 2D in-plane field monitor will be added.
    """

    # Geometry must be placed in GDS cells to import into Tidy3D
    coup_cell = gdstk.Cell("Coupler")

    substrate = gdstk.rectangle(
        (-device_length / 2, -wg_spacing_in / 2 - 10),
        (device_length / 2, wg_spacing_in / 2 + 10),
        layer=0,
    )
    coup_cell.add(substrate)

    # Add the coupler to a gdstk cell
    gds_coup = make_coupler(
        device_length,
        wg_spacing_in,
        wg_width,
        wg_spacing_coup,
        coup_length,
        bend_length,
    )
    coup_cell.add(gds_coup)

    # Substrate
    (oxide_geo,) = td.PolySlab.from_gds(
        gds_cell=coup_cell,
        gds_layer=0,
        gds_dtype=0,
        slab_bounds=(-10, 0),
        reference_plane=reference_plane,
        axis=2,
    )

    oxide = td.Structure(geometry=oxide_geo, medium=mat_sub)

    # Waveguides (import all datatypes if gds_dtype not specified)
    coupler1_geo, coupler2_geo = td.PolySlab.from_gds(
        gds_cell=coup_cell,
        gds_layer=1,
        slab_bounds=(0, wg_height),
        sidewall_angle=sidewall_angle,
        reference_plane=reference_plane,
        axis=2,
    )

    coupler1 = td.Structure(geometry=coupler1_geo, medium=mat_wg)
    coupler2 = td.Structure(geometry=coupler2_geo, medium=mat_wg)

    # Simulation size along propagation direction
    sim_length = 2 + 2 * bend_length + coup_length

    # Spacing between waveguides and PML
    sim_size = [
        sim_length,
        wg_spacing_in + wg_width + 2 * pml_spacing,
        wg_height + 2 * pml_spacing,
    ]

    # Source
    src_pos = -sim_length / 2 + 0.5
    msource = td.ModeSource(
        center=[src_pos, wg_spacing_in / 2, wg_height / 2],
        size=[0, 3, 2],
        source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth),
        direction="+",
        mode_spec=td.ModeSpec(num_modes=1),
        mode_index=0,
    )

    mon_in = td.ModeMonitor(
        center=[(src_pos + 0.5), wg_spacing_in / 2, wg_height / 2],
        size=[0, 3, 2],
        freqs=freqs,
        mode_spec=td.ModeSpec(num_modes=1),
        name="in",
    )
    mon_ref_bot = td.ModeMonitor(
        center=[(src_pos + 0.5), -wg_spacing_in / 2, wg_height / 2],
        size=[0, 3, 2],
        freqs=freqs,
        mode_spec=td.ModeSpec(num_modes=1),
        name="reflect_bottom",
    )
    mon_top = td.ModeMonitor(
        center=[-(src_pos + 0.5), wg_spacing_in / 2, wg_height / 2],
        size=[0, 3, 2],
        freqs=freqs,
        mode_spec=td.ModeSpec(num_modes=1),
        name="top",
    )
    mon_bot = td.ModeMonitor(
        center=[-(src_pos + 0.5), -wg_spacing_in / 2, wg_height / 2],
        size=[0, 3, 2],
        freqs=freqs,
        mode_spec=td.ModeSpec(num_modes=1),
        name="bottom",
    )
    monitors = [mon_in, mon_ref_bot, mon_top, mon_bot]

    if add_field_monitor:
        field_monitor = td.FieldMonitor(
            center=[0, 0, wg_height / 2],
            size=[td.inf, td.inf, 0],
            freqs=freqs,
            name="field",
        )
        monitors.append(field_monitor)

    # Initialize the simulation
    sim = td.Simulation(
        size=sim_size,
        grid_spec=td.GridSpec.auto(min_steps_per_wvl=grid_cells_per_wvl),
        structures=[oxide, coupler1, coupler2],
        sources=[msource],
        monitors=monitors,
        run_time=50 / fwidth,
        boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()),
    )

    return sim

# Set up the parameters we want to sweep in the Monte Carlo search
param_spc = tdd.ParameterFloat(name="wg_spacing_coup", span=wg_spacing_coup_range)
param_len = tdd.ParameterFloat(name="coup_length", span=coup_length_range)
parameters = [param_spc, param_len]

monte_carlo = tdd.MethodMonteCarlo(
    num_points=num_monte_carlo_sims,
    seed=2,
)
design_space = tdd.DesignSpace(
    method=monte_carlo,
    parameters=parameters,
    task_name="DC_MC_Surrogate",
    path_dir=str(mc_data_dir),
)


def measure_coupling_amplitude(sim_data):
    """Constructs a "row" of the scattering matrix when sourced from top left port"""

    input_amp = sim_data["in"].amps.sel(direction="+")

    amps_dict = {}
    directions = ("-", "-", "+", "+")
    for monitor, direction in zip(sim_data.simulation.monitors[:4], directions):
        amp = sim_data[monitor.name].amps.sel(direction=direction)
        amp_normalized = amp / input_amp
        amps_dict[monitor.name] = amp_normalized

    return amps_dict


def fn_post(*args, **kwargs):
    """Post processing function, but label the outputs using a dictionary."""
    amps_dict = measure_coupling_amplitude(*args, **kwargs)
    return amps_dict

# Neural Network parameters
number_samples = len(expanded_results)
train_samples = int(np.floor(0.6 * number_samples))
validation_samples = int(np.floor(0.2 * number_samples))
test_samples = number_samples - train_samples - validation_samples
batch_size = 16
epochs = 120
final_training_epochs = 1000

training_dataset, validation_dataset, test_dataset = torch.utils.data.random_split(
    expanded_results, [train_samples, validation_samples, test_samples]
)

training_dataloader = DataLoader(training_dataset, batch_size=batch_size, shuffle=True)
validation_dataloader = DataLoader(validation_dataset, batch_size=batch_size, shuffle=False)
testing_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# Assuming the features are coming in as (wg_spacing_coup, coup_length, freq)
# and we can normalize by the min/max possible values to center each feature in [-1, 1]
min_freq = np.min(freqs)
max_freq = np.max(freqs)

subtract_min = torch.tensor(
    [wg_spacing_coup_range[0], coup_length_range[0], min_freq], dtype=torch.float
)
normalize_max = torch.tensor(
    [
        wg_spacing_coup_range[1] - wg_spacing_coup_range[0],
        coup_length_range[1] - coup_length_range[0],
        max_freq - min_freq,
    ],
    dtype=torch.float,
)

network_input_size = 3


# Define a fully connected model in PyTorch with ReLU activation
# where we can set a variable number of hidden layers through
# neurons list. Allow option to specify dropout percentage.
class VariableANN(nn.Module):
    def __init__(self, input_size, neurons, dropout) -> None:
        super().__init__()
        self.layers = nn.ModuleList()
        self.layers.append(nn.Linear(input_size, neurons[0]))

        for i in range(1, len(neurons)):
            self.layers.append(nn.Linear(neurons[i - 1], neurons[i]))

        self.layers.append(nn.Linear(neurons[-1], 2))

        self.drop = nn.Dropout(dropout)
        self.relu = nn.ReLU()

    def forward(self, x):
        subtract_min_unsqueeze = torch.unsqueeze(subtract_min, dim=0)
        normalize_max_unsqueeze = torch.unsqueeze(normalize_max, dim=0)

        # normalize inputs to be between [-1, 1]
        x = x - subtract_min_unsqueeze
        x = x / normalize_max_unsqueeze
        x = 2.0 * (x - 0.5)

        for layer in self.layers[:-1]:
            x = self.drop(self.relu(layer(x)))

        x = self.layers[-1](x)

        return x

loss_fn = nn.MSELoss()


def evaluate_loss(data, network, loss_fn_user=None):
    """Extract inputs/labels from the data and allow for loss evaluation on user defined loss
    function (otherwise use the default loss_fn which is set to MSE)."""

    inputs, labels = data

    if loss_fn_user is None:
        loss_fn_user = loss_fn

    loss = loss_fn_user(network(inputs.type(torch.FloatTensor)), labels.type(torch.FloatTensor))

    return loss


def evaluate_loss_from_loader(data_loader, network, loss_fn_user=None):
    """Given a data_loader, run all the batches through the network and evaluate loss."""
    loss_by_iteration = []
    for data_idx, data in enumerate(data_loader):
        loss = evaluate_loss(data, network, loss_fn_user)
        loss_by_iteration.append(loss.item())

    avg_loss = np.mean(loss_by_iteration)
    return avg_loss


def train_network(hyperparameters, epochs):
    """Train a network according to hyperparameters specification for a given number of epochs."""

    # build network and set to training mode
    network = VariableANN(network_input_size, hyperparameters["hidden"], hyperparameters["dropout"])
    network.train()
    opt = optim.Adam(
        network.parameters(),
        lr=hyperparameters["learning_rate"],
        weight_decay=hyperparameters["weight_decay"],
    )

    training_loss_by_epoch = np.zeros(1 + epochs)
    validation_loss_by_epoch = np.zeros(1 + epochs)

    avg_training_loss = evaluate_loss_from_loader(training_dataloader, network)
    avg_validation_loss = evaluate_loss_from_loader(validation_dataloader, network)

    training_loss_by_epoch[0] = avg_training_loss
    validation_loss_by_epoch[0] = avg_validation_loss

    for epoch_idx in range(0, epochs):
        training_loss_by_iteration = []

        for data_idx, data in enumerate(training_dataloader):
            opt.zero_grad()
            loss = evaluate_loss(data, network)

            training_loss_by_iteration.append(loss.item())

            loss.backward()
            opt.step()

        # Remove training freatures like dropout for the evaluation section and then turn it back on for training
        network.eval()
        avg_validation_loss = evaluate_loss_from_loader(validation_dataloader, network)
        network.train()

        training_loss_by_epoch[1 + epoch_idx] = np.mean(training_loss_by_iteration)
        validation_loss_by_epoch[1 + epoch_idx] = avg_validation_loss

    return training_loss_by_epoch, validation_loss_by_epoch, network

allowed_hidden_layer_sizes = np.arange(16, 257, 16)
allowed_dropout = np.arange(0, 0.55, 0.05)
num_hidden_layers = 2
allowed_learning_rate_log_range = (-5, -3)
allowed_weight_decay_log_range = (-5, -4)

num_trial_networks = 75


def hyperparemeter_fn(dropout, learning_rate_log, weight_decay_log, **hidden_layer_sizes):
    hidden_layer_sizes = [
        hidden_layer_sizes[f"hidden_layer_size_{idx}"] for idx in range(0, len(hidden_layer_sizes))
    ]

    hyperparam = {
        "hidden": hidden_layer_sizes,
        "dropout": dropout,
        "learning_rate": 10**learning_rate_log,
        "weight_decay": 10**weight_decay_log,
    }

    training_loss, validation_loss, network = train_network(hyperparam, epochs)

    return -validation_loss[-1]


param_dropout = tdd.ParameterAny(name="dropout", allowed_values=tuple(allowed_dropout))
param_lr_log = tdd.ParameterFloat(name="learning_rate_log", span=allowed_learning_rate_log_range)
param_wd_log = tdd.ParameterFloat(name="weight_decay_log", span=allowed_weight_decay_log_range)
params_hyperparameter = [
    tdd.ParameterAny(
        name=f"hidden_layer_size_{idx}",
        allowed_values=tuple(allowed_hidden_layer_sizes),
    )
    for idx in range(0, num_hidden_layers)
]
hyperparameters = [param_dropout, param_lr_log, param_wd_log] + params_hyperparameter

hyperparam_optimizer = tdd.MethodMonteCarlo(
    num_points=num_trial_networks,
    seed=2,
)
design_space = tdd.DesignSpace(
    method=hyperparam_optimizer,
    parameters=hyperparameters,
    task_name="DC_Network_MC",
    path_dir=str(mc_data_dir),
)
network_hyperparameter_results = design_space.run(hyperparemeter_fn, verbose=True)