evaluate.py

"""Evaluation with advanced academic benchmarks for phase-2 experiments."""

from __future__ import annotations

import math
from pathlib import Path
from typing import Dict, Iterable, Tuple

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.utils.data import DataLoader

from config import ExperimentConfig, ensure_output_dirs, update_config_from_args
from dataset import JitteredCovarianceDataset
from model import JitterRobustUnfoldedSBLNet


def resolve_device(args) -> torch.device:
    """Resolve device in a strict and explicit way."""

    requested = str(getattr(args, "device", "auto")).lower()
    if getattr(args, "cpu", False):
        requested = "cpu"

    if requested == "auto":
        requested = "cuda" if torch.cuda.is_available() else "cpu"

    if requested == "cuda":
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA was requested but is unavailable.")
        gpu_id = int(getattr(args, "gpu_id", 0))
        if gpu_id < 0 or gpu_id >= torch.cuda.device_count():
            raise ValueError(
                f"Invalid --gpu_id={gpu_id}. Available GPU count: {torch.cuda.device_count()}."
            )
        device = torch.device(f"cuda:{gpu_id}")
        torch.cuda.set_device(device)
        return device

    if requested == "cpu":
        return torch.device("cpu")

    raise ValueError(f"Unsupported --device value: {requested}")


def ri_to_complex(cov_ri: torch.Tensor) -> torch.Tensor:
    """Convert covariance channels to complex matrix.

    Args:
        cov_ri: [2, M, M] tensor.
    """

    return cov_ri[0].to(torch.complex64) + 1j * cov_ri[1].to(torch.complex64)


def torch_steering(
    angles_deg: torch.Tensor,
    sensor_positions_half_lambda: torch.Tensor,
) -> torch.Tensor:
    """Build complex steering matrix in PyTorch.

    a_m(theta) = exp(j * pi * d_m * sin(theta)).
    """

    theta = torch.deg2rad(angles_deg)
    phase = math.pi * sensor_positions_half_lambda[:, None] * torch.sin(theta)[None, :]
    return torch.exp(1j * phase.to(torch.complex64))


def normalize_spectrum(spec: torch.Tensor) -> torch.Tensor:
    """Normalize non-negative spectrum to [0, 1]."""

    spec = torch.real(spec)
    spec = torch.clamp(spec, min=0.0)
    return spec / (torch.max(spec) + 1e-8)


def spectrum_rmse(pred: torch.Tensor, target: torch.Tensor) -> float:
    """Spectrum-domain RMSE."""

    return float(torch.sqrt(torch.mean((pred - target) ** 2)).item())


def music_spectrum(
    covariance: torch.Tensor,
    scan_angles_deg: torch.Tensor,
    sensor_positions_half_lambda: torch.Tensor,
    num_sources: int,
) -> torch.Tensor:
    """Standard MUSIC pseudo-spectrum."""

    m = covariance.size(0)
    k_eff = max(1, min(int(num_sources), m - 1))
    cov_h = (covariance + covariance.conj().T) * 0.5
    evals, evecs = torch.linalg.eigh(cov_h)
    noise_subspace = evecs[:, : m - k_eff]
    proj_n = noise_subspace @ noise_subspace.conj().T

    A = torch_steering(scan_angles_deg, sensor_positions_half_lambda)  # [M, N]
    proj_a = proj_n @ A
    denom = torch.real(torch.sum(A.conj() * proj_a, dim=0))
    pseudo = 1.0 / torch.clamp(denom, min=1e-8)
    return normalize_spectrum(pseudo)


def build_lag_profile(
    covariance: torch.Tensor,
    sensor_positions_half_lambda: torch.Tensor,
) -> dict[int, torch.Tensor]:
    """Average covariance entries with identical difference co-array lag."""

    positions = sensor_positions_half_lambda.to(torch.int64)
    lag_values: dict[int, list[torch.Tensor]] = {}
    m = covariance.size(0)
    for i in range(m):
        for j in range(m):
            lag = int((positions[i] - positions[j]).item())
            lag_values.setdefault(lag, []).append(covariance[i, j])

    lag_profile: dict[int, torch.Tensor] = {}
    for lag, vals in lag_values.items():
        stacked = torch.stack(vals)
        lag_profile[lag] = torch.mean(stacked)
    return lag_profile


def contiguous_center_lag_span(lag_profile: dict[int, torch.Tensor]) -> int:
    """Find largest contiguous lag span around zero: [-L, ..., L]."""

    l = 0
    while (l + 1 in lag_profile) and (-(l + 1) in lag_profile):
        l += 1
    return l


def lag_profile_to_toeplitz(
    lag_profile: dict[int, torch.Tensor],
    span: int,
    device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Create Toeplitz virtual covariance and observation mask from lag profile."""

    n = span + 1  # ULA size from contiguous lag span
    R = torch.zeros((n, n), device=device, dtype=torch.complex64)
    mask = torch.zeros((n, n), device=device, dtype=torch.float32)
    for i in range(n):
        for j in range(n):
            lag = i - j
            if lag in lag_profile:
                R[i, j] = lag_profile[lag]
                mask[i, j] = 1.0

    R = 0.5 * (R + R.conj().T)
    return R, mask


def spatial_smoothing(covariance: torch.Tensor, subarray_len: int) -> torch.Tensor:
    """Forward spatial smoothing over overlapping virtual ULA subarrays."""

    n = covariance.size(0)
    p = int(subarray_len)
    if p > n:
        raise ValueError("subarray_len must not exceed covariance size.")
    num_sub = n - p + 1
    smoothed = torch.zeros((p, p), device=covariance.device, dtype=covariance.dtype)
    for s in range(num_sub):
        smoothed += covariance[s : s + p, s : s + p]
    smoothed = smoothed / float(num_sub)
    smoothed = 0.5 * (smoothed + smoothed.conj().T)
    return smoothed


def coarray_ss_music_spectrum(
    covariance: torch.Tensor,
    cfg: ExperimentConfig,
    device: torch.device,
) -> torch.Tensor:
    """Co-array SS-MUSIC for sparse coprime array."""

    positions = torch.tensor(cfg.array_positions_half_lambda, device=device, dtype=torch.float32)
    lag_profile = build_lag_profile(covariance, positions)
    span = contiguous_center_lag_span(lag_profile)
    if span < 2:
        scan = torch.tensor(cfg.angle_grid_deg, device=device)
        return music_spectrum(covariance, scan, positions, cfg.num_targets)

    Rv, _ = lag_profile_to_toeplitz(lag_profile, span=span, device=device)
    nv = Rv.size(0)
    sub_len = min(nv, max(cfg.num_targets + 1, nv // 2))
    Rss = spatial_smoothing(Rv, subarray_len=sub_len)

    scan = torch.tensor(cfg.angle_grid_deg, device=device)
    virtual_pos = torch.arange(sub_len, device=device, dtype=torch.float32)
    k_eff = min(cfg.num_targets, sub_len - 1)
    return music_spectrum(Rss, scan, virtual_pos, k_eff)


def svt_matrix_completion(
    observed_toeplitz: torch.Tensor,
    mask: torch.Tensor,
    num_iter: int = 30,
    tau: float = 0.08,
) -> torch.Tensor:
    """Low-rank matrix completion via iterative singular value thresholding."""

    x = observed_toeplitz.clone()
    mask_c = mask.to(torch.complex64)
    eps_eye = 1e-6 * torch.eye(x.size(0), device=x.device, dtype=x.dtype)

    for _ in range(num_iter):
        x = 0.5 * (x + x.conj().T)
        u, s, vh = torch.linalg.svd(x + eps_eye, full_matrices=False)
        s_new = torch.relu(s - tau)
        x_low_rank = (u * s_new.unsqueeze(0)) @ vh
        x = mask_c * observed_toeplitz + (1.0 - mask_c) * x_low_rank
        x = 0.5 * (x + x.conj().T)

    return x


def matrix_completion_ss_music_spectrum(
    covariance: torch.Tensor,
    cfg: ExperimentConfig,
    device: torch.device,
) -> torch.Tensor:
    """Matrix Completion + SS-MUSIC benchmark."""

    positions = torch.tensor(cfg.array_positions_half_lambda, device=device, dtype=torch.float32)
    lag_profile = build_lag_profile(covariance, positions)
    max_lag = int(torch.max(positions).item())
    if max_lag < 2:
        scan = torch.tensor(cfg.angle_grid_deg, device=device)
        return music_spectrum(covariance, scan, positions, cfg.num_targets)

    R_obs, mask = lag_profile_to_toeplitz(lag_profile, span=max_lag, device=device)
    R_comp = svt_matrix_completion(R_obs, mask, num_iter=35, tau=0.07)

    nv = R_comp.size(0)
    sub_len = min(nv, max(cfg.num_targets + 1, nv // 2))
    Rss = spatial_smoothing(R_comp, subarray_len=sub_len)

    scan = torch.tensor(cfg.angle_grid_deg, device=device)
    virtual_pos = torch.arange(sub_len, device=device, dtype=torch.float32)
    k_eff = min(cfg.num_targets, sub_len - 1)
    return music_spectrum(Rss, scan, virtual_pos, k_eff)


def gridless_sbl_spectrum(
    covariance: torch.Tensor,
    cfg: ExperimentConfig,
    device: torch.device,
    fine_step_deg: float = 0.25,
    num_iter: int = 18,
) -> torch.Tensor:
    """Fast iterative gridless-like SBL with off-grid peak refinement.

    Procedure:
        1) Covariance-domain SBL updates on a fine grid.
        2) Quadratic off-grid interpolation around dominant peaks.
        3) Reconstruct narrow spectrum on target 1-degree grid.
    """

    positions = torch.tensor(cfg.array_positions_half_lambda, device=device, dtype=torch.float32)
    fine_grid = torch.arange(
        cfg.angle_min_deg,
        cfg.angle_max_deg + 0.5 * fine_step_deg,
        fine_step_deg,
        device=device,
        dtype=torch.float32,
    )
    A = torch_steering(fine_grid, positions)  # [M, Nf]
    m, nf = A.size()
    eye = torch.eye(m, device=device, dtype=torch.complex64)

    gamma = torch.full((nf,), 0.1, device=device, dtype=torch.float32)
    sigma2 = torch.clamp(torch.real(torch.trace(covariance)) / (20.0 * m), min=1e-6)

    for _ in range(num_iter):
        Ag = A * gamma.unsqueeze(0).to(torch.complex64)
        sigma_y = Ag @ A.conj().T + sigma2 * eye + 1e-5 * eye
        sigma_inv = torch.linalg.pinv(sigma_y)
        middle = sigma_inv @ covariance @ sigma_inv

        q = torch.real(torch.sum(A.conj() * (sigma_inv @ A), dim=0)).clamp_min(1e-8)
        p = torch.real(torch.sum(A.conj() * (middle @ A), dim=0)).clamp_min(1e-8)
        gamma = torch.clamp(gamma * torch.sqrt(p / q), min=1e-8, max=1e4)

        cov_model = (A * gamma.unsqueeze(0).to(torch.complex64)) @ A.conj().T
        residual = covariance - cov_model
        sigma2 = torch.clamp(torch.real(torch.trace(residual)) / float(m), min=1e-6)

    k = min(cfg.num_targets, nf)
    vals, idx = torch.topk(gamma, k=k)
    idx = idx.sort().values
    vals = gamma[idx]

    # Off-grid quadratic interpolation for sub-grid angle estimation.
    refined_angles = []
    for peak_idx in idx:
        i = int(peak_idx.item())
        theta_c = fine_grid[i]
        if 1 <= i < (nf - 1):
            yl = gamma[i - 1]
            yc = gamma[i]
            yr = gamma[i + 1]
            denom = yl - 2.0 * yc + yr
            delta = 0.5 * (yl - yr) / (denom + 1e-8)
            delta = torch.clamp(delta, -0.5, 0.5)
            theta_c = theta_c + delta * fine_step_deg
        refined_angles.append(theta_c)
    refined_angles_t = torch.stack(refined_angles)

    # Render on target grid as narrow Gaussian needles.
    scan = torch.tensor(cfg.angle_grid_deg, device=device, dtype=torch.float32)
    spec = torch.zeros_like(scan)
    width = 0.45
    for amp, theta in zip(vals, refined_angles_t):
        spec += amp * torch.exp(-0.5 * ((scan - theta) / width) ** 2)
    return normalize_spectrum(spec)


def infer_all_spectra(
    model: torch.nn.Module,
    cov_ri: torch.Tensor,
    cfg: ExperimentConfig,
    device: torch.device,
) -> dict[str, torch.Tensor]:
    """Run proposed model and all benchmark algorithms for one sample."""

    model_spec = model(cov_ri.unsqueeze(0)).squeeze(0)
    cov = ri_to_complex(cov_ri)
    scan = torch.tensor(cfg.angle_grid_deg, device=device)
    positions = torch.tensor(cfg.array_positions_half_lambda, device=device, dtype=torch.float32)

    spec_music = music_spectrum(cov, scan, positions, cfg.num_targets)
    spec_ss = coarray_ss_music_spectrum(cov, cfg, device)
    spec_mc_ss = matrix_completion_ss_music_spectrum(cov, cfg, device)
    spec_gridless = gridless_sbl_spectrum(cov, cfg, device)

    return {
        "unfolded": normalize_spectrum(model_spec),
        "mc_ss_music": normalize_spectrum(spec_mc_ss),
        "gridless_sbl": normalize_spectrum(spec_gridless),
        "coarray_ss_music": normalize_spectrum(spec_ss),
        "music": normalize_spectrum(spec_music),
    }


@torch.no_grad()
def evaluate_loader(
    model: torch.nn.Module,
    loader: DataLoader,
    cfg: ExperimentConfig,
    device: torch.device,
) -> dict[str, float]:
    """Evaluate all methods and return mean RMSE dictionary."""

    model.eval()
    keys = ["unfolded", "mc_ss_music", "gridless_sbl", "coarray_ss_music", "music"]
    errors = {k: [] for k in keys}

    for batch in loader:
        cov_batch = batch["cov_ri"].to(device)
        target_batch = batch["spectrum"].to(device)
        b = cov_batch.size(0)

        for i in range(b):
            spectra = infer_all_spectra(model, cov_batch[i], cfg, device)
            gt = target_batch[i]
            for k in keys:
                errors[k].append(spectrum_rmse(spectra[k], gt))

    return {k: float(np.mean(v)) for k, v in errors.items()}


def plot_rmse_curve(
    x_values: Iterable[float],
    series: dict[str, Iterable[float]],
    xlabel: str,
    title: str,
    output_path: Path,
) -> None:
    """Plot paper-ready RMSE comparison for all five methods."""

    x = np.array(list(x_values), dtype=np.float32)

    style = {
        "unfolded": dict(label="Unfolded SBL Net", color="blue", linestyle="-"),
        "mc_ss_music": dict(label="Matrix Completion + SS-MUSIC", color="red", linestyle="--"),
        "gridless_sbl": dict(label="Gridless SBL", color="black", linestyle=":"),
        "coarray_ss_music": dict(label="Co-array SS-MUSIC", color="orange", linestyle=":"),
        "music": dict(label="Standard MUSIC", color="green", linestyle=":"),
    }

    plt.figure(figsize=(9, 6))
    for key in ["unfolded", "mc_ss_music", "gridless_sbl", "coarray_ss_music", "music"]:
        y = np.array(list(series[key]), dtype=np.float32)
        plt.plot(
            x,
            y,
            linewidth=2.5,
            color=style[key]["color"],
            linestyle=style[key]["linestyle"],
            label=style[key]["label"],
        )

    plt.xlabel(xlabel, fontsize=14)
    plt.ylabel("Spectrum RMSE", fontsize=14)
    plt.title(title, fontsize=14)
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)
    plt.grid(True)
    plt.legend(fontsize=11)
    plt.tight_layout()
    plt.savefig(output_path, dpi=220)
    plt.close()


@torch.no_grad()
def plot_sample_spatial_spectrum(
    model: torch.nn.Module,
    cfg: ExperimentConfig,
    device: torch.device,
    output_path: Path,
    snr_db: float = 10.0,
    num_snapshots: int = 10,
) -> None:
    """Plot one sample spectrum with all five methods."""

    sample_ds = JitteredCovarianceDataset(
        num_samples=1,
        cfg=cfg,
        fixed_snr_db=snr_db,
        fixed_num_snapshots=num_snapshots,
        seed=cfg.seed + 123456,
    )
    sample = sample_ds[0]
    cov_ri = sample["cov_ri"].to(device)
    gt = sample["spectrum"].cpu().numpy()
    spectra = infer_all_spectra(model, cov_ri, cfg, device)
    angles = cfg.angle_grid_deg

    style = {
        "unfolded": dict(label="Unfolded SBL Net", color="blue", linestyle="-"),
        "mc_ss_music": dict(label="Matrix Completion + SS-MUSIC", color="red", linestyle="--"),
        "gridless_sbl": dict(label="Gridless SBL", color="black", linestyle=":"),
        "coarray_ss_music": dict(label="Co-array SS-MUSIC", color="orange", linestyle=":"),
        "music": dict(label="Standard MUSIC", color="green", linestyle=":"),
    }

    plt.figure(figsize=(10, 6))
    for key in ["unfolded", "mc_ss_music", "gridless_sbl", "coarray_ss_music", "music"]:
        plt.plot(
            angles,
            spectra[key].detach().cpu().numpy(),
            linewidth=2.5,
            color=style[key]["color"],
            linestyle=style[key]["linestyle"],
            label=style[key]["label"],
        )
    plt.plot(
        angles,
        gt,
        linewidth=2.5,
        color="magenta",
        linestyle="-.",
        label="Ground Truth",
    )

    plt.xlabel("Angle (deg)", fontsize=14)
    plt.ylabel("Normalized Spectrum", fontsize=14)
    plt.title(
        f"Sample Spatial Spectrum (SNR={snr_db:.1f} dB, L={num_snapshots}, K={cfg.num_targets})",
        fontsize=14,
    )
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)
    plt.grid(True)
    plt.legend(fontsize=10)
    plt.tight_layout()
    plt.savefig(output_path, dpi=220)
    plt.close()


def _load_checkpoint(checkpoint_path: Path, args) -> tuple[ExperimentConfig, dict]:
    """Load model checkpoint and recover training config."""

    if not checkpoint_path.exists():
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)

    cfg = ExperimentConfig()
    if "config" in ckpt and isinstance(ckpt["config"], dict):
        cfg = ExperimentConfig(**ckpt["config"])
    cfg = update_config_from_args(cfg, args)
    return cfg, ckpt


def run_evaluation(args) -> None:
    """Run full benchmark evaluation and generate paper-ready figures."""

    base_cfg = update_config_from_args(ExperimentConfig(), args)
    ckpt_path_override = getattr(args, "checkpoint_path_override", None)
    ckpt_path = Path(ckpt_path_override) if ckpt_path_override else Path(base_cfg.checkpoint_path)

    cfg, ckpt = _load_checkpoint(ckpt_path, args)
    ensure_output_dirs(cfg)

    device = resolve_device(args)
    if device.type == "cuda":
        print(f"[Eval] Device: {device} ({torch.cuda.get_device_name(device)})")
    else:
        print(f"[Eval] Device: {device}")
    print(f"[Eval] Checkpoint: {ckpt_path}")

    model = JitterRobustUnfoldedSBLNet(cfg).to(device)
    model.load_state_dict(ckpt["model_state_dict"])
    model.eval()

    eval_samples = int(getattr(args, "eval_samples_per_setting", 160))
    eval_bs = int(getattr(args, "eval_batch_size", 32))
    eval_workers = int(getattr(args, "eval_num_workers", 0))
    fixed_snr = float(getattr(args, "snapshot_curve_snr_db", 10.0))

    # RMSE vs SNR
    snr_values = np.arange(cfg.snr_min_db, cfg.snr_max_db + 1e-6, 5.0, dtype=np.float32)
    snr_curves = {
        "unfolded": [],
        "mc_ss_music": [],
        "gridless_sbl": [],
        "coarray_ss_music": [],
        "music": [],
    }
    for snr in snr_values:
        ds = JitteredCovarianceDataset(
            num_samples=eval_samples,
            cfg=cfg,
            fixed_snr_db=float(snr),
            fixed_num_snapshots=cfg.num_snapshots,
            seed=cfg.seed + int(100 * snr) + 3000,
        )
        loader = DataLoader(
            ds,
            batch_size=eval_bs,
            shuffle=False,
            num_workers=eval_workers,
            pin_memory=(device.type == "cuda"),
            persistent_workers=(eval_workers > 0),
        )
        metrics = evaluate_loader(model, loader, cfg, device)
        for key in snr_curves:
            snr_curves[key].append(metrics[key])
        print(
            f"[Eval][SNR={snr:.1f}dB] "
            f"Unfolded={metrics['unfolded']:.4f}, "
            f"MC+SS={metrics['mc_ss_music']:.4f}, "
            f"Gridless={metrics['gridless_sbl']:.4f}, "
            f"SS-MUSIC={metrics['coarray_ss_music']:.4f}, "
            f"MUSIC={metrics['music']:.4f}"
        )

    rmse_snr_path = Path(cfg.figure_dir) / "rmse_vs_snr_phase2.png"
    plot_rmse_curve(
        x_values=snr_values,
        series=snr_curves,
        xlabel="SNR (dB)",
        title=f"RMSE vs SNR (L={cfg.num_snapshots})",
        output_path=rmse_snr_path,
    )

    # RMSE vs snapshots
    snapshot_values = [2, 4, 6, 8, 10, 12, 16, 20]
    l_curves = {
        "unfolded": [],
        "mc_ss_music": [],
        "gridless_sbl": [],
        "coarray_ss_music": [],
        "music": [],
    }
    for l in snapshot_values:
        ds = JitteredCovarianceDataset(
            num_samples=eval_samples,
            cfg=cfg,
            fixed_snr_db=fixed_snr,
            fixed_num_snapshots=int(l),
            seed=cfg.seed + l + 9000,
        )
        loader = DataLoader(
            ds,
            batch_size=eval_bs,
            shuffle=False,
            num_workers=eval_workers,
            pin_memory=(device.type == "cuda"),
            persistent_workers=(eval_workers > 0),
        )
        metrics = evaluate_loader(model, loader, cfg, device)
        for key in l_curves:
            l_curves[key].append(metrics[key])
        print(
            f"[Eval][L={l}] "
            f"Unfolded={metrics['unfolded']:.4f}, "
            f"MC+SS={metrics['mc_ss_music']:.4f}, "
            f"Gridless={metrics['gridless_sbl']:.4f}, "
            f"SS-MUSIC={metrics['coarray_ss_music']:.4f}, "
            f"MUSIC={metrics['music']:.4f}"
        )

    rmse_l_path = Path(cfg.figure_dir) / "rmse_vs_snapshots_phase2.png"
    plot_rmse_curve(
        x_values=snapshot_values,
        series=l_curves,
        xlabel="Number of Snapshots L",
        title=f"RMSE vs Snapshots (SNR={fixed_snr:.1f} dB)",
        output_path=rmse_l_path,
    )

    sample_path = Path(cfg.figure_dir) / "sample_spatial_spectrum_phase2.png"
    plot_sample_spatial_spectrum(
        model=model,
        cfg=cfg,
        device=device,
        output_path=sample_path,
        snr_db=fixed_snr,
        num_snapshots=cfg.num_snapshots,
    )

    print(f"[Eval] Saved: {rmse_snr_path}")
    print(f"[Eval] Saved: {rmse_l_path}")
    print(f"[Eval] Saved: {sample_path}")