src/pino/evaluate.py

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn as nn
from scipy.spatial.distance import cdist
from scipy.stats import spearmanr
from torch.utils.data import DataLoader

from pino.heads import PIMTHeads
from pino.pimt_model import FragranceTrajectoryDataset, PhysicsInformedMixtureTransformer
from pino.train import pad_trajectory_collate

logger = logging.getLogger("pino.evaluate")


@torch.no_grad()
def extract_embeddings(
    model: nn.Module,
    heads: nn.Module,
    loader: DataLoader,
    device: torch.device,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Return predicted/target descriptor profiles plus genre labels."""
    model.eval()
    heads.eval()
    predictions: list[np.ndarray] = []
    targets: list[np.ndarray] = []
    genre_labels: list[np.ndarray] = []
    tier_predictions: list[np.ndarray] = []
    for batch in loader:
        tokens = batch["tokens"].to(device)
        physics = batch["physics"].to(device)
        src_key_padding_mask = batch["src_key_padding_mask"].to(device)
        target_obj = batch["target_obj"].to(device)
        latent = model(tokens, physics, src_key_padding_mask)
        pred = heads(latent, physics, src_key_padding_mask)

        tier_pred = pred["objective"].cpu().numpy()
        tier_target = target_obj.cpu().numpy()
        tier_predictions.append(tier_pred)

        # Collapse the top/middle/base pyramid into one perceptual descriptor
        # profile for retrieval/distance metrics.
        predictions.append(tier_pred.mean(axis=1))
        targets.append(tier_target.mean(axis=1))
        genre_labels.append(batch["genre_label"].cpu().numpy())

    return (
        np.vstack(predictions),
        np.vstack(targets),
        np.concatenate(genre_labels),
        np.concatenate(tier_predictions, axis=0),
    )


def spearman_distance_correlation(embeddings: np.ndarray, targets: np.ndarray) -> dict[str, float]:
    """
    Compute Spearman correlation between latent-space distances and target-space distances.
    Returns the mean and per-sample Spearman rho between latent and descriptor distances.
    """
    latent_dist = cdist(embeddings, embeddings, metric="euclidean")
    target_dist = cdist(targets, targets, metric="euclidean")
    # Use upper triangular indices to avoid self-pairs and duplicates.
    n = len(embeddings)
    idx = np.triu_indices(n, k=1)
    lat_vec = latent_dist[idx]
    tgt_vec = target_dist[idx]
    rho, p_value = spearmanr(lat_vec, tgt_vec)
    return {
        "spearman_rho": float(rho),
        "spearman_p_value": float(p_value),
        "n_pairs": int(len(lat_vec)),
    }


def mean_average_precision(embeddings: np.ndarray, labels: np.ndarray, k: int | None = None) -> dict[str, float]:
    """
    Compute retrieval mAP in the latent space using genre labels as relevance.
    For each sample, positives are same-genre samples; negatives are all others.
    """
    k = k or len(embeddings) - 1
    distances = cdist(embeddings, embeddings, metric="euclidean")
    np.fill_diagonal(distances, np.inf)
    aps: list[float] = []
    for i in range(len(embeddings)):
        sorted_idx = np.argsort(distances[i])[:k]
        relevant = (labels[sorted_idx] == labels[i]).astype(np.int32)
        if relevant.sum() == 0:
            continue
        # Cumulative precision at each relevant rank.
        precisions = []
        num_relevant = 0
        for rank, is_rel in enumerate(relevant, start=1):
            if is_rel:
                num_relevant += 1
                precisions.append(num_relevant / rank)
        aps.append(np.mean(precisions) if precisions else 0.0)
    return {
        "mAP": float(np.mean(aps)) if aps else 0.0,
        "AP_std": float(np.std(aps)) if aps else 0.0,
        "n_queries": len(aps),
    }


def _infer_checkpoint_config(checkpoint: dict[str, Any]) -> dict[str, int]:
    """Infer architecture knobs saved by the current trainer."""
    model_state = checkpoint["model_state_dict"]
    input_proj = model_state["input_proj.weight"]
    hidden_dim = int(input_proj.shape[0])
    embedding_dim = int(input_proj.shape[1])
    state_dim = int(model_state["gating.physics_scaler"].shape[0])

    layer_prefixes = {
        key.split(".")[2]
        for key in model_state
        if key.startswith("encoder.layers.") and key.endswith("self_attn.in_proj_weight")
    }
    num_layers = len(layer_prefixes) or 4

    # num_heads is not represented in the state dict.  Current HF training uses
    # four heads; fall back to the largest common divisor up to 8 for older runs.
    num_heads = 4 if hidden_dim % 4 == 0 else next(h for h in (8, 2, 1) if hidden_dim % h == 0)
    return {
        "hidden_dim": hidden_dim,
        "embedding_dim": embedding_dim,
        "state_dim": state_dim,
        "num_layers": num_layers,
        "num_heads": num_heads,
    }


def evaluate(checkpoint_path: str, data_path: str, batch_size: int = 16, hidden_dim: int | None = None) -> dict[str, Any]:
    """Load a checkpoint and run perceptual evaluation metrics."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    checkpoint = torch.load(checkpoint_path, map_location=device)
    cfg = _infer_checkpoint_config(checkpoint)
    if hidden_dim is not None and hidden_dim != cfg["hidden_dim"]:
        logger.warning("Ignoring --hidden-dim=%d; checkpoint requires hidden_dim=%d", hidden_dim, cfg["hidden_dim"])

    all_genres = ["citrus_cologne", "fougere", "floral_woody", "amber_oriental", "wildcard"]
    genre_map = {genre: idx for idx, genre in enumerate(all_genres)}
    dataset = FragranceTrajectoryDataset(data_path=data_path, use_embedding_fallback=True, genre_map=genre_map)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=pad_trajectory_collate)

    model = PhysicsInformedMixtureTransformer(
        embedding_dim=cfg["embedding_dim"],
        state_dim=cfg["state_dim"],
        hidden_dim=cfg["hidden_dim"],
        num_heads=cfg["num_heads"],
        num_layers=cfg["num_layers"],
    ).to(device)
    heads = PIMTHeads(hidden_dim=cfg["hidden_dim"], objective_dim=138).to(device)
    model.load_state_dict(checkpoint["model_state_dict"])
    heads.load_state_dict(checkpoint["heads_state_dict"])

    predictions, targets, genre_labels, tier_predictions = extract_embeddings(model, heads, loader, device)

    spearman_metrics = spearman_distance_correlation(predictions, targets)
    map_metrics = mean_average_precision(predictions, genre_labels)
    mse = float(np.mean((predictions - targets) ** 2))
    mae = float(np.mean(np.abs(predictions - targets)))

    metrics = {
        "spearman": spearman_metrics,
        "retrieval_mAP": map_metrics,
        "prediction": {
            "profile_mse": mse,
            "profile_mae": mae,
            "prediction_variance": float(np.var(predictions)),
            "tier_prediction_variance": float(np.var(tier_predictions)),
        },
        "n_samples": len(predictions),
        "embedding_dim": predictions.shape[1],
        "checkpoint": {
            "epoch": checkpoint.get("epoch"),
            "val_loss": checkpoint.get("val_loss"),
            **cfg,
        },
    }

    logger.info("=" * 50)
    logger.info("PERCEPTUAL EVALUATION RESULTS")
    logger.info("=" * 50)
    logger.info("Spearman distance correlation: rho=%.4f | p=%.4e | n_pairs=%d",
                spearman_metrics["spearman_rho"], spearman_metrics["spearman_p_value"], spearman_metrics["n_pairs"])
    logger.info("Retrieval mAP@all: %.4f +/- %.4f | n_queries=%d",
                map_metrics["mAP"], map_metrics["AP_std"], map_metrics["n_queries"])
    logger.info("Profile prediction: mse=%.4f | mae=%.4f | variance=%.6f", mse, mae, metrics["prediction"]["prediction_variance"])
    logger.info("=" * 50)
    return metrics


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate PINO perceptual metrics")
    parser.add_argument("--checkpoint", default="models/pimt_v1.pt", help="Path to trained checkpoint")
    parser.add_argument("--data", default="data/synthetic_dataset_v2.jsonl", help="Dataset JSON-L")
    parser.add_argument("--batch-size", type=int, default=16, help="Evaluation batch size")
    parser.add_argument("--hidden-dim", type=int, default=None, help="Deprecated; architecture is inferred from checkpoint")
    parser.add_argument("--output", default="metrics.json", help="Path to write JSON metrics")
    parser.add_argument("--log-level", default="INFO")
    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level.upper(), logging.INFO),
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )

    metrics = evaluate(args.checkpoint, args.data, args.batch_size, args.hidden_dim)
    Path(args.output).write_text(json.dumps(metrics, indent=2))
    print(json.dumps(metrics, indent=2))