from __future__ import annotations

import time
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset

from .ae import CircularConv1d
from .ks import params_from_config, power_spectrum
from .ks_data import _spectrum_rel_l1, generate_autoencoder_data
from .utils import best_device, ensure_dir, format_seconds, save_json, set_seed


class KSDataset(Dataset):
    def __init__(self, states: np.ndarray, mean: float, std: float, augment_shift: bool) -> None:
        self.states = ((states.astype(np.float32) - float(mean)) / float(std)).astype(np.float32)
        self.augment_shift = bool(augment_shift)
        self.n = int(states.shape[-1])

    def __len__(self) -> int:
        return int(self.states.shape[0])

    def __getitem__(self, idx: int) -> torch.Tensor:
        x = torch.from_numpy(self.states[idx]).unsqueeze(0)
        if self.augment_shift:
            shift = int(torch.randint(0, self.n, ()).item())
            x = torch.roll(x, shifts=shift, dims=-1)
        return x


class KSConvAutoencoder(nn.Module):
    def __init__(
        self,
        n: int = 1024,
        latent_channels: int = 2,
        kz: int = 64,
        base_channels: int = 32,
    ) -> None:
        super().__init__()
        if n % kz != 0:
            raise ValueError("N must be divisible by Kz.")
        downsample = n // kz
        if downsample <= 0 or downsample & (downsample - 1):
            raise ValueError("N / Kz must be a power of two.")
        self.n = int(n)
        self.latent_channels = int(latent_channels)
        self.kz = int(kz)
        self.base_channels = int(base_channels)
        stages = int(np.log2(downsample))
        b = int(base_channels)
        stage_channels = [min(b * (2 ** min(stage, 2)), 4 * b) for stage in range(stages)]

        encoder: list[nn.Module] = [CircularConv1d(1, b, 7), nn.GELU()]
        in_ch = b
        for out_ch in stage_channels:
            encoder.extend(
                [
                    CircularConv1d(in_ch, out_ch, 5, stride=2),
                    nn.GELU(),
                    CircularConv1d(out_ch, out_ch, 5),
                    nn.GELU(),
                ]
            )
            in_ch = out_ch
        encoder.extend(
            [
                CircularConv1d(in_ch, 2 * b, 3),
                nn.GELU(),
                CircularConv1d(2 * b, self.latent_channels, 3),
            ]
        )
        self.encoder = nn.Sequential(*encoder)

        hidden = stage_channels[-1] if stage_channels else b
        decoder: list[nn.Module] = [CircularConv1d(self.latent_channels, hidden, 3), nn.GELU()]
        in_ch = hidden
        for out_ch in reversed(stage_channels):
            decoder.extend(
                [
                    nn.Upsample(scale_factor=2, mode="nearest"),
                    CircularConv1d(in_ch, out_ch, 5),
                    nn.GELU(),
                    CircularConv1d(out_ch, out_ch, 5),
                    nn.GELU(),
                ]
            )
            in_ch = out_ch
        decoder.extend([CircularConv1d(in_ch, b, 5), nn.GELU(), CircularConv1d(b, 1, 3)])
        self.decoder = nn.Sequential(*decoder)

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        return self.encoder(x)

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        return self.decoder(z)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.decode(self.encode(x))


class KSGlobalAutoencoder(nn.Module):
    def __init__(
        self,
        n: int = 1024,
        latent_dim: int = 256,
        bottleneck_grid: int = 64,
        base_channels: int = 32,
    ) -> None:
        super().__init__()
        if n % bottleneck_grid != 0:
            raise ValueError("N must be divisible by global bottleneck_grid.")
        downsample = n // bottleneck_grid
        if downsample <= 0 or downsample & (downsample - 1):
            raise ValueError("N / bottleneck_grid must be a power of two.")
        self.n = int(n)
        self.latent_dim = int(latent_dim)
        self.bottleneck_grid = int(bottleneck_grid)
        self.base_channels = int(base_channels)
        stages = int(np.log2(downsample))
        b = int(base_channels)
        stage_channels = [min(b * (2 ** min(stage, 2)), 4 * b) for stage in range(stages)]

        encoder: list[nn.Module] = [CircularConv1d(1, b, 7), nn.GELU()]
        in_ch = b
        for out_ch in stage_channels:
            encoder.extend(
                [
                    CircularConv1d(in_ch, out_ch, 5, stride=2),
                    nn.GELU(),
                    CircularConv1d(out_ch, out_ch, 5),
                    nn.GELU(),
                ]
            )
            in_ch = out_ch
        hidden = stage_channels[-1] if stage_channels else b
        encoder.extend([CircularConv1d(in_ch, hidden, 3), nn.GELU()])
        self.encoder_conv = nn.Sequential(*encoder)
        self.encoder_dense = nn.Linear(hidden * self.bottleneck_grid, self.latent_dim)

        self.decoder_dense = nn.Linear(self.latent_dim, hidden * self.bottleneck_grid)
        decoder: list[nn.Module] = [CircularConv1d(hidden, hidden, 3), nn.GELU()]
        in_ch = hidden
        for out_ch in reversed(stage_channels):
            decoder.extend(
                [
                    nn.Upsample(scale_factor=2, mode="nearest"),
                    CircularConv1d(in_ch, out_ch, 5),
                    nn.GELU(),
                    CircularConv1d(out_ch, out_ch, 5),
                    nn.GELU(),
                ]
            )
            in_ch = out_ch
        decoder.extend([CircularConv1d(in_ch, b, 5), nn.GELU(), CircularConv1d(b, 1, 3)])
        self.decoder_conv = nn.Sequential(*decoder)

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        h = self.encoder_conv(x)
        return self.encoder_dense(h.flatten(start_dim=1))

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        h = self.decoder_dense(z)
        h = h.reshape(z.shape[0], -1, self.bottleneck_grid)
        return self.decoder_conv(h)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.decode(self.encode(x))


class KSGlobalPCAAutoencoder(nn.Module):
    def __init__(self, n: int = 1024, latent_dim: int = 256) -> None:
        super().__init__()
        self.n = int(n)
        self.latent_dim = int(latent_dim)
        self.bottleneck_grid = 0
        self.base_channels = 0
        self.register_buffer("pca_mean", torch.zeros(self.n, dtype=torch.float32))
        self.register_buffer("components", torch.zeros(self.latent_dim, self.n, dtype=torch.float32))

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        flat = x.squeeze(1)
        return (flat - self.pca_mean) @ self.components.T

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        flat = z @ self.components + self.pca_mean
        return flat.unsqueeze(1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.decode(self.encode(x))


def checkpoint_path(config: dict[str, Any]) -> Path:
    model = config["model"]
    ae = config["ae"]
    length = float(model["length"])
    length_label = int(length) if float(length).is_integer() else str(length).replace(".", "p")
    return Path(config["paths"]["model_dir"]) / (
        f"ks_ae_N{int(model['N'])}_L{length_label}_C{int(ae['latent_channels'])}_Kz{int(ae['Kz'])}.pt"
    )


def global_checkpoint_path(config: dict[str, Any]) -> Path:
    model = config["model"]
    ae = config["global_ae"]
    length = float(model["length"])
    length_label = int(length) if float(length).is_integer() else str(length).replace(".", "p")
    return Path(config["paths"]["model_dir"]) / (
        f"ks_global_ae_N{int(model['N'])}_L{length_label}_Z{int(ae['latent_dim'])}.pt"
    )


def _metrics_path(config: dict[str, Any]) -> Path:
    return Path(config["paths"]["model_dir"]) / "ks_ae_metrics.json"


def _global_metrics_path(config: dict[str, Any]) -> Path:
    return Path(config["paths"]["model_dir"]) / "ks_global_ae_metrics.json"


def _diagnostics_path(config: dict[str, Any]) -> Path:
    return Path(config["paths"]["model_dir"]) / "ks_ae_diagnostics.npz"


def _global_diagnostics_path(config: dict[str, Any]) -> Path:
    return Path(config["paths"]["model_dir"]) / "ks_global_ae_diagnostics.npz"


def _spectra_path(config: dict[str, Any]) -> Path:
    return Path(config["paths"]["model_dir"]) / "ks_ae_spectra_summary.npz"


def _global_spectra_path(config: dict[str, Any]) -> Path:
    return Path(config["paths"]["model_dir"]) / "ks_global_ae_spectra_summary.npz"


def load_autoencoder(path: str | Path, device: str = "cpu") -> tuple[KSConvAutoencoder, float, float]:
    payload = torch.load(Path(path), map_location=device, weights_only=False)
    model = KSConvAutoencoder(
        n=int(payload["n"]),
        latent_channels=int(payload["latent_channels"]),
        kz=int(payload["kz"]),
        base_channels=int(payload["base_channels"]),
    )
    model.load_state_dict(payload["state_dict"])
    model.to(device)
    model.eval()
    return model, float(payload["mean"]), float(payload["std"])


def load_global_autoencoder(path: str | Path, device: str = "cpu") -> tuple[nn.Module, float, float]:
    payload = torch.load(Path(path), map_location=device, weights_only=False)
    architecture = str(payload.get("architecture", "conv_dense"))
    if architecture == "pca_global":
        model: nn.Module = KSGlobalPCAAutoencoder(
            n=int(payload["n"]),
            latent_dim=int(payload["latent_dim"]),
        )
    else:
        model = KSGlobalAutoencoder(
            n=int(payload["n"]),
            latent_dim=int(payload["latent_dim"]),
            bottleneck_grid=int(payload["bottleneck_grid"]),
            base_channels=int(payload["base_channels"]),
        )
    model.load_state_dict(payload["state_dict"])
    model.to(device)
    model.eval()
    return model, float(payload["mean"]), float(payload["std"])


def encode_physical(
    model: KSConvAutoencoder,
    states: np.ndarray,
    mean: float,
    std: float,
    device: str,
) -> np.ndarray:
    x = torch.as_tensor(((states.astype(np.float32) - mean) / std), device=device).unsqueeze(1)
    model.eval()
    with torch.no_grad():
        z = model.encode(x).detach().cpu().numpy()
    return z.reshape(states.shape[0], -1).astype(np.float64)


def encode_global_physical(
    model: nn.Module,
    states: np.ndarray,
    mean: float,
    std: float,
    device: str,
) -> np.ndarray:
    x = torch.as_tensor(((states.astype(np.float32) - mean) / std), device=device).unsqueeze(1)
    model.eval()
    with torch.no_grad():
        z = model.encode(x).detach().cpu().numpy()
    return z.astype(np.float64)


def decode_latent(
    model: KSConvAutoencoder,
    latent_flat: np.ndarray,
    mean: float,
    std: float,
    device: str,
) -> np.ndarray:
    n_members = int(latent_flat.shape[0])
    z = torch.as_tensor(
        latent_flat.reshape(n_members, model.latent_channels, model.kz).astype(np.float32),
        device=device,
    )
    model.eval()
    with torch.no_grad():
        decoded = model.decode(z).squeeze(1).detach().cpu().numpy()
    return (decoded * std + mean).astype(np.float64)


def decode_global_latent(
    model: nn.Module,
    latent: np.ndarray,
    mean: float,
    std: float,
    device: str,
) -> np.ndarray:
    z = torch.as_tensor(latent.astype(np.float32), device=device)
    model.eval()
    with torch.no_grad():
        decoded = model.decode(z).squeeze(1).detach().cpu().numpy()
    return (decoded * std + mean).astype(np.float64)


def _spectral_loss(recon: torch.Tensor, truth: torch.Tensor, max_mode: int) -> torch.Tensor:
    recon_hat = torch.fft.rfft(recon.squeeze(1), dim=-1)
    truth_hat = torch.fft.rfft(truth.squeeze(1), dim=-1)
    max_idx = min(int(max_mode), recon_hat.shape[-1] - 1)
    if max_idx <= 1:
        return torch.zeros((), dtype=recon.dtype, device=recon.device)
    recon_log = torch.log1p(torch.abs(recon_hat[..., 1 : max_idx + 1]))
    truth_log = torch.log1p(torch.abs(truth_hat[..., 1 : max_idx + 1]))
    return nn.functional.mse_loss(recon_log, truth_log)


def _loss_terms(model: nn.Module, x: torch.Tensor, ae_cfg: dict[str, Any]) -> dict[str, torch.Tensor]:
    recon = model(x)
    mse = nn.functional.mse_loss(recon, x)
    derivative = nn.functional.mse_loss(recon - torch.roll(recon, shifts=1, dims=-1), x - torch.roll(x, shifts=1, dims=-1))
    spectrum = _spectral_loss(recon, x, int(ae_cfg.get("spectrum_loss_max_mode", x.shape[-1] // 3)))
    total = (
        mse
        + float(ae_cfg.get("derivative_loss_weight", 0.0)) * derivative
        + float(ae_cfg.get("spectrum_loss_weight", 0.0)) * spectrum
    )
    return {"total": total, "mse": mse, "derivative": derivative, "spectrum": spectrum}


def _evaluate_losses(model: nn.Module, loader: DataLoader, device: str, ae_cfg: dict[str, Any]) -> dict[str, float]:
    model.eval()
    totals = {"total": 0.0, "mse": 0.0, "derivative": 0.0, "spectrum": 0.0}
    count = 0
    with torch.no_grad():
        for x in loader:
            x = x.to(device)
            terms = _loss_terms(model, x, ae_cfg)
            batch = int(x.shape[0])
            for key in totals:
                totals[key] += float(terms[key].item()) * batch
            count += batch
    return {key: value / max(count, 1) for key, value in totals.items()}


def _save_checkpoint(
    model: KSConvAutoencoder,
    config: dict[str, Any],
    mean: float,
    std: float,
    best_val_loss: float,
    phase: str,
    epoch: int,
) -> None:
    ae = config["ae"]
    torch.save(
        {
            "state_dict": model.state_dict(),
            "n": int(config["model"]["N"]),
            "length": float(config["model"]["length"]),
            "latent_channels": int(ae["latent_channels"]),
            "kz": int(ae["Kz"]),
            "base_channels": int(ae["base_channels"]),
            "mean": float(mean),
            "std": float(std),
            "best_val_loss": float(best_val_loss),
            "best_epoch": int(epoch),
            "train_phase": phase,
        },
        checkpoint_path(config),
    )


def _save_global_checkpoint(
    model: nn.Module,
    config: dict[str, Any],
    mean: float,
    std: float,
    best_val_loss: float,
    phase: str,
    epoch: int,
) -> None:
    ae = config["global_ae"]
    torch.save(
        {
            "state_dict": model.state_dict(),
            "n": int(config["model"]["N"]),
            "length": float(config["model"]["length"]),
            "latent_dim": int(ae["latent_dim"]),
            "architecture": str(ae.get("architecture", "conv_dense")),
            "bottleneck_grid": int(ae.get("bottleneck_grid", 0)),
            "base_channels": int(ae.get("base_channels", 0)),
            "mean": float(mean),
            "std": float(std),
            "best_val_loss": float(best_val_loss),
            "best_epoch": int(epoch),
            "train_phase": phase,
        },
        global_checkpoint_path(config),
    )


def _train_phase(
    model: KSConvAutoencoder,
    train_loader: DataLoader,
    val_loader: DataLoader,
    device: str,
    config: dict[str, Any],
    mean: float,
    std: float,
    phase: str,
    learning_rate: float,
    epochs: int,
    patience: int,
    initial_best: float,
) -> dict[str, Any]:
    ae_cfg = config["ae"]
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=float(learning_rate),
        weight_decay=float(ae_cfg["weight_decay"]),
    )
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode="min",
        factor=float(ae_cfg.get("scheduler_factor", 0.5)),
        patience=int(ae_cfg.get("scheduler_patience", 5)),
        min_lr=float(ae_cfg.get("min_learning_rate", 1.0e-5)),
    )
    best_val = float(initial_best)
    best_epoch = 0
    epochs_since_improve = 0
    history: list[dict[str, Any]] = []
    start = time.time()
    grad_clip = float(ae_cfg.get("gradient_clip_norm", 0.0))
    for epoch in range(1, int(epochs) + 1):
        model.train()
        train_totals = {"total": 0.0, "mse": 0.0, "derivative": 0.0, "spectrum": 0.0}
        train_count = 0
        max_grad_norm = 0.0
        for x in train_loader:
            x = x.to(device)
            optimizer.zero_grad(set_to_none=True)
            terms = _loss_terms(model, x, ae_cfg)
            if not torch.isfinite(terms["total"]):
                raise FloatingPointError(f"Non-finite KS AE loss in phase {phase}, epoch {epoch}")
            terms["total"].backward()
            if grad_clip > 0.0:
                grad_norm = float(torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip).item())
                max_grad_norm = max(max_grad_norm, grad_norm)
            optimizer.step()
            batch = int(x.shape[0])
            for key in train_totals:
                train_totals[key] += float(terms[key].item()) * batch
            train_count += batch
        train_metrics = {key: value / max(train_count, 1) for key, value in train_totals.items()}
        val_metrics = _evaluate_losses(model, val_loader, device, ae_cfg)
        scheduler.step(val_metrics["total"])
        row = {
            "phase": phase,
            "epoch": int(epoch),
            "learning_rate": float(optimizer.param_groups[0]["lr"]),
            "train_loss": float(train_metrics["total"]),
            "train_mse": float(train_metrics["mse"]),
            "train_derivative_loss": float(train_metrics["derivative"]),
            "train_spectrum_loss": float(train_metrics["spectrum"]),
            "val_loss": float(val_metrics["total"]),
            "val_mse": float(val_metrics["mse"]),
            "val_normalized_rmse": float(np.sqrt(val_metrics["mse"])),
            "val_derivative_loss": float(val_metrics["derivative"]),
            "val_spectrum_loss": float(val_metrics["spectrum"]),
            "max_grad_norm": float(max_grad_norm),
        }
        history.append(row)
        if val_metrics["total"] < best_val:
            best_val = float(val_metrics["total"])
            best_epoch = int(epoch)
            epochs_since_improve = 0
            _save_checkpoint(model, config, mean, std, best_val, phase, epoch)
        else:
            epochs_since_improve += 1
        if epochs_since_improve >= int(patience):
            break
    return {
        "phase": phase,
        "best_val_loss": float(best_val),
        "best_epoch": int(best_epoch),
        "epochs_run": len(history),
        "elapsed": format_seconds(time.time() - start),
        "history": history,
    }


def _train_global_phase(
    model: KSGlobalAutoencoder,
    train_loader: DataLoader,
    val_loader: DataLoader,
    device: str,
    config: dict[str, Any],
    mean: float,
    std: float,
    phase: str,
    learning_rate: float,
    epochs: int,
    patience: int,
    initial_best: float,
) -> dict[str, Any]:
    ae_cfg = config["global_ae"]
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=float(learning_rate),
        weight_decay=float(ae_cfg["weight_decay"]),
    )
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode="min",
        factor=float(ae_cfg.get("scheduler_factor", 0.5)),
        patience=int(ae_cfg.get("scheduler_patience", 5)),
        min_lr=float(ae_cfg.get("min_learning_rate", 1.0e-5)),
    )
    best_val = float(initial_best)
    best_epoch = 0
    epochs_since_improve = 0
    history: list[dict[str, Any]] = []
    start = time.time()
    grad_clip = float(ae_cfg.get("gradient_clip_norm", 0.0))
    for epoch in range(1, int(epochs) + 1):
        model.train()
        train_totals = {"total": 0.0, "mse": 0.0, "derivative": 0.0, "spectrum": 0.0}
        train_count = 0
        max_grad_norm = 0.0
        for x in train_loader:
            x = x.to(device)
            optimizer.zero_grad(set_to_none=True)
            terms = _loss_terms(model, x, ae_cfg)
            if not torch.isfinite(terms["total"]):
                raise FloatingPointError(f"Non-finite KS global AE loss in phase {phase}, epoch {epoch}")
            terms["total"].backward()
            if grad_clip > 0.0:
                grad_norm = float(torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip).item())
                max_grad_norm = max(max_grad_norm, grad_norm)
            optimizer.step()
            batch = int(x.shape[0])
            for key in train_totals:
                train_totals[key] += float(terms[key].item()) * batch
            train_count += batch
        train_metrics = {key: value / max(train_count, 1) for key, value in train_totals.items()}
        val_metrics = _evaluate_losses(model, val_loader, device, ae_cfg)
        scheduler.step(val_metrics["total"])
        row = {
            "phase": phase,
            "epoch": int(epoch),
            "learning_rate": float(optimizer.param_groups[0]["lr"]),
            "train_loss": float(train_metrics["total"]),
            "train_mse": float(train_metrics["mse"]),
            "train_derivative_loss": float(train_metrics["derivative"]),
            "train_spectrum_loss": float(train_metrics["spectrum"]),
            "val_loss": float(val_metrics["total"]),
            "val_mse": float(val_metrics["mse"]),
            "val_normalized_rmse": float(np.sqrt(val_metrics["mse"])),
            "val_derivative_loss": float(val_metrics["derivative"]),
            "val_spectrum_loss": float(val_metrics["spectrum"]),
            "max_grad_norm": float(max_grad_norm),
        }
        history.append(row)
        if val_metrics["total"] < best_val:
            best_val = float(val_metrics["total"])
            best_epoch = int(epoch)
            epochs_since_improve = 0
            _save_global_checkpoint(model, config, mean, std, best_val, phase, epoch)
        else:
            epochs_since_improve += 1
        if epochs_since_improve >= int(patience):
            break
    return {
        "phase": phase,
        "best_val_loss": float(best_val),
        "best_epoch": int(best_epoch),
        "epochs_run": len(history),
        "elapsed": format_seconds(time.time() - start),
        "history": history,
    }


def _reconstruct(model: nn.Module, states: np.ndarray, mean: float, std: float, device: str, batch_size: int) -> np.ndarray:
    model.eval()
    outputs = []
    with torch.no_grad():
        for start in range(0, states.shape[0], int(batch_size)):
            batch = states[start : start + int(batch_size)].astype(np.float32)
            x = torch.as_tensor((batch - mean) / std, device=device).unsqueeze(1)
            recon = model(x).squeeze(1).detach().cpu().numpy()
            outputs.append(recon * std + mean)
    return np.concatenate(outputs, axis=0).astype(np.float32)


def _metrics_for_split(
    config: dict[str, Any],
    split_name: str,
    truth: np.ndarray,
    recon: np.ndarray,
    normalization_std: float,
) -> dict[str, float]:
    params = params_from_config(config)
    gates = config["gates"]
    truth_spec = power_spectrum(truth, params)["power_spectrum"]
    recon_spec = power_spectrum(recon, params)["power_spectrum"]
    low_k_max = int(gates["low_k_max"])
    active_k_max = min(int(gates["active_k_max"]), params.n // 3)
    normalized = (recon - truth) / float(normalization_std)
    return {
        f"{split_name}_normalized_reconstruction_rmse": float(np.sqrt(np.mean(normalized**2))),
        f"{split_name}_physical_rmse": float(np.sqrt(np.mean((recon - truth) ** 2))),
        f"{split_name}_mean_abs_error": float(abs(np.mean(recon) - np.mean(truth))),
        f"{split_name}_variance_rel_error": _relative_scalar_error(float(np.var(recon)), float(np.var(truth))),
        f"{split_name}_low_k_spectrum_rel_l1": _spectrum_rel_l1(recon_spec, truth_spec, 1, low_k_max),
        f"{split_name}_active_band_spectrum_rel_l1": _spectrum_rel_l1(recon_spec, truth_spec, 1, active_k_max),
    }


def _relative_scalar_error(value: float, reference: float) -> float:
    return float(abs(value - reference) / max(abs(reference), 1.0e-12))


def _quality_gates(config: dict[str, Any], diagnostics: dict[str, float]) -> dict[str, Any]:
    gates = config["gates"]
    split = str(gates.get("selection_split", "val"))
    if split not in {"val", "holdout"}:
        raise ValueError("gates.selection_split must be 'val' or 'holdout'.")
    nrmse = float(diagnostics[f"{split}_normalized_reconstruction_rmse"])
    low = float(diagnostics[f"{split}_low_k_spectrum_rel_l1"])
    active = float(diagnostics[f"{split}_active_band_spectrum_rel_l1"])
    finite = bool(diagnostics["val_outputs_finite"] and diagnostics["holdout_outputs_finite"])
    primary = bool(
        finite
        and nrmse <= float(gates["primary_normalized_rmse_max"])
        and low <= float(gates["primary_low_k_spectrum_rel_error_max"])
        and active <= float(gates["primary_active_band_spectrum_rel_error_max"])
    )
    backup = bool(
        finite
        and nrmse <= float(gates["backup_normalized_rmse_max"])
        and low <= float(gates["backup_low_k_spectrum_rel_error_max"])
        and active <= float(gates["backup_active_band_spectrum_rel_error_max"])
    )
    return {
        "primary_pass": primary,
        "backup_pass": backup,
        "selection_split": split,
        "thresholds": gates,
    }


def _write_diagnostics(
    config: dict[str, Any],
    model: KSConvAutoencoder,
    data: dict[str, np.ndarray],
    mean: float,
    std: float,
    device: str,
) -> dict[str, Any]:
    ae_cfg = config["ae"]
    batch_size = int(ae_cfg["batch_size"])
    val = data["val"]
    holdout = data["holdout"]
    val_recon = _reconstruct(model, val, mean, std, device, batch_size)
    holdout_recon = _reconstruct(model, holdout, mean, std, device, batch_size)
    diagnostics: dict[str, Any] = {}
    diagnostics.update(_metrics_for_split(config, "val", val, val_recon, std))
    diagnostics.update(_metrics_for_split(config, "holdout", holdout, holdout_recon, std))
    diagnostics["val_outputs_finite"] = bool(np.all(np.isfinite(val_recon)))
    diagnostics["holdout_outputs_finite"] = bool(np.all(np.isfinite(holdout_recon)))

    params = params_from_config(config)
    val_truth_spec = power_spectrum(val, params)
    val_recon_spec = power_spectrum(val_recon, params)
    holdout_truth_spec = power_spectrum(holdout, params)
    holdout_recon_spec = power_spectrum(holdout_recon, params)

    sample_count = min(int(ae_cfg["diagnostic_samples"]), int(val.shape[0]))
    sample = val[:sample_count].astype(np.float32)
    sample_recon = val_recon[:sample_count].astype(np.float32)
    x_norm = torch.as_tensor((sample - mean) / std, device=device).unsqueeze(1)
    model.eval()
    responses = []
    with torch.no_grad():
        z0 = model.encode(x_norm[:1])
        base = model.decode(z0).squeeze(1)
        center = model.kz // 2
        for channel in range(model.latent_channels):
            zp = z0.clone()
            zp[:, channel, center] += float(ae_cfg["diagnostic_perturbation"])
            diff = (model.decode(zp).squeeze(1) - base).detach().cpu().numpy()[0] * std
            responses.append(diff)

    np.savez_compressed(
        _diagnostics_path(config),
        truth=sample,
        reconstruction=sample_recon,
        latent_cell_responses=np.asarray(responses, dtype=np.float32),
        latent_cell_index=np.asarray([center], dtype=np.int64),
        latent_cell_physical_coord=np.asarray([center * model.n / model.kz], dtype=np.float32),
    )
    np.savez_compressed(
        _spectra_path(config),
        wavenumber=val_truth_spec["wavenumber"],
        val_truth_power_spectrum=val_truth_spec["power_spectrum"],
        val_reconstruction_power_spectrum=val_recon_spec["power_spectrum"],
        holdout_truth_power_spectrum=holdout_truth_spec["power_spectrum"],
        holdout_reconstruction_power_spectrum=holdout_recon_spec["power_spectrum"],
    )
    diagnostics["quality_gates"] = _quality_gates(config, diagnostics)
    return diagnostics


def _write_global_diagnostics(
    config: dict[str, Any],
    model: KSGlobalAutoencoder,
    data: dict[str, np.ndarray],
    mean: float,
    std: float,
    device: str,
) -> dict[str, Any]:
    ae_cfg = config["global_ae"]
    batch_size = int(ae_cfg["batch_size"])
    val = data["val"]
    holdout = data["holdout"]
    val_recon = _reconstruct(model, val, mean, std, device, batch_size)
    holdout_recon = _reconstruct(model, holdout, mean, std, device, batch_size)
    diagnostics: dict[str, Any] = {}
    diagnostics.update(_metrics_for_split(config, "val", val, val_recon, std))
    diagnostics.update(_metrics_for_split(config, "holdout", holdout, holdout_recon, std))
    diagnostics["val_outputs_finite"] = bool(np.all(np.isfinite(val_recon)))
    diagnostics["holdout_outputs_finite"] = bool(np.all(np.isfinite(holdout_recon)))

    params = params_from_config(config)
    val_truth_spec = power_spectrum(val, params)
    val_recon_spec = power_spectrum(val_recon, params)
    holdout_truth_spec = power_spectrum(holdout, params)
    holdout_recon_spec = power_spectrum(holdout_recon, params)

    sample_count = min(int(ae_cfg["diagnostic_samples"]), int(val.shape[0]))
    sample = val[:sample_count].astype(np.float32)
    sample_recon = val_recon[:sample_count].astype(np.float32)
    x_norm = torch.as_tensor((sample - mean) / std, device=device).unsqueeze(1)
    model.eval()
    responses = []
    with torch.no_grad():
        z0 = model.encode(x_norm[:1])
        base = model.decode(z0).squeeze(1)
        indices = np.asarray(
            [
                0,
                model.latent_dim // 4,
                model.latent_dim // 2,
                (3 * model.latent_dim) // 4,
            ],
            dtype=np.int64,
        )
        for index in indices:
            zp = z0.clone()
            zp[:, int(index)] += float(ae_cfg["diagnostic_perturbation"])
            diff = (model.decode(zp).squeeze(1) - base).detach().cpu().numpy()[0] * std
            responses.append(diff)

    np.savez_compressed(
        _global_diagnostics_path(config),
        truth=sample,
        reconstruction=sample_recon,
        latent_unit_responses=np.asarray(responses, dtype=np.float32),
        latent_unit_index=indices,
    )
    np.savez_compressed(
        _global_spectra_path(config),
        wavenumber=val_truth_spec["wavenumber"],
        val_truth_power_spectrum=val_truth_spec["power_spectrum"],
        val_reconstruction_power_spectrum=val_recon_spec["power_spectrum"],
        holdout_truth_power_spectrum=holdout_truth_spec["power_spectrum"],
        holdout_reconstruction_power_spectrum=holdout_recon_spec["power_spectrum"],
    )
    diagnostics["quality_gates"] = _quality_gates(config, diagnostics)
    return diagnostics


def _parameter_count(model: nn.Module) -> int:
    params = int(sum(param.numel() for param in model.parameters()))
    if params > 0:
        return params
    return int(sum(buffer.numel() for buffer in model.buffers()))


def train_autoencoder(config: dict[str, Any], overwrite_data: bool = False, overwrite_ae: bool = False) -> Path:
    ensure_dir(config["paths"]["model_dir"])
    ensure_dir(config["paths"]["figures_dir"])
    out_path = checkpoint_path(config)
    metrics_path = _metrics_path(config)
    if out_path.exists() and metrics_path.exists() and not overwrite_ae:
        return out_path

    generate_autoencoder_data(config, overwrite=overwrite_data)
    data_path = Path(config["paths"]["data_dir"]) / "ks_ae_data.npz"
    payload = np.load(data_path)
    data = {
        "train": payload["train"],
        "val": payload["val"],
        "holdout": payload["holdout"],
    }
    mean = float(payload["mean"])
    std = float(payload["std"])
    ae_cfg = config["ae"]
    set_seed(int(config["seeds"].get("ae_train", 44001)))
    device = best_device(str(ae_cfg.get("device", "auto")))
    model = KSConvAutoencoder(
        n=int(config["model"]["N"]),
        latent_channels=int(ae_cfg["latent_channels"]),
        kz=int(ae_cfg["Kz"]),
        base_channels=int(ae_cfg["base_channels"]),
    ).to(device)
    train_loader = DataLoader(
        KSDataset(data["train"], mean, std, augment_shift=True),
        batch_size=int(ae_cfg["batch_size"]),
        shuffle=True,
        num_workers=0,
    )
    val_loader = DataLoader(
        KSDataset(data["val"], mean, std, augment_shift=False),
        batch_size=int(ae_cfg["batch_size"]),
        shuffle=False,
        num_workers=0,
    )

    overall_start = time.time()
    phases = [
        _train_phase(
            model,
            train_loader,
            val_loader,
            device,
            config,
            mean,
            std,
            "main",
            float(ae_cfg["learning_rate"]),
            int(ae_cfg["epochs"]),
            int(ae_cfg["patience"]),
            float("inf"),
        )
    ]
    model, mean, std = load_autoencoder(out_path, device=device)
    diagnostics = _write_diagnostics(config, model, data, mean, std, device)
    if (
        bool(ae_cfg.get("finetune_if_backup_passes", True))
        and not bool(diagnostics["quality_gates"]["primary_pass"])
        and bool(diagnostics["quality_gates"]["backup_pass"])
    ):
        phases.append(
            _train_phase(
                model,
                train_loader,
                val_loader,
                device,
                config,
                mean,
                std,
                "finetune",
                float(ae_cfg["finetune_learning_rate"]),
                int(ae_cfg["finetune_epochs"]),
                int(ae_cfg["finetune_patience"]),
                float(phases[-1]["best_val_loss"]),
            )
        )
        model, mean, std = load_autoencoder(out_path, device=device)
        diagnostics = _write_diagnostics(config, model, data, mean, std, device)

    history = [row for phase in phases for row in phase["history"]]
    metrics = {
        "checkpoint": out_path,
        "data_path": data_path,
        "device": device,
        "architecture": str(ae_cfg.get("architecture", "conv_dense")),
        "elapsed": format_seconds(time.time() - overall_start),
        "phases": phases,
        "history": history,
        "diagnostics": {k: v for k, v in diagnostics.items() if k != "quality_gates"},
        "quality_gates": diagnostics["quality_gates"],
        "mean": mean,
        "std": std,
        "latent_shape": [int(ae_cfg["latent_channels"]), int(ae_cfg["Kz"])],
        "latent_dimension": int(ae_cfg["latent_channels"]) * int(ae_cfg["Kz"]),
        "physical_dimension": int(config["model"]["N"]),
        "parameter_count": _parameter_count(model),
        "compression_ratio": float(int(config["model"]["N"]) / (int(ae_cfg["latent_channels"]) * int(ae_cfg["Kz"]))),
    }
    save_json(metrics, metrics_path)
    save_json(
        {
            "checkpoint": out_path,
            "quality_gates": diagnostics["quality_gates"],
            "diagnostics": metrics["diagnostics"],
        },
        Path(config["paths"]["figures_dir"]) / "ks_ae_quality_summary.json",
    )
    return out_path


def train_global_autoencoder(config: dict[str, Any], overwrite_data: bool = False, overwrite_ae: bool = False) -> Path:
    ensure_dir(config["paths"]["model_dir"])
    ensure_dir(config["paths"]["figures_dir"])
    out_path = global_checkpoint_path(config)
    metrics_path = _global_metrics_path(config)
    if out_path.exists() and metrics_path.exists() and not overwrite_ae:
        return out_path

    if "global_ae" not in config:
        raise KeyError("Config must define global_ae to train the global KS autoencoder.")
    ae_cfg = config["global_ae"]
    if str(ae_cfg.get("architecture", "conv_dense")) == "pca_global":
        return train_global_pca_autoencoder(config, overwrite_data=overwrite_data, overwrite_ae=overwrite_ae)
    generate_autoencoder_data(config, overwrite=overwrite_data)
    data_path = Path(config["paths"]["data_dir"]) / "ks_ae_data.npz"
    payload = np.load(data_path)
    data = {
        "train": payload["train"],
        "val": payload["val"],
        "holdout": payload["holdout"],
    }
    mean = float(payload["mean"])
    std = float(payload["std"])
    set_seed(int(config["seeds"].get("global_ae_train", config["seeds"].get("ae_train", 44001) + 17)))
    device = best_device(str(ae_cfg.get("device", config.get("ae", {}).get("device", "auto"))))
    model = KSGlobalAutoencoder(
        n=int(config["model"]["N"]),
        latent_dim=int(ae_cfg["latent_dim"]),
        bottleneck_grid=int(ae_cfg["bottleneck_grid"]),
        base_channels=int(ae_cfg["base_channels"]),
    ).to(device)
    train_loader = DataLoader(
        KSDataset(data["train"], mean, std, augment_shift=True),
        batch_size=int(ae_cfg["batch_size"]),
        shuffle=True,
        num_workers=0,
    )
    val_loader = DataLoader(
        KSDataset(data["val"], mean, std, augment_shift=False),
        batch_size=int(ae_cfg["batch_size"]),
        shuffle=False,
        num_workers=0,
    )

    overall_start = time.time()
    phases = [
        _train_global_phase(
            model,
            train_loader,
            val_loader,
            device,
            config,
            mean,
            std,
            "main",
            float(ae_cfg["learning_rate"]),
            int(ae_cfg["epochs"]),
            int(ae_cfg["patience"]),
            float("inf"),
        )
    ]
    model, mean, std = load_global_autoencoder(out_path, device=device)
    diagnostics = _write_global_diagnostics(config, model, data, mean, std, device)
    if (
        bool(ae_cfg.get("finetune_if_backup_passes", True))
        and not bool(diagnostics["quality_gates"]["primary_pass"])
        and bool(diagnostics["quality_gates"]["backup_pass"])
    ):
        phases.append(
            _train_global_phase(
                model,
                train_loader,
                val_loader,
                device,
                config,
                mean,
                std,
                "finetune",
                float(ae_cfg["finetune_learning_rate"]),
                int(ae_cfg["finetune_epochs"]),
                int(ae_cfg["finetune_patience"]),
                float(phases[-1]["best_val_loss"]),
            )
        )
        model, mean, std = load_global_autoencoder(out_path, device=device)
        diagnostics = _write_global_diagnostics(config, model, data, mean, std, device)

    history = [row for phase in phases for row in phase["history"]]
    metrics = {
        "checkpoint": out_path,
        "data_path": data_path,
        "device": device,
        "architecture": str(ae_cfg.get("architecture", "conv_dense")),
        "elapsed": format_seconds(time.time() - overall_start),
        "phases": phases,
        "history": history,
        "diagnostics": {k: v for k, v in diagnostics.items() if k != "quality_gates"},
        "quality_gates": diagnostics["quality_gates"],
        "mean": mean,
        "std": std,
        "latent_shape": [int(ae_cfg["latent_dim"])],
        "latent_dimension": int(ae_cfg["latent_dim"]),
        "bottleneck_grid": int(ae_cfg["bottleneck_grid"]),
        "base_channels": int(ae_cfg["base_channels"]),
        "physical_dimension": int(config["model"]["N"]),
        "parameter_count": _parameter_count(model),
        "compression_ratio": float(int(config["model"]["N"]) / int(ae_cfg["latent_dim"])),
    }
    save_json(metrics, metrics_path)
    save_json(
        {
            "checkpoint": out_path,
            "quality_gates": diagnostics["quality_gates"],
            "diagnostics": metrics["diagnostics"],
        },
        Path(config["paths"]["figures_dir"]) / "ks_global_ae_quality_summary.json",
    )
    return out_path


def train_global_pca_autoencoder(
    config: dict[str, Any],
    overwrite_data: bool = False,
    overwrite_ae: bool = False,
) -> Path:
    ensure_dir(config["paths"]["model_dir"])
    ensure_dir(config["paths"]["figures_dir"])
    out_path = global_checkpoint_path(config)
    metrics_path = _global_metrics_path(config)
    if out_path.exists() and metrics_path.exists() and not overwrite_ae:
        return out_path

    generate_autoencoder_data(config, overwrite=overwrite_data)
    data_path = Path(config["paths"]["data_dir"]) / "ks_ae_data.npz"
    payload = np.load(data_path)
    data = {
        "train": payload["train"],
        "val": payload["val"],
        "holdout": payload["holdout"],
    }
    mean = float(payload["mean"])
    std = float(payload["std"])
    ae_cfg = config["global_ae"]
    latent_dim = int(ae_cfg["latent_dim"])
    device = best_device(str(ae_cfg.get("device", config.get("ae", {}).get("device", "auto"))))

    start = time.time()
    train_norm = ((data["train"].astype(np.float64) - mean) / std).astype(np.float64)
    pca_mean = np.mean(train_norm, axis=0)
    centered = train_norm - pca_mean[None, :]
    covariance = (centered.T @ centered) / max(int(centered.shape[0]) - 1, 1)
    eigvals, eigvecs = np.linalg.eigh(covariance)
    order = np.argsort(eigvals)[::-1]
    selected = order[:latent_dim]
    total_variance = float(np.sum(np.maximum(eigvals, 0.0)))
    explained = float(np.sum(np.maximum(eigvals[selected], 0.0)) / max(total_variance, 1.0e-20))

    model = KSGlobalPCAAutoencoder(n=int(config["model"]["N"]), latent_dim=latent_dim).to(device)
    with torch.no_grad():
        model.pca_mean.copy_(torch.as_tensor(pca_mean.astype(np.float32), device=device))
        model.components.copy_(torch.as_tensor(eigvecs[:, selected].T.astype(np.float32), device=device))

    _save_global_checkpoint(model, config, mean, std, 0.0, "pca", 1)
    model, mean, std = load_global_autoencoder(out_path, device=device)
    diagnostics = _write_global_diagnostics(config, model, data, mean, std, device)
    elapsed = format_seconds(time.time() - start)
    metrics = {
        "checkpoint": out_path,
        "data_path": data_path,
        "device": device,
        "elapsed": elapsed,
        "architecture": "pca_global",
        "explained_variance_ratio": explained,
        "phases": [
            {
                "phase": "pca",
                "best_val_loss": 0.0,
                "best_epoch": 1,
                "epochs_run": 1,
                "elapsed": elapsed,
                "explained_variance_ratio": explained,
                "history": [],
            }
        ],
        "history": [],
        "diagnostics": {k: v for k, v in diagnostics.items() if k != "quality_gates"},
        "quality_gates": diagnostics["quality_gates"],
        "mean": mean,
        "std": std,
        "latent_shape": [latent_dim],
        "latent_dimension": latent_dim,
        "physical_dimension": int(config["model"]["N"]),
        "parameter_count": _parameter_count(model),
        "compression_ratio": float(int(config["model"]["N"]) / latent_dim),
    }
    save_json(metrics, metrics_path)
    save_json(
        {
            "checkpoint": out_path,
            "architecture": "pca_global",
            "explained_variance_ratio": explained,
            "quality_gates": diagnostics["quality_gates"],
            "diagnostics": metrics["diagnostics"],
        },
        Path(config["paths"]["figures_dir"]) / "ks_global_ae_quality_summary.json",
    )
    return out_path


def compare_autoencoder_quality(config: dict[str, Any]) -> Path:
    spatial_path = _metrics_path(config)
    global_path = _global_metrics_path(config)
    if not spatial_path.exists():
        raise FileNotFoundError(f"Spatial AE metrics are missing: {spatial_path}")
    if not global_path.exists():
        raise FileNotFoundError(f"Global AE metrics are missing: {global_path}")
    import json

    with spatial_path.open("r", encoding="utf-8") as f:
        spatial = json.load(f)
    with global_path.open("r", encoding="utf-8") as f:
        global_ = json.load(f)
    gates = config.get("gates", {})
    split = str(gates.get("selection_split", "val"))
    key = f"{split}_normalized_reconstruction_rmse"
    spatial_nrmse = float(spatial["diagnostics"][key])
    global_nrmse = float(global_["diagnostics"][key])
    best = min(spatial_nrmse, global_nrmse)
    worst = max(spatial_nrmse, global_nrmse)
    rel_ratio = float(worst / max(best, 1.0e-12))
    abs_gap = float(worst - best)
    max_ratio = float(gates.get("comparability_max_nrmse_ratio", 1.25))
    max_abs_gap = float(gates.get("comparability_max_abs_nrmse_gap", 0.005))
    spatial_pass = bool(
        spatial.get("quality_gates", {}).get("primary_pass", False)
        or spatial.get("quality_gates", {}).get("backup_pass", False)
    )
    global_pass = bool(
        global_.get("quality_gates", {}).get("primary_pass", False)
        or global_.get("quality_gates", {}).get("backup_pass", False)
    )
    comparable = bool(spatial_pass and global_pass and (rel_ratio <= max_ratio or abs_gap <= max_abs_gap))
    summary = {
        "selection_split": split,
        "spatial_checkpoint": spatial.get("checkpoint"),
        "global_checkpoint": global_.get("checkpoint"),
        "spatial_nrmse": spatial_nrmse,
        "global_nrmse": global_nrmse,
        "nrmse_ratio_worst_over_best": rel_ratio,
        "nrmse_abs_gap": abs_gap,
        "spatial_gate_pass": spatial_pass,
        "global_gate_pass": global_pass,
        "comparable": comparable,
        "thresholds": {
            "max_ratio": max_ratio,
            "max_abs_gap": max_abs_gap,
        },
        "spatial_parameter_count": spatial.get("parameter_count"),
        "global_parameter_count": global_.get("parameter_count"),
        "spatial_elapsed": spatial.get("elapsed"),
        "global_elapsed": global_.get("elapsed"),
    }
    out_path = Path(config["paths"]["figures_dir"]) / "ks_ae_comparison_summary.json"
    save_json(summary, out_path)
    if not comparable:
        raise RuntimeError(f"Spatial/global AE quality is not comparable; see {out_path}")
    return out_path
