# Imports, runtime status, and global output paths.
from __future__ import annotations

import os
import random
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

import autograd as ag
import autograd.numpy as anp
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn

import tidy3d as td
import tidy3d.web as web
from tidy3d.plugins.autograd import (
    make_erosion_dilation_penalty,
    make_filter_and_project,
    rescale,
    smooth_min,
)

ASSET_DIR = Path("assets")
OUT_DIR = Path("outputs/wdm_2d_cnn_ig")
FIG_DIR = OUT_DIR / "figs"
IG_DIR = OUT_DIR / "ig"
for p in (ASSET_DIR, OUT_DIR, FIG_DIR, IG_DIR):
    p.mkdir(parents=True, exist_ok=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Python runtime ready")
print("Tidy3D:", getattr(td, "__version__", "unknown"))
print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("CUDA device:", torch.cuda.get_device_name(0))
print("Device used for CNN:", device)

# Optional trained CNN checkpoint for the CNN + IG section.
# A trained CNN checkpoint is required for CNN prediction and Integrated Gradients,
# but it is not included with this public notebook.
#
# To enable the CNN + IG cells, provide your own compatible checkpoint by either:
#   1. setting CNN_CHECKPOINT_PATH, or
#   2. changing CHECKPOINT_PATH below.
CHECKPOINT_PATH = Path(
    os.environ.get("CNN_CHECKPOINT_PATH", "your_trained_cnn_checkpoint.pt")
)
RUN_CNN_IG = CHECKPOINT_PATH.exists()

if RUN_CNN_IG:
    print(
        f"Using CNN checkpoint: {CHECKPOINT_PATH} ({CHECKPOINT_PATH.stat().st_size:,} bytes)"
    )
else:
    print(
        "A trained CNN checkpoint is required for the CNN + IG section but is not included."
    )
    print(
        "Set CNN_CHECKPOINT_PATH or edit CHECKPOINT_PATH to point to your own compatible checkpoint."
    )

# CNN definition and optional checkpoint load.
# The architecture is included so the IG workflow is clear, but model weights are not included.
class SmallCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(16, 16, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(),
        )
        self.head = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(64, 1)

    def forward(self, x):
        z = self.net(x)
        z = self.head(z).squeeze(-1).squeeze(-1)
        return self.fc(z)


def load_state_dict(path: Path):
    try:
        return torch.load(path, map_location=device, weights_only=True)
    except TypeError:
        return torch.load(path, map_location=device)


def to_bchw_1ch(x_np, device=device):
    x = torch.as_tensor(x_np, dtype=torch.float32)
    if x.ndim == 2:
        x = x.unsqueeze(0).unsqueeze(0)
    elif x.ndim == 3:
        x = x.unsqueeze(0)
    elif x.ndim == 4:
        pass
    else:
        raise ValueError(f"Unexpected input shape: {tuple(x.shape)}")
    mx = torch.max(x)
    if mx > 0:
        x = x / mx
    return x.to(device)


model = SmallCNN().to(device)
if RUN_CNN_IG:
    state = load_state_dict(CHECKPOINT_PATH)
    missing, unexpected = model.load_state_dict(state, strict=False)
    print("Missing checkpoint keys:", missing)
    print("Unexpected checkpoint keys:", unexpected)
    assert not missing and not unexpected, "Checkpoint does not match SmallCNN."
    model.eval()
    print("SmallCNN parameters:", sum(p.numel() for p in model.parameters()))
else:
    model.eval()
    print("SmallCNN architecture is defined, but no trained weights were loaded.")

# Main design/runtime configuration.
RANDOM_SEED = 7
NUM_OPT_STEPS = 25  # Increase for a better final device; reduce for a faster demo.
LEARNING_RATE = 0.10
BETA_MIN = 1.0
BETA_MAX = 50.0
PENALTY_WEIGHT = 1.0
LEAK_WEIGHT_START = 0.0
LEAK_WEIGHT_END = 1.0

# 2D WDM material constants, matching the Flexcompute 2D example style.
# In the 2D example, SiO2 is documented but not used as a substrate/background.
n_si = 3.49
n_sio2 = 1.45  # not used in this 2D simulation
n_air = 1.0

# Physical WDM target wavelengths in microns.
wvls_design = np.array([1.310, 1.550])
channel_names = ["1310 nm", "1550 nm"]

# Geometry.
lx = 2.0  # design region x span, um
ly = 2.0  # design region y span, um
ly_single = ly / len(wvls_design)
lz = td.inf  # official 2D convention: structures are infinite in z
wg_width = 0.50  # waveguide width, um
wg_length = 1.50  # straight port length on each side, um
output_port_separation = 1.10  # center-to-center output spacing, um
cnn_span = 5.0  # final CNN raster window, um
cnn_pixels = 125  # 5 um / 0.04 um = 125
cnn_pixel_um = cnn_span / cnn_pixels

# Optimization grid. This matches the CNN pixel pitch inside the 2 um design region.
dl_design_region = 0.040
nx = int(round(lx / dl_design_region))
ny = int(round(ly / dl_design_region))
assert nx == 50 and ny == 50

# Fabrication-style smoothness constraint.
filter_radius = 0.100
beta_penalty = 10.0

# Simulation domain. z-size 0 means a 2D simulation.
Lx = lx + wg_length * 2
Ly = ly + (cnn_span - lx)  # 2 um design + 1.5 um top/bottom buffers = 5 um
Lz = 0.0
assert abs(Lx - cnn_span) < 1e-12 and abs(Ly - cnn_span) < 1e-12
min_steps_per_wvl = 18

np.random.seed(RANDOM_SEED)
random.seed(RANDOM_SEED)
print(
    f"Optimization grid: {nx} x {ny}, design pixel = {dl_design_region * 1000:.0f} nm"
)
print(f"CNN raster: {cnn_pixels} x {cnn_pixels}, pixel = {cnn_pixel_um * 1000:.0f} nm")
print("Wavelengths:", [f"{w * 1000:.0f} nm" for w in wvls_design])
print("Tidy3D dimensionality check: Lz =", Lz, ", lz = td.inf")

# Build the static 2D Tidy3D WDM simulation: ports, monitors, source.
freqs_design = td.C_0 / wvls_design
num_freqs_design = len(freqs_design)
freq_max = np.max(freqs_design)
freq_min = np.min(freqs_design)
freq0 = np.mean(freqs_design)
df_design = abs(np.mean(np.diff(freqs_design)))
fwidth = freq_max - freq_min
run_time = 120 / fwidth

channel_fwidth = df_design / 2.0
channel_bounds = [
    (f - channel_fwidth / 2, f + channel_fwidth / 2) for f in freqs_design
]
num_freqs_channel = 5
channel_freqs = []
for fmin, fmax in channel_bounds:
    channel_freqs.extend(np.linspace(fmin, fmax, num_freqs_channel).tolist())

medium_si = td.Medium(permittivity=n_si**2)

if num_freqs_design == 2:
    centers_y = np.array([-output_port_separation / 2, output_port_separation / 2])
else:
    centers_y = np.linspace(
        -ly / 2.0 + ly_single / 2.0, +ly / 2.0 - ly_single / 2.0, num_freqs_design
    )

# Official 2D convention: structures have size td.inf in z while the simulation has Lz = 0.
wg_in = td.Structure(
    geometry=td.Box(center=(-Lx / 2, 0, 0), size=(wg_length * 2, wg_width, lz)),
    medium=medium_si,
)

wgs_out = []
for cy in centers_y:
    wgs_out.append(
        td.Structure(
            geometry=td.Box(
                center=(+Lx / 2, cy, 0), size=(wg_length * 2, wg_width, lz)
            ),
            medium=medium_si,
        )
    )

mode_span_y = max(0.9 * ly_single, 1.8 * wg_width)
mode_size = (0, mode_span_y, td.inf)
mode_spec = td.ModeSpec(num_modes=1)

mode_src = td.ModeSource(
    center=(-Lx / 2 + wg_length / 2, 0, 0),
    size=mode_size,
    source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth),
    direction="+",
    mode_index=0,
    mode_spec=mode_spec,
)

mnts_mode = []
mnts_flux = []
for i, cy in enumerate(centers_y):
    mnts_mode.append(
        td.ModeMonitor(
            center=(Lx / 2 - wg_length / 2, cy, 0),
            size=mode_size,
            freqs=channel_freqs,
            mode_spec=mode_spec,
            name=f"mode_{i}",
        )
    )
    mnts_flux.append(
        td.FluxMonitor(
            center=(Lx / 2 - wg_length / 2, cy, 0),
            size=mode_size,
            freqs=channel_freqs,
            name=f"flux_{i}",
        )
    )

fld_mnt = td.FieldMonitor(
    center=(0, 0, 0), size=(td.inf, td.inf, 0), freqs=freqs_design, name="field"
)

sim_static = td.Simulation(
    size=(Lx, Ly, Lz),
    grid_spec=td.GridSpec.auto(
        min_steps_per_wvl=min_steps_per_wvl, wavelength=np.min(wvls_design)
    ),
    structures=[wg_in] + wgs_out,
    sources=[mode_src],
    monitors=mnts_mode + mnts_flux + [fld_mnt],
    boundary_spec=td.BoundarySpec.pml(x=True, y=True, z=True if Lz else False),
    run_time=run_time,
)

print("2D simulation built:")
print("  sim size =", sim_static.size)
print("  mode_size =", mode_size)
print("  centers_y =", centers_y)
print("  boundary_spec =", sim_static.boundary_spec)

fig, ax = plt.subplots(1, 1, figsize=(6, 5), dpi=150)
sim_static.plot_eps(z=0, ax=ax)
ax.set_aspect("equal")
ax.set_title("Static 2D simulation: input/output ports")
plt.show()

# Differentiable design-region parameterization.
filter_project = make_filter_and_project(filter_radius, dl_design_region)
design_region_geo = td.Box(size=(lx, ly, lz), center=(0, 0, 0))


def get_density(params: np.ndarray, beta: float) -> np.ndarray:
    return filter_project(params, beta=beta)


def make_eps(params: np.ndarray, beta: float) -> np.ndarray:
    density = get_density(params, beta=beta)
    # Official 2D WDM example rescales between air and silicon.
    return rescale(density, n_air**2, n_si**2)


def make_custom_medium(params: np.ndarray, beta: float) -> td.Structure:
    eps = make_eps(params, beta).reshape((nx, ny, 1))
    xs = anp.linspace(-lx / 2, lx / 2, nx)
    ys = anp.linspace(-ly / 2, ly / 2, ny)
    coords = dict(x=xs, y=ys, z=[0.0])
    eps_arr = td.ScalarFieldDataArray(data=eps, coords=coords)
    medium = td.CustomMedium(permittivity=eps_arr)
    return td.Structure(geometry=design_region_geo, medium=medium)


def get_sim(params, beta, include_extra_mnts: bool = True):
    design_region = make_custom_medium(params, beta=beta)
    design_override = td.MeshOverrideStructure(
        geometry=design_region.geometry,
        dl=[dl_design_region, dl_design_region, dl_design_region],
    )
    grid_spec = sim_static.grid_spec.updated_copy(
        override_structures=list(sim_static.grid_spec.override_structures)
        + [design_override]
    )
    update_dict = dict(
        structures=list(sim_static.structures) + [design_region], grid_spec=grid_spec
    )
    if not include_extra_mnts:
        update_dict["monitors"] = mnts_mode
    return sim_static.updated_copy(**update_dict)


rng = np.random.default_rng(RANDOM_SEED)
params0 = rng.random((nx, ny))
params = np.array(params0, dtype=float)

fig, axs = plt.subplots(1, 2, figsize=(8, 3.5), dpi=150, constrained_layout=True)
axs[0].imshow(params0.T, origin="lower", cmap="gray", vmin=0, vmax=1)
axs[0].set_title("Initial raw parameters")
axs[0].axis("off")
axs[1].imshow(
    get_density(params0, beta=BETA_MIN).T, origin="lower", cmap="gray", vmin=0, vmax=1
)
axs[1].set_title("Initial filtered density")
axs[1].axis("off")
plt.show()

sim0 = get_sim(params0, beta=BETA_MIN)
fig, ax = plt.subplots(1, 1, figsize=(6, 5), dpi=150)
sim0.plot_eps(z=0, ax=ax, monitor_alpha=0.15, source_alpha=0.5)
ax.set_aspect("equal")
ax.set_title("Initial 2D WDM permittivity")
plt.show()

# Objective, metrics, and autograd value/gradient wrapper.
import xarray as xr


def average_over_channel(spectrum: xr.DataArray, fmin: float, fmax: float):
    freqs = spectrum.f
    freqs_in_channel = np.logical_and(freqs >= fmin, freqs <= fmax).values
    num_freqs = np.sum(freqs_in_channel)
    return spectrum.values @ freqs_in_channel / num_freqs


def get_power(sim_data: td.SimulationData, mnt_index: int, freq_index: int):
    mnt_name = mnts_mode[mnt_index].name
    mnt_data = sim_data[mnt_name]
    fmin_channel, fmax_channel = channel_bounds[freq_index]
    amp = mnt_data.amps.sel(direction="+", mode_index=0)
    power_spectrum = anp.abs(amp) ** 2
    return average_over_channel(power_spectrum, fmin=fmin_channel, fmax=fmax_channel)


def get_metric(sim_data: td.SimulationData, mnt_index: int, leak_weight: float = 1.0):
    power_all = [
        get_power(sim_data=sim_data, mnt_index=mnt_index, freq_index=j)
        for j in range(num_freqs_design)
    ]
    power_transmitted = power_all[mnt_index]
    power_leaked = sum(power_all) - power_transmitted
    avg_power_leaked = power_leaked / (num_freqs_design - 1)
    return power_transmitted - leak_weight * avg_power_leaked


penalty = make_erosion_dilation_penalty(
    filter_radius, dl_design_region, beta=beta_penalty
)


def objective(
    params, beta: float, penalty_weight: float = 1.0, leak_weight: float = 0.0
):
    sim = get_sim(params, beta=beta, include_extra_mnts=False)
    sim_data = web.run(sim, task_name="wdm_2d_opt_step", verbose=False)
    metrics = []
    for mnt_index in range(num_freqs_design):
        metrics.append(
            get_metric(sim_data=sim_data, mnt_index=mnt_index, leak_weight=leak_weight)
        )
    metric = smooth_min(anp.array(metrics))
    penalty_value = penalty(params)
    return metric - penalty_weight * penalty_value


grad_fn = ag.value_and_grad(objective)
print("Objective and gradient function are ready.")

# Run the 2D adjoint optimization.
# This cell submits Tidy3D jobs. It is the slow/cost-bearing part of the notebook.
def adam_update(params, grad, state, lr=0.1, beta1=0.9, beta2=0.999, eps=1e-8):
    t, m, v = state
    t += 1
    m = beta1 * m + (1 - beta1) * grad
    v = beta2 * v + (1 - beta2) * (grad * grad)
    m_hat = m / (1 - beta1**t)
    v_hat = v / (1 - beta2**t)
    params = params + lr * m_hat / (np.sqrt(v_hat) + eps)
    params = np.clip(params, 0.0, 1.0)
    return params, (t, m, v)


params = np.array(params0, dtype=float)
adam_state = (0, np.zeros_like(params), np.zeros_like(params))
Js = []
params_history = [params.copy()]
beta_history = []
grad_norm_history = []

for i in range(NUM_OPT_STEPS):
    perc_done = 0.0 if NUM_OPT_STEPS == 1 else i / (NUM_OPT_STEPS - 1)
    beta_i = BETA_MIN * (1 - perc_done) + BETA_MAX * perc_done
    leak_weight = LEAK_WEIGHT_START * (1 - perc_done) + LEAK_WEIGHT_END * perc_done

    value, gradient = grad_fn(
        params,
        beta=beta_i,
        penalty_weight=PENALTY_WEIGHT,
        leak_weight=leak_weight,
    )
    value = float(value)
    gradient = np.array(gradient, dtype=float)
    grad_norm = float(np.linalg.norm(gradient))

    params, adam_state = adam_update(params, gradient, adam_state, lr=LEARNING_RATE)

    Js.append(value)
    beta_history.append(beta_i)
    grad_norm_history.append(grad_norm)
    params_history.append(params.copy())

    print(
        f"step {i + 1:03d}/{NUM_OPT_STEPS} | J={value:.5e} | beta={beta_i:.2f} | leak={leak_weight:.2f} | grad_norm={grad_norm:.3e}"
    )

    if (
        (i == 0)
        or ((i + 1) % max(1, NUM_OPT_STEPS // 5) == 0)
        or (i + 1 == NUM_OPT_STEPS)
    ):
        density_i = get_density(params, beta_i)
        plt.figure(figsize=(3, 3), dpi=140)
        plt.imshow(density_i.T, origin="lower", cmap="gray", vmin=0, vmax=1)
        plt.title(f"Density after step {i + 1}")
        plt.axis("off")
        plt.show()

fig, axs = plt.subplots(1, 2, figsize=(10, 3.6), dpi=150, constrained_layout=True)
axs[0].plot(np.arange(1, len(Js) + 1), Js, marker="o")
axs[0].set_title("Optimization objective")
axs[0].set_xlabel("Step")
axs[0].set_ylabel("J")
axs[0].grid(True, alpha=0.3)
axs[1].plot(
    np.arange(1, len(grad_norm_history) + 1),
    grad_norm_history,
    marker="o",
    color="tab:orange",
)
axs[1].set_title("Gradient norm")
axs[1].set_xlabel("Step")
axs[1].set_ylabel("||grad||")
axs[1].grid(True, alpha=0.3)
plt.show()

# Final Tidy3D evaluation with spectra and fields.
params_final = params_history[-1]
beta_final = beta_history[-1] if beta_history else BETA_MAX

density_final = np.array(get_density(params_final, beta_final)).reshape(nx, ny)
binary_design = (density_final >= 0.5).astype(np.uint8)

num_freqs_measure = 151
freqs_measure = np.linspace(
    freq_min - df_design, freq_max + df_design, num_freqs_measure
)
sim_final = get_sim(params_final, beta=beta_final, include_extra_mnts=True)
for i in range(num_freqs_design):
    sim_final = sim_final.updated_copy(freqs=freqs_measure, path=f"monitors/{i}")
    sim_final = sim_final.updated_copy(
        freqs=freqs_measure, path=f"monitors/{i + num_freqs_design}"
    )

fig, axs = plt.subplots(1, 3, figsize=(13, 3.8), dpi=160, constrained_layout=True)
axs[0].imshow(density_final.T, origin="lower", cmap="gray", vmin=0, vmax=1)
axs[0].set_title("Final density")
axs[0].axis("off")
axs[1].imshow(binary_design.T, origin="lower", cmap="gray_r", vmin=0, vmax=1)
axs[1].set_title("Final binary design")
axs[1].axis("off")
sim_final.plot_eps(z=0, ax=axs[2], monitor_alpha=0, source_alpha=0)
axs[2].set_title("Final 2D permittivity")
axs[2].set_aspect("equal")
plt.show()

print("Submitting final evaluation simulation...")
sim_data_final = web.run(sim_final, task_name="wdm_2d_final_eval", verbose=False)

plt.figure(figsize=(8, 4), dpi=150)
final_powers = {}
for i, freq in enumerate(freqs_design):
    mnt = sim_data_final[mnts_mode[i].name]
    powers = np.abs(mnt.amps.sel(direction="+", mode_index=0)) ** 2
    freqs_sim = np.array(powers.f)
    wavelengths_nm = 1000 * td.C_0 / freqs_sim
    power_at_design = float(powers.sel(f=freq, method="nearest"))
    final_powers[channel_names[i]] = power_at_design
    loss_db_spectrum = 10 * np.log10(np.maximum(np.array(powers.values), 1e-12))
    loss_db_design = 10 * np.log10(max(power_at_design, 1e-12))
    wvl_nm = 1000 * td.C_0 / freq
    fmin, fmax = channel_bounds[i]
    plt.axvspan(1000 * td.C_0 / fmin, 1000 * td.C_0 / fmax, alpha=0.12)
    plt.plot(
        wavelengths_nm, loss_db_spectrum, label=f"output {i} target {channel_names[i]}"
    )
    plt.scatter([wvl_nm], [loss_db_design], 80, marker="*")
plt.xlabel("Wavelength (nm)")
plt.ylabel("Mode power (dB)")
plt.title("Final WDM modal transmission")
plt.grid(True, alpha=0.3)
plt.legend()
plt.show()
print("Final design powers:", final_powers)

fig, axes = plt.subplots(
    1,
    num_freqs_design,
    figsize=(6 * num_freqs_design, 4.5),
    dpi=170,
    constrained_layout=True,
)
if num_freqs_design == 1:
    axes = [axes]
for freq, name, ax in zip(freqs_design, channel_names, axes):
    sim_data_final.plot_field("field", "E", "abs^2", f=freq, ax=ax)
    ax.set_title(f"Field intensity at {name}")
    ax.set_xlabel("x (um)")
    ax.set_ylabel("y (um)")
plt.show()

# Convert final optimized 2D design to the CNN training-mask convention: 125 x 125 over 5 um x 5 um.
def rasterize_for_cnn(binary_design: np.ndarray) -> np.ndarray:
    assert binary_design.shape == (nx, ny)
    pixel = cnn_span / cnn_pixels
    xs = -cnn_span / 2 + (np.arange(cnn_pixels) + 0.5) * pixel
    ys = -cnn_span / 2 + (np.arange(cnn_pixels) + 0.5) * pixel
    X, Y = np.meshgrid(xs, ys, indexing="ij")
    mask = np.zeros((cnn_pixels, cnn_pixels), dtype=np.uint8)

    in_design = (np.abs(X) <= lx / 2) & (np.abs(Y) <= ly / 2)
    ix = np.floor((X[in_design] + lx / 2) / lx * nx).astype(int)
    iy = np.floor((Y[in_design] + ly / 2) / ly * ny).astype(int)
    ix = np.clip(ix, 0, nx - 1)
    iy = np.clip(iy, 0, ny - 1)
    mask[in_design] = binary_design[ix, iy]

    input_port = (X >= -cnn_span / 2) & (X <= -lx / 2) & (np.abs(Y) <= wg_width / 2)
    mask[input_port] = 1
    for cy in centers_y:
        out_port = (
            (X >= lx / 2) & (X <= cnn_span / 2) & (np.abs(Y - cy) <= wg_width / 2)
        )
        mask[out_port] = 1
    return mask


cnn_mask = rasterize_for_cnn(binary_design)
assert cnn_mask.shape == (125, 125)
assert set(np.unique(cnn_mask)).issubset({0, 1})
np.save(OUT_DIR / "optimized_wdm_cnn_mask.npy", cnn_mask)

fig, axs = plt.subplots(1, 3, figsize=(12, 3.8), dpi=160, constrained_layout=True)
axs[0].imshow(density_final.T, origin="lower", cmap="gray", vmin=0, vmax=1)
axs[0].set_title("Optimized density\n2 um design region")
axs[0].axis("off")
axs[1].imshow(binary_design.T, origin="lower", cmap="gray_r", vmin=0, vmax=1)
axs[1].set_title(f"Optimized binary design\n{nx} x {ny}")
axs[1].axis("off")
axs[2].imshow(cnn_mask, cmap="gray_r", origin="upper", vmin=0, vmax=1)
axs[2].set_title(
    f"CNN-ready optimized mask\n{cnn_mask.shape}, fill={cnn_mask.mean():.3f}"
)
axs[2].axis("off")
plt.suptitle("Optimized WDM design to CNN-ready mask")
plt.show()

print("CNN mask shape:", cnn_mask.shape)
print("CNN mask fill fraction:", float(cnn_mask.mean()))
print("Saved:", OUT_DIR / "optimized_wdm_cnn_mask.npy")

# Simple Integrated Gradients through the CNN.
def integrated_gradients(
    model: torch.nn.Module,
    x: torch.Tensor,
    target_fn: Callable[[torch.Tensor], torch.Tensor] = lambda y: y.squeeze(1),
    baseline: Optional[torch.Tensor] = None,
    steps: int = 1000,
    batch_size: int = 64,
) -> torch.Tensor:
    model.eval()
    x = x.detach().to(device)
    if baseline is None:
        baseline = torch.zeros_like(x, device=device)
    else:
        baseline = baseline.detach().to(device)

    alphas = torch.linspace(0.0, 1.0, steps, device=device)
    weights = torch.ones(steps, device=device)
    weights[0] = 0.5
    weights[-1] = 0.5
    weights = weights / (steps - 1)
    total_grad = torch.zeros_like(x, device=device)

    for start in range(0, steps, batch_size):
        end = min(steps, start + batch_size)
        a = alphas[start:end].view(-1, 1, 1, 1)
        w = weights[start:end].view(-1, 1, 1, 1)
        x_batch = baseline + a * (x - baseline)
        x_batch.requires_grad_(True)
        out = model(x_batch)
        tgt = target_fn(out)
        grads = torch.autograd.grad(
            tgt,
            x_batch,
            torch.ones_like(tgt),
            retain_graph=False,
            create_graph=False,
        )[0]
        total_grad += (grads * w).sum(dim=0, keepdim=True)
    return (x - baseline) * total_grad


IG_STEPS = 1000
if RUN_CNN_IG:
    ig_opt = integrated_gradients(model, x_opt, steps=IG_STEPS, batch_size=64)
    ig_np = ig_opt.squeeze().detach().cpu().numpy()
    print("IG shape:", ig_np.shape)
    print("IG finite:", bool(np.isfinite(ig_np).all()))
    print("IG min/max:", float(ig_np.min()), float(ig_np.max()))
    assert np.isfinite(ig_np).all()
else:
    ig_np = None
    print("Skipping IG computation: trained checkpoint not supplied.")

# IG visualizations: raw, absolute, hotspot overlay, and signed contribution overlay.
if RUN_CNN_IG:

    def normalize_01(arr, eps=1e-12):
        arr = np.asarray(arr, dtype=float)
        mn, mx = float(arr.min()), float(arr.max())
        if mx <= mn + eps:
            return np.zeros_like(arr)
        return (arr - mn) / (mx - mn + eps)

    x_np = x_opt.squeeze().detach().cpu().numpy()
    ig_abs = np.abs(ig_np)
    clip_v = max(float(np.percentile(ig_abs, 99.7)), 1e-12)
    hot_thr = float(np.percentile(ig_abs, 99.7))
    pos_vals = ig_np[ig_np > 0]
    neg_vals = -ig_np[ig_np < 0]
    pos_thr = float(np.percentile(pos_vals, 99.5)) if pos_vals.size else np.inf
    neg_thr = float(np.percentile(neg_vals, 99.5)) if neg_vals.size else np.inf
    pos_hot = (ig_np > 0) & (ig_np >= pos_thr)
    neg_hot = (ig_np < 0) & ((-ig_np) >= neg_thr)

    fig, axs = plt.subplots(2, 3, figsize=(14, 8.5), dpi=180, constrained_layout=True)
    axs[0, 0].imshow(x_np, cmap="gray_r", origin="upper", vmin=0, vmax=1)
    axs[0, 0].set_title("CNN input mask")
    axs[0, 0].axis("off")
    im1 = axs[0, 1].imshow(
        ig_np, cmap="seismic", origin="upper", vmin=-clip_v, vmax=clip_v
    )
    axs[0, 1].set_title("Raw signed IG")
    axs[0, 1].axis("off")
    plt.colorbar(im1, ax=axs[0, 1], fraction=0.046, pad=0.04)
    im2 = axs[0, 2].imshow(
        np.clip(ig_abs, 0, clip_v), cmap="magma", origin="upper", vmin=0, vmax=clip_v
    )
    axs[0, 2].set_title("|IG| clipped at 99.7 percentile")
    axs[0, 2].axis("off")
    plt.colorbar(im2, ax=axs[0, 2], fraction=0.046, pad=0.04)

    axs[1, 0].imshow(
        np.ma.masked_where(ig_abs < hot_thr, ig_abs),
        cmap="magma",
        origin="upper",
        vmin=0,
        vmax=clip_v,
    )
    axs[1, 0].set_title("Top |IG| hotspots")
    axs[1, 0].axis("off")
    axs[1, 1].imshow(x_np, cmap="gray_r", origin="upper", vmin=0, vmax=1)
    heat = np.ma.masked_where(
        ig_abs < hot_thr, normalize_01(np.clip(ig_abs, 0, clip_v))
    )
    axs[1, 1].imshow(heat, cmap="magma", origin="upper", alpha=0.72, vmin=0, vmax=1)
    axs[1, 1].contour(
        (ig_abs >= hot_thr).astype(float), levels=[0.5], colors="red", linewidths=1.0
    )
    axs[1, 1].set_title("Hotspot overlay on optimized mask")
    axs[1, 1].axis("off")
    axs[1, 2].imshow(x_np, cmap="gray_r", origin="upper", vmin=0, vmax=1)
    axs[1, 2].imshow(
        np.ma.masked_where(~pos_hot, pos_hot.astype(float)),
        cmap="Reds",
        origin="upper",
        alpha=0.75,
        vmin=0,
        vmax=1,
    )
    axs[1, 2].imshow(
        np.ma.masked_where(~neg_hot, neg_hot.astype(float)),
        cmap="Blues",
        origin="upper",
        alpha=0.75,
        vmin=0,
        vmax=1,
    )
    axs[1, 2].contour(pos_hot.astype(float), levels=[0.5], colors="red", linewidths=0.8)
    axs[1, 2].contour(
        neg_hot.astype(float), levels=[0.5], colors="blue", linewidths=0.8
    )
    axs[1, 2].set_title("Signed tails: red positive, blue negative")
    axs[1, 2].axis("off")
    plt.suptitle("CNN IG for 1550 nm CNN output")
    plt.show()

    np.save(IG_DIR / "optimized_wdm_ig_power1550.npy", ig_np)
    print("Saved IG array:", IG_DIR / "optimized_wdm_ig_power1550.npy")
else:
    print("Skipping IG visualizations: trained checkpoint not supplied.")

# Alpha-path story and completeness trace for the same CNN IG path.
if RUN_CNN_IG:

    def ig_path_story(
        model: torch.nn.Module,
        x: torch.Tensor,
        target_fn: Callable[[torch.Tensor], torch.Tensor] = lambda y: y.squeeze(1),
        baseline: Optional[torch.Tensor] = None,
        steps: int = 201,
    ) -> Dict[str, Any]:
        model.eval()
        x = x.detach().to(device)
        if baseline is None:
            baseline = torch.zeros_like(x, device=device)
        else:
            baseline = baseline.detach().to(device)
        direction = x - baseline
        alphas = torch.linspace(0, 1, steps, device=device)

        snaps: List[np.ndarray] = []
        f_vals: List[float] = []
        s_vals: List[float] = []
        for a in alphas:
            aa = a.view(1, 1, 1, 1)
            x_a = (baseline + aa * direction).clone().detach().requires_grad_(True)
            out = model(x_a)
            tgt = target_fn(out)
            grad = torch.autograd.grad(
                tgt.sum(), x_a, retain_graph=False, create_graph=False
            )[0]
            step_attr = (direction * grad).squeeze().detach().cpu().numpy()
            snaps.append(np.abs(step_attr))
            f_vals.append(float(tgt.detach().cpu().reshape(-1)[0]))
            s_vals.append(float((grad * direction).sum().detach().cpu().item()))

        alphas_np = alphas.detach().cpu().numpy()
        f_vals_np = np.array(f_vals)
        s_vals_np = np.array(s_vals)
        cum_trapz = np.zeros_like(s_vals_np)
        for i in range(1, len(s_vals_np)):
            da = float(alphas_np[i] - alphas_np[i - 1])
            cum_trapz[i] = (
                cum_trapz[i - 1] + 0.5 * (s_vals_np[i - 1] + s_vals_np[i]) * da
            )
        return {
            "snaps": snaps,
            "alphas": alphas_np,
            "f_vals": f_vals_np,
            "s_vals": s_vals_np,
            "cum_trapz": cum_trapz,
            "delta_f": float(f_vals_np[-1] - f_vals_np[0]),
        }

    story = ig_path_story(model, x_opt, steps=201)
    idxs = np.linspace(0, len(story["alphas"]) - 1, 8, dtype=int)
    all_vals = np.concatenate([h.ravel() for h in story["snaps"]])
    gmax = max(float(np.percentile(all_vals, 99.9)), 1e-12)
    base = x_np

    fig, axs = plt.subplots(2, 4, figsize=(14, 6.5), dpi=170, constrained_layout=True)
    for ax, idx in zip(axs.flatten(), idxs):
        alpha = story["alphas"][idx]
        heat = np.clip(story["snaps"][idx] / gmax, 0, 1)
        ax.imshow(base, cmap="gray_r", origin="upper", vmin=0, vmax=1)
        ax.imshow(heat, cmap="magma", origin="upper", alpha=0.65, vmin=0, vmax=1)
        ax.set_title(f"alpha={alpha:.2f}")
        ax.axis("off")
    plt.suptitle("IG path heat snapshots: |(x - baseline) * grad f(x_alpha)|")
    plt.show()

    integral_final = float(story["cum_trapz"][-1])
    delta_final = float(story["delta_f"])
    relative_gap = abs(integral_final - delta_final) / max(abs(delta_final), 1e-12)
    print(f"IG completeness relative gap: {100 * relative_gap:.2f}%")
else:
    print("Skipping IG path snapshots: trained checkpoint not supplied.")