import json
from copy import deepcopy
from pathlib import Path

import autograd.numpy as np
import matplotlib.pyplot as plt
import tidy3d as td
from autograd import value_and_grad
from optim import adam_update, apply_updates, clip_params, init_adam
from setup import (
    center_wavelength,
    first_gap_sin,
    get_mode_monitor_power,
    make_simulation,
    max_gap_si,
    max_gap_sin,
    max_width_si,
    max_width_sin,
    min_gap_si,
    min_gap_sin,
    min_width_si,
    min_width_sin,
    num_elements,
)
from tidy3d import web

def objective(params):
    """Objective function for adjoint optimization.

    Takes a dictionary of geometry parameters and returns a scalar loss.
    The function is differentiable via autograd so the adjoint method can
    supply gradients for every parameter in one shot.

    Parameters
    ----------
    params:
        Dictionary holding the current grating geometry arrays.

    Returns
    -------
    float
        Negative of the coupling efficiency so gradient descent maximizes power.
    """
    # Build the tidy3d simulation with the current parameters. Autograd traces
    # everything through the power extraction so the adjoint gradient can be
    # computed efficiently.
    sim = make_simulation(
        params["widths_si"],
        params["gaps_si"],
        params["widths_sin"],
        params["gaps_sin"],
        first_gap_si=params["first_gap_si"],
        first_gap_sin=params["first_gap_sin"],
    )

    sim_data = web.run(sim, task_name="gc_adjoint", verbose=False)

    # Convert the mode monitor result into a scalar objective (negative power)
    # so minimization increases the coupled power at the target wavelength.
    power_da = get_mode_monitor_power(sim_data)
    freq0 = td.C_0 / center_wavelength
    target_power = power_da.sel(f=freq0, method="nearest")
    return -target_power.item()

scalars = {}
with Path("./results/gc_bayes_opt_best.json").open("r", encoding="utf-8") as f:
    data = json.load(f)
    for key in data:
        scalars[key] = float(data[key])

num_iters = 80

params0 = {
    "widths_si": np.full(num_elements, scalars["width_si"]),
    "gaps_si": np.full(num_elements, scalars["gap_si"]),
    "widths_sin": np.full(num_elements, scalars["width_sin"]),
    "gaps_sin": np.full(num_elements, scalars["gap_sin"]),
    "first_gap_si": scalars["first_gap_si"],
    "first_gap_sin": first_gap_sin,
}

bounds = {
    "widths_si": (min_width_si, max_width_si),
    "gaps_si": (min_gap_si, max_gap_si),
    "widths_sin": (min_width_sin, max_width_sin),
    "gaps_sin": (min_gap_sin, max_gap_sin),
    "first_gap_si": (None, None),
    "first_gap_sin": (min_gap_sin, None),
}

vg_fun = value_and_grad(objective)
params = deepcopy(params0)
opt_state = init_adam(params, lr=1e-2)
target_powers = []

for n in range(num_iters):
    value, grad = vg_fun(params)
    target_power = -value

    target_powers.append(target_power)
    print(f"iter {n}: target_power={target_power:.4f}")

    updates, opt_state = adam_update(grad, opt_state)
    params = apply_updates(params, updates)
    params = clip_params(params, bounds)

fig, ax = plt.subplots(figsize=(6, 4))
ax.plot(np.arange(len(target_powers)), target_powers, marker="o")
ax.set_xlabel("Iteration")
ax.set_ylabel("Target Power")
ax.set_title("Adjoint Optimization History")
ax.grid(True, alpha=0.3)
plt.show()

def compute_spectrum(param_set, task_name):
    sim = make_simulation(
        param_set["widths_si"],
        param_set["gaps_si"],
        param_set["widths_sin"],
        param_set["gaps_sin"],
        first_gap_si=param_set["first_gap_si"],
        first_gap_sin=param_set["first_gap_sin"],
    )
    sim_data = web.run(sim, task_name=task_name, verbose=False)
    power_da = get_mode_monitor_power(sim_data)
    freqs = power_da.coords["f"].values
    wavelengths = td.C_0 / freqs
    power = np.squeeze(power_da.data)
    sort_idx = np.argsort(wavelengths)
    wavelengths = wavelengths[sort_idx]
    power = np.array(power)[sort_idx]
    return wavelengths, power

w_before, p_before = compute_spectrum(params0, "gc_adjoint_before")
w_after, p_after = compute_spectrum(params, "gc_adjoint_after")

before_db = 10 * np.log10(p_before)
after_db = 10 * np.log10(p_after)

fig, ax = plt.subplots(figsize=(6, 4))
ax.plot(w_before, before_db, label="Before adjoint")
ax.plot(w_after, after_db, label="After adjoint")
ax.set_xlabel("Wavelength (µm)")
ax.set_ylabel("Transmission (dB)")
ax.set_title("Mode Monitor Spectrum (Before vs After)")
ax.grid(True, alpha=0.3)
ax.legend()
plt.show()

final_sim = make_simulation(
    params["widths_si"],
    params["gaps_si"],
    params["widths_sin"],
    params["gaps_sin"],
    first_gap_si=params["first_gap_si"],
    first_gap_sin=params["first_gap_sin"],
)
ax = final_sim.plot(y=0)
ax.set_title("Apodized grating geometry after adjoint optimization (y=0)")
plt.show()

def serialize_params(param_dict):
    """Detach autograd containers into JSON-serializable Python objects."""
    return {
        "widths_si": [float(value) for value in param_dict["widths_si"]],
        "gaps_si": [float(value) for value in param_dict["gaps_si"]],
        "widths_sin": [float(value) for value in param_dict["widths_sin"]],
        "gaps_sin": [float(value) for value in param_dict["gaps_sin"]],
        "first_gap_si": float(param_dict["first_gap_si"]),
        "first_gap_sin": float(param_dict["first_gap_sin"]),
    }


export_path = Path("./results/gc_adjoint_best.json")
export_path.parent.mkdir(parents=True, exist_ok=True)

payload = serialize_params(params)
payload["target_power"] = float(target_powers[-1]) if target_powers else None

with export_path.open("w", encoding="utf-8") as f:
    json.dump(payload, f, indent=2)

print(f"Saved adjoint design to {export_path.resolve()}")