evaluate.py

"""
Evaluation script for NeurophysicsLab models.

Loads trained PhysicsNet checkpoints and computes comprehensive metrics
including prediction MSE, energy violation, R² score, and latent space
statistics across all three physical systems.

Usage::

    python evaluate.py
    python evaluate.py --systems projectile spring
    python evaluate.py --systems orbit --n_test 500
"""

import argparse
import json
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.metrics import r2_score

# Ensure the project root is importable
PROJECT_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(PROJECT_ROOT))

from models.physics_net import PhysicsNet, TrajectoryDataset, get_system_dims
from models.losses import energy_violation_metric
from utils import set_seed, load_model, load_trajectories





@torch.no_grad()
def evaluate_system(
    model: PhysicsNet,
    trajectories: list[np.ndarray],
    system: str,
    window: int = 20,
    device: str = "cpu",
) -> dict[str, float]:
    """Evaluate a trained model on a set of trajectories.

    Computes prediction MSE, R² score, energy violation, and latent space
    statistics over all sliding windows from the provided trajectories.

    Args:
        model: Trained PhysicsNet in eval mode.
        trajectories: List of raw trajectory arrays for testing.
        system: Physical system name.
        window: Sliding window size. Default ``20``.
        device: Computation device.

    Returns:
        Dictionary with keys: ``mse``, ``r2``, ``energy_violation``,
        ``latent_mean_std``, ``latent_mean_abs``, ``n_samples``.
    """
    model.eval()
    input_dim, state_cols, energy_col = get_system_dims(system)

    all_pred = []
    all_true = []
    all_pred_energy = []
    all_true_energy = []
    all_latents = []

    for traj in trajectories:
        if traj.shape[0] <= window:
            continue

        states = traj[:, state_cols].astype(np.float32)
        energies = traj[:, energy_col].astype(np.float32)

        for start in range(traj.shape[0] - window):
            win = torch.from_numpy(states[start:start + window]).unsqueeze(0).to(device)
            target = states[start + window]
            true_e = energies[start + window]

            pred_state, latent, pred_e = model(win)

            all_pred.append(pred_state.cpu().numpy().flatten())
            all_true.append(target)
            all_pred_energy.append(pred_e.cpu().item())
            all_true_energy.append(true_e)
            all_latents.append(latent.cpu().numpy().flatten())

    all_pred = np.array(all_pred)
    all_true = np.array(all_true)
    all_pred_energy = np.array(all_pred_energy)
    all_true_energy = np.array(all_true_energy)
    all_latents = np.array(all_latents)

    # MSE
    mse = float(np.mean((all_pred - all_true) ** 2))

    # R² score
    r2 = float(r2_score(all_true.flatten(), all_pred.flatten()))

    # Energy violation
    ev = float(np.mean(np.abs(all_pred_energy - all_true_energy)))

    # Latent space statistics
    latent_std_per_dim = np.std(all_latents, axis=0)
    latent_mean_std = float(np.mean(latent_std_per_dim))
    latent_mean_abs = float(np.mean(np.abs(all_latents)))

    return {
        "mse": mse,
        "r2": r2,
        "energy_violation": ev,
        "latent_mean_std": latent_mean_std,
        "latent_mean_abs": latent_mean_abs,
        "n_samples": len(all_pred),
    }


def plot_comparison(
    results: dict[str, dict[str, float]],
    output_dir: Path,
) -> None:
    """Generate comparison bar charts across systems.

    Creates a figure with three subplots comparing MSE, energy violation,
    and R² score across all evaluated systems.

    Args:
        results: Dictionary mapping system names to their metric dicts.
        output_dir: Directory to save the output figure.
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    systems = list(results.keys())
    if not systems:
        return

    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    fig.suptitle("Model Evaluation Comparison", fontsize=14, fontweight="bold")

    colors = ["#6366f1", "#22c55e", "#f97316"]

    # MSE comparison
    mse_values = [results[s]["mse"] for s in systems]
    bars = axes[0].bar(systems, mse_values, color=colors[:len(systems)], alpha=0.85)
    axes[0].set_ylabel("MSE")
    axes[0].set_title("Prediction MSE")
    axes[0].set_yscale("log")
    for bar, val in zip(bars, mse_values):
        axes[0].text(
            bar.get_x() + bar.get_width() / 2, bar.get_height(),
            f"{val:.2e}", ha="center", va="bottom", fontsize=9,
        )

    # Energy violation
    ev_values = [results[s]["energy_violation"] for s in systems]
    bars = axes[1].bar(systems, ev_values, color=colors[:len(systems)], alpha=0.85)
    axes[1].set_ylabel("Mean |ΔE|")
    axes[1].set_title("Energy Violation")
    for bar, val in zip(bars, ev_values):
        axes[1].text(
            bar.get_x() + bar.get_width() / 2, bar.get_height(),
            f"{val:.4f}", ha="center", va="bottom", fontsize=9,
        )

    # R² score
    r2_values = [results[s]["r2"] for s in systems]
    bars = axes[2].bar(systems, r2_values, color=colors[:len(systems)], alpha=0.85)
    axes[2].set_ylabel("R²")
    axes[2].set_title("Prediction R² Score")
    axes[2].set_ylim(min(0, min(r2_values) - 0.1), 1.05)
    for bar, val in zip(bars, r2_values):
        axes[2].text(
            bar.get_x() + bar.get_width() / 2, bar.get_height(),
            f"{val:.4f}", ha="center", va="bottom", fontsize=9,
        )

    plt.tight_layout()
    fig_path = output_dir / "evaluation_comparison.png"
    plt.savefig(fig_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"[PLOT] Comparison chart saved to {fig_path}")


def plot_latent_distributions(
    model: PhysicsNet,
    trajectories: list[np.ndarray],
    system: str,
    output_dir: Path,
    window: int = 20,
    max_trajs: int = 50,
) -> None:
    """Plot latent space distribution for a given system.

    Generates a 2D scatter plot of the first two PCA components of the
    latent vectors, colored by total energy.

    Args:
        model: Trained PhysicsNet model.
        trajectories: Trajectory arrays.
        system: Physical system name.
        output_dir: Directory to save the figure.
        window: Sliding window size.
        max_trajs: Maximum number of trajectories to plot.
    """
    from sklearn.decomposition import PCA

    output_dir.mkdir(parents=True, exist_ok=True)
    _, state_cols, energy_col = get_system_dims(system)

    all_latents = []
    all_energies = []

    for traj in trajectories[:max_trajs]:
        if traj.shape[0] < window:
            continue
        states = traj[:, state_cols].astype(np.float32)
        latents = model.get_latent(states, window=window)
        energies = traj[window - 1:window - 1 + latents.shape[0], energy_col]
        all_latents.append(latents)
        all_energies.append(energies)

    if not all_latents:
        return

    latents = np.concatenate(all_latents, axis=0)
    energies = np.concatenate(all_energies, axis=0)

    # PCA to 2D
    pca = PCA(n_components=2)
    latent_2d = pca.fit_transform(latents)

    fig, ax = plt.subplots(figsize=(8, 6))
    scatter = ax.scatter(
        latent_2d[:, 0], latent_2d[:, 1],
        c=energies, cmap="coolwarm", alpha=0.5, s=5,
    )
    plt.colorbar(scatter, label="Total Energy")
    ax.set_xlabel(f"PCA 1 ({pca.explained_variance_ratio_[0]:.1%} var)")
    ax.set_ylabel(f"PCA 2 ({pca.explained_variance_ratio_[1]:.1%} var)")
    ax.set_title(f"Latent Space — {system.capitalize()}")
    ax.grid(True, alpha=0.3)

    fig_path = output_dir / f"{system}_latent_space.png"
    plt.savefig(fig_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"[PLOT] Latent space plot saved to {fig_path}")


def main() -> None:
    """Entry point for the evaluation script.

    Evaluates trained models across specified systems, prints a summary
    table, and generates comparison plots.
    """
    parser = argparse.ArgumentParser(
        description="Evaluate trained PhysicsNet models across physical systems."
    )
    parser.add_argument(
        "--systems",
        nargs="+",
        default=["projectile", "spring", "orbit"],
        choices=["projectile", "spring", "orbit"],
        help="Systems to evaluate.",
    )
    parser.add_argument("--n_test", type=int, default=200, help="Number of test trajectories per system.")
    parser.add_argument("--window", type=int, default=20, help="Sliding window size.")
    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
    args = parser.parse_args()

    set_seed(args.seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    data_dir = PROJECT_ROOT / "data"
    checkpoint_dir = PROJECT_ROOT / "checkpoints"
    plots_dir = PROJECT_ROOT / "plots"
    results_dir = PROJECT_ROOT / "results"
    results_dir.mkdir(parents=True, exist_ok=True)

    results = {}

    for system in args.systems:
        print(f"\n{'='*60}")
        print(f"  Evaluating: {system.capitalize()}")
        print(f"{'='*60}")

        try:
            model = load_model(system, checkpoint_dir, device=device)
        except FileNotFoundError as e:
            print(f"  [SKIP] {e}")
            continue

        # Load test trajectories (last n_test from the dataset)
        all_trajs = load_trajectories(system, data_dir)
        test_trajs = all_trajs[-args.n_test:]
        print(f"  [DATA] Using {len(test_trajs)} test trajectories")

        # Evaluate
        metrics = evaluate_system(
            model, test_trajs, system,
            window=args.window, device=device,
        )
        results[system] = metrics

        print(f"  MSE:              {metrics['mse']:.6f}")
        print(f"  R²:               {metrics['r2']:.6f}")
        print(f"  Energy Violation:  {metrics['energy_violation']:.6f}")
        print(f"  Latent Std:        {metrics['latent_mean_std']:.6f}")
        print(f"  Samples:           {metrics['n_samples']}")

        # Plot latent space
        try:
            plot_latent_distributions(
                model, test_trajs, system, plots_dir,
                window=args.window, max_trajs=50,
            )
        except Exception as e:
            print(f"  [WARN] Latent plot failed: {e}")

    # Print summary table
    if results:
        print(f"\n{'='*80}")
        print(f"  EVALUATION SUMMARY")
        print(f"{'='*80}")
        print(
            f"  {'System':<15s} | {'MSE':>12s} | {'Energy Viol.':>14s} | "
            f"{'R²':>8s} | {'Latent Std':>12s} | {'Samples':>8s}"
        )
        print(f"  {'-'*15}-+-{'-'*12}-+-{'-'*14}-+-{'-'*8}-+-{'-'*12}-+-{'-'*8}")
        for system, m in results.items():
            print(
                f"  {system:<15s} | {m['mse']:>12.6f} | {m['energy_violation']:>14.6f} | "
                f"{m['r2']:>8.4f} | {m['latent_mean_std']:>12.6f} | {m['n_samples']:>8d}"
            )
        print(f"{'='*80}\n")

        # Generate comparison plots
        plot_comparison(results, plots_dir)

    # Save results
    metrics_path = results_dir / "evaluation_metrics.json"
    with open(metrics_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"[RESULTS] Evaluation metrics saved to {metrics_path}")


if __name__ == "__main__":
    main()