import json
from pathlib import Path

import autograd.numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import tidy3d as td
from autograd import value_and_grad
from scipy.stats import norm
from setup import (
    center_wavelength,
    default_spacer_thickness,
    get_mode_monitor_power,
    make_simulation,
)
from tidy3d import web

def load_nominal_parameters(path):
    """Load a design JSON (Bayes or adjoint) into numpy-friendly fields."""
    data = json.loads(Path(path).read_text(encoding="utf-8"))
    return {
        "widths_si": np.array(data["widths_si"]),
        "gaps_si": np.array(data["gaps_si"]),
        "widths_sin": np.array(data["widths_sin"]),
        "gaps_sin": np.array(data["gaps_sin"]),
        "first_gap_si": data["first_gap_si"],
        "first_gap_sin": data["first_gap_sin"],
        "spacer_thickness": default_spacer_thickness,
    }

def make_variation_builder(nominal):
    """Return a closure that maps process deltas to a tidy3d Simulation."""
    base_widths_si = np.array(nominal["widths_si"])
    base_gaps_si = np.array(nominal["gaps_si"])

    def builder(overlay_delta=0.0, spacer_delta=0.0, etch_bias=0.0):
        # Etch bias widens features when positive and narrows them when
        # negative, so widths grow with the bias while gaps shrink, mirroring
        # the fabrication effect of over/under etching.
        pert_widths_si = base_widths_si + etch_bias
        pert_gaps_si = base_gaps_si - etch_bias

        return make_simulation(
            pert_widths_si,
            pert_gaps_si,
            nominal["widths_sin"],
            nominal["gaps_sin"],
            first_gap_si=nominal["first_gap_si"] + overlay_delta,
            first_gap_sin=nominal["first_gap_sin"],
            spacer_thickness=nominal["spacer_thickness"] + spacer_delta,
        )

    return builder

design_path = Path("./results") / "gc_adjoint_best.json"

# Load the best apodized design from the previous notebook.
# This will be our nominal, or central, design point for the analysis.
nominal = load_nominal_parameters(design_path)
builder = make_variation_builder(nominal)

# Define the fabrication bias in microns (20 nm).
bias = 0.02

# Create simulations for each fabrication scenario: over-etched, nominal,
# and under-etched. Positive bias widens features, while a negative bias
# corresponds to over-etching that narrows them.
bias_cases = {
    "Over-etched (-20 nm)": builder(etch_bias=-bias),
    "Nominal": builder(),
    "Under-etched (+20 nm)": builder(etch_bias=bias),
}

bias_data = web.run_async(bias_cases, verbose=False)

bias_wavelengths = None
bias_spectra = {}

for label, sim_data in bias_data.items():
    power_da = get_mode_monitor_power(sim_data)
    freqs = power_da.coords["f"].values
    wavelengths = td.C_0 / freqs
    power = np.asarray(power_da.data).squeeze()
    order = np.argsort(wavelengths)
    wavelengths = wavelengths[order]
    power = power[order]

    if bias_wavelengths is None:
        bias_wavelengths = wavelengths

    bias_spectra[label] = power

fig, ax = plt.subplots(figsize=(6, 4))
colors = {
    "Over-etched (-20 nm)": "tab:orange",
    "Nominal": "tab:green",
    "Under-etched (+20 nm)": "tab:blue",
}

for label, spectrum in bias_spectra.items():
    ax.plot(
        bias_wavelengths,
        10 * np.log10(spectrum),
        label=label,
        color=colors.get(label, None),
        linewidth=2 if label == "Nominal" else 1.8,
        alpha=0.9,
    )

ax.set_xlabel("Wavelength (µm)")
ax.set_ylabel("Transmission (dB)")
ax.set_title("Impact of ±20 nm Fabrication Bias")
ax.legend()
ax.grid(True, alpha=0.3)
plt.show()

sigma_spec = {
    "overlay": 0.025,
    "spacer": 0.02,
    "widths_si": 0.01,
}

seed = 42
num_mc_samples = 100
design_path = Path("./results") / "gc_adjoint_best.json"

nominal = load_nominal_parameters(design_path)
builder = make_variation_builder(nominal)

sigma_vector = np.array([sigma_spec["overlay"], sigma_spec["spacer"], sigma_spec["widths_si"]])
rng = np.random.default_rng(seed)
samples = rng.standard_normal(size=(num_mc_samples, len(sigma_vector))) * sigma_vector

sims = {"nominal": builder()}
sims.update({f"sample_{idx + 1}": builder(*tuple(sample)) for idx, sample in enumerate(samples)})

batch_data = web.run_async(sims, verbose=False)

ordered_names = list(sims.keys())
wavelengths = None
linear_spectra = []

for name in ordered_names:
    sim_data = batch_data[name]
    power_da = get_mode_monitor_power(sim_data)
    freqs = power_da.coords["f"].values
    wl = td.C_0 / freqs
    power = np.asarray(power_da.data).squeeze()
    order = np.argsort(wl)
    wl = wl[order]
    power = power[order]

    if wavelengths is None:
        wavelengths = wl

    linear_spectra.append(power)

linear_array = np.vstack(linear_spectra)
nominal_index = ordered_names.index("nominal")
nominal_spectrum = linear_array[nominal_index]

fig, ax = plt.subplots(figsize=(6, 4))

for name, spectrum in zip(ordered_names, linear_array):
    if name == "nominal":
        continue
    spectrum_db = 10 * np.log10(spectrum)
    ax.plot(
        wavelengths,
        spectrum_db,
        color="lightgray",
        alpha=0.6,
        linewidth=1,
        zorder=1,
    )

mean_spectrum = linear_array.mean(axis=0)
p10_spectrum = np.percentile(linear_array, 10, axis=0)
p90_spectrum = np.percentile(linear_array, 90, axis=0)

ax.fill_between(
    wavelengths,
    10 * np.log10(p10_spectrum),
    10 * np.log10(p90_spectrum),
    color="tab:blue",
    alpha=0.18,
    label="P10–P90",
    zorder=2,
)
ax.plot(
    wavelengths,
    10 * np.log10(mean_spectrum),
    color="tab:blue",
    linewidth=2,
    linestyle="--",
    label="Mean",
    zorder=3,
)
ax.plot(
    wavelengths,
    10 * np.log10(nominal_spectrum),
    color="black",
    linewidth=2.5,
    label="Nominal",
    zorder=4,
)

ax.set_xlabel("Wavelength (µm)")
ax.set_ylabel("Transmission (dB)")
ax.set_title("Monte Carlo Ensemble of Spectra")
ax.legend()
ax.grid(True, alpha=0.3)
plt.show()

def linear_to_loss_db(x):
    return -10 * np.log10(x)


idx_center = np.argmin(np.abs(wavelengths - center_wavelength))
eta_center = linear_array[:, idx_center]

mc_summary = {
    "wavelength_um": wavelengths[idx_center],
    "mean_linear": eta_center.mean(),
    "mean_db": linear_to_loss_db(eta_center.mean()),
    "std_linear": eta_center.std(ddof=1),
    "std_db": linear_to_loss_db(eta_center.mean() - eta_center.std(ddof=1))
    - linear_to_loss_db(eta_center.mean()),
    "p10_linear": np.percentile(eta_center, 10),
    "p10_db": linear_to_loss_db(np.percentile(eta_center, 10)),
    "p90_linear": np.percentile(eta_center, 90),
    "p90_db": linear_to_loss_db(np.percentile(eta_center, 90)),
    "sigma_overlay": sigma_spec["overlay"],
    "sigma_spacer": sigma_spec["spacer"],
    "sigma_widths_si": sigma_spec["widths_si"],
}

pd.Series(mc_summary).to_frame("value")

eta_center_db = linear_to_loss_db(eta_center)
fig, ax = plt.subplots(figsize=(6, 4))
ax.hist(eta_center_db[1:], bins="auto", color="tab:blue", alpha=0.7, label="Samples")
ax.axvline(eta_center_db[0], color="black", linewidth=2, label="Nominal")
ax.set_xlabel(f"Transmission at {wavelengths[idx_center]:.3f} µm (dB)")
ax.set_ylabel("Count")
ax.set_title("Monte Carlo Transmission at Center Wavelength")
ax.grid(True, alpha=0.3)
ax.legend()
plt.show()