main.py

"""
Standalone test / visualisation script.

Loads two saved models (NNModel and NNPhysicsModel), evaluates both on the
held-out test set, and displays their trajectory predictions side by side
against the ground truth.

Usage (from src/):
    python test.py --nn_model models/nn_....pt --physics_model models/hybrid_....pt
    python test.py   # auto-picks the two most-recent .pt files
"""

import os
import sys
import glob
import copy
import argparse
import numpy as np
import matplotlib.pyplot as plt
import torch

# Make sure imports resolve when running from src/
sys.path.insert(0, os.path.dirname(__file__))

from lib.models import PushNetFactory
from lib.physics import PushPhysics
from helpers.utils import load_data, prepare_dataloaders_split
from helpers.config import load_config
from colorama import init, Fore, Style

init()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def print_header(text: str):
    print(f"\n{Fore.CYAN}{Style.BRIGHT}{text}{Style.RESET_ALL}")

def print_success(text: str):
    print(f"{Fore.GREEN}{text}{Style.RESET_ALL}")


def print_info(text: str):
    print(f"{Fore.YELLOW}{text}{Style.RESET_ALL}")


# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------

def find_latest_models(models_dir: str):
    """Return the two most-recently-saved .pt files (oldest-of-two, newest)."""
    files = sorted(glob.glob(os.path.join(models_dir, "*.pt")), key=os.path.getmtime)
    if len(files) < 2:
        raise FileNotFoundError(
            f"Need at least 2 .pt files in '{models_dir}'. "
            "Run main.py with both NNModel and NNPhysicsModel configs first."
        )
    return files[-2], files[-1]


def load_model(model_path: str, model_cfg: dict, device: torch.device):
    model = PushNetFactory.create(model_cfg).to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()
    return model


# ---------------------------------------------------------------------------
# Plot 1 – Training / Validation Loss Curves
# ---------------------------------------------------------------------------

def plot_loss_curves():
    train_path = "results/train_losses.npy"
    val_path   = "results/val_losses.npy"
    if not (os.path.exists(train_path) and os.path.exists(val_path)):
        print_info("Loss history not found in results/. Skipping loss curve plot.")
        return

    print_header("Plot 1: Training vs Validation Loss Curves")
    train_losses = np.load(train_path)
    val_losses   = np.load(val_path)

    plt.figure(figsize=(10, 5))
    plt.plot(train_losses, label="Train Loss",      color="blue")
    plt.plot(val_losses,   label="Validation Loss", color="orange")
    plt.xlabel("Epoch")
    plt.ylabel("MSE Loss")
    plt.title("Training vs Validation Loss")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    os.makedirs("results", exist_ok=True)
    plt.savefig("results/training_curves.png", dpi=150)
    plt.show()
    print_success("Saved → results/training_curves.png")


# ---------------------------------------------------------------------------
# Plot 2 – Side-by-side trajectory comparison
# ---------------------------------------------------------------------------

def _draw_trajectory(ax, gt_x, gt_y, gt_th, pr_x, pr_y, pr_th,
                     pred_label: str, pred_color: str, arrow_len: float = 0.015):
    """Draw GT and model predictions on a single Axes."""
    ax.plot(gt_x, gt_y, "o-", color="royalblue", label="Ground Truth", zorder=2)
    for i, (xi, yi, ti) in enumerate(zip(gt_x, gt_y, gt_th)):
        ax.annotate(
            "",
            xy=(xi + arrow_len * np.cos(ti), yi + arrow_len * np.sin(ti)),
            xytext=(xi, yi),
            arrowprops=dict(arrowstyle="-|>", color="royalblue", lw=1.5),
            zorder=3,
        )
        ax.text(xi, yi + 0.003, str(i), fontsize=7, ha="center", color="royalblue")

    ax.plot(pr_x, pr_y, "s--", color=pred_color, label=pred_label, zorder=2)
    for i, (xi, yi, ti) in enumerate(zip(pr_x, pr_y, pr_th)):
        ax.annotate(
            "",
            xy=(xi + arrow_len * np.cos(ti), yi + arrow_len * np.sin(ti)),
            xytext=(xi, yi),
            arrowprops=dict(arrowstyle="-|>", color=pred_color, lw=1.5),
            zorder=3,
        )
        ax.text(xi, yi - 0.005, str(i), fontsize=7, ha="center", color=pred_color)

    ax.set_xlabel("X position (m)")
    ax.set_ylabel("Y position (m)")
    ax.legend()
    ax.set_aspect("equal")
    ax.grid(True, linestyle="--", alpha=0.5)


def plot_trajectories_side_by_side(nn_model, physics_model, phy_preds_all: np.ndarray,
                                   x_test: np.ndarray, y_test: np.ndarray,
                                   device: torch.device, num_pushes: int = 10):
    print_header(f"Plot 2: Side-by-side Trajectory Comparison ({num_pushes} pushes, test set)")

    x_tensor = torch.FloatTensor(x_test).to(device)
    indices  = np.arange(num_pushes)

    with torch.no_grad():
        nn_preds = nn_model(x_tensor).cpu().numpy()
        ph_preds = physics_model(x_tensor).cpu().numpy()

    gt_x,  gt_y,  gt_th  = y_test[indices, 0],           y_test[indices, 1],           y_test[indices, 2]
    nn_x,  nn_y,  nn_th  = nn_preds[indices, 0],         nn_preds[indices, 1],         nn_preds[indices, 2]
    ph_x,  ph_y,  ph_th  = ph_preds[indices, 0],         ph_preds[indices, 1],         ph_preds[indices, 2]
    py_x,  py_y,  py_th  = phy_preds_all[indices, 0],    phy_preds_all[indices, 1],    phy_preds_all[indices, 2]

    fig, (ax_nn, ax_ph, ax_py) = plt.subplots(1, 3, figsize=(22, 7))

    _draw_trajectory(ax_nn, gt_x, gt_y, gt_th, nn_x, nn_y, nn_th,
                     pred_label="NNModel", pred_color="tomato")
    ax_nn.set_title("Ground Truth vs NNModel\nArrows = orientation (θ)")

    _draw_trajectory(ax_ph, gt_x, gt_y, gt_th, ph_x, ph_y, ph_th,
                     pred_label="NNPhysicsModel", pred_color="seagreen")
    ax_ph.set_title("Ground Truth vs NNPhysicsModel\nArrows = orientation (θ)")

    _draw_trajectory(ax_py, gt_x, gt_y, gt_th, py_x, py_y, py_th,
                     pred_label="PhysicsModel", pred_color="darkorange")
    ax_py.set_title("Ground Truth vs PhysicsModel (analytical)\nArrows = orientation (θ)")

    fig.suptitle(f"Push Trajectory Comparison — {num_pushes} test-set pushes", fontsize=13)
    plt.tight_layout()
    os.makedirs("results", exist_ok=True)
    plt.savefig("results/trajectory_comparison.png", dpi=150)
    plt.show()
    print_success("Saved → results/trajectory_comparison.png")

    # Save each model's trajectory as a separate figure
    individual = [
        ("NNModel",        "tomato",     "NNModel",                  nn_x, nn_y, nn_th, "results/trajectory_nn.png"),
        ("NNPhysicsModel", "seagreen",   "NNPhysicsModel",           ph_x, ph_y, ph_th, "results/trajectory_hybrid.png"),
        ("PhysicsModel",   "darkorange", "PhysicsModel (analytical)", py_x, py_y, py_th, "results/trajectory_physics.png"),
    ]
    for label, pred_color, pred_label, px, py, pth, fname in individual:
        fig_s, ax_s = plt.subplots(figsize=(9, 7))
        _draw_trajectory(ax_s, gt_x, gt_y, gt_th, px, py, pth,
                         pred_label=pred_label, pred_color=pred_color)
        ax_s.set_title(f"Ground Truth vs {label}\nArrows = orientation (θ)")
        fig_s.suptitle(f"{label} — {num_pushes} test-set pushes", fontsize=12)
        plt.tight_layout()
        plt.savefig(fname, dpi=150)
        plt.show()
        print_success(f"Saved → {fname}")


# ---------------------------------------------------------------------------
# Evaluation helper
# ---------------------------------------------------------------------------

def evaluate(model, x_test: np.ndarray, y_test: np.ndarray,
             device: torch.device, label: str):
    x_t = torch.FloatTensor(x_test).to(device)
    y_t = torch.FloatTensor(y_test).to(device)
    with torch.no_grad():
        preds   = model(x_t)
        loss    = model.loss(preds, y_t)
        metrics = model.accuracy(preds, y_t)
    print_success(f"[{label}] Test MSE Loss    : {loss.item():.4f}")
    print_info(   f"[{label}] Position Error   : {metrics['position_error']:.4f} m")
    print_info(   f"[{label}] Angle Error (rad): {metrics['angle_error_rad']:.4f}")


def evaluate_physics(physics: PushPhysics, x_test: np.ndarray, y_test: np.ndarray):
    """Evaluate the plain analytical physics model (no learned weights)."""
    x_t = torch.FloatTensor(x_test)
    y_t = torch.FloatTensor(y_test)
    with torch.no_grad():
        preds      = physics.compute_motion(x_t)
        mse        = torch.mean((preds - y_t) ** 2).item()
        pos_err    = torch.norm(preds[:, :2] - y_t[:, :2], dim=1).mean().item()
        angle_diff = preds[:, 2] - y_t[:, 2]
        angle_err  = torch.abs(
            torch.atan2(torch.sin(angle_diff), torch.cos(angle_diff))
        ).mean().item()
    print_success(f"[PhysicsModel] Test MSE Loss    : {mse:.4f}")
    print_info(   f"[PhysicsModel] Position Error   : {pos_err:.4f} m")
    print_info(   f"[PhysicsModel] Angle Error (rad): {angle_err:.4f}")
    return preds.numpy()


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args():
    parser = argparse.ArgumentParser(description="Compare NNModel vs NNPhysicsModel on test set")
    parser.add_argument("--config",        type=str, default=None,
                        help="Path to config YAML (default: config/default.yaml)")
    parser.add_argument("--nn_model",      type=str, default="/home/prithvi/MLR/Assignment_2_Fillable/src/models/nn_20260305_155807.pt",
                        help="Path to the NNModel .pt checkpoint")
    parser.add_argument("--physics_model", type=str, default="/home/prithvi/MLR/Assignment_2_Fillable/src/models/hybrid_20260305_160615.pt",
                        help="Path to the NNPhysicsModel .pt checkpoint")
    parser.add_argument("--num_pushes",    type=int, default=10,
                        help="Number of pushes to show in the trajectory plot")
    return parser.parse_args()


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    args = parse_args()

    # ── Config & device ───────────────────────────────────────────────────
    config = load_config(args.config)
    device = config.get_device()
    print_info(f"Using device: {device}")

    # ── Data → test split ─────────────────────────────────────────────────
    print_header("Loading Data")
    x_data, y_data = load_data(config)
    print_info(f"Full data shapes: x={x_data.shape}, y={y_data.shape}")

    _, _, test_loader = prepare_dataloaders_split(
        x_data, y_data, config, val_split=0.15, test_split=0.15
    )
    x_test = np.concatenate([xb.numpy() for xb, _ in test_loader])
    y_test = np.concatenate([yb.numpy() for _, yb in test_loader])
    print_info(f"Test split : x={x_test.shape}, y={y_test.shape}")

    # ── Model paths ───────────────────────────────────────────────────────
    nn_path, ph_path = args.nn_model, args.physics_model
    if nn_path is None or ph_path is None:
        auto_a, auto_b = find_latest_models("models")
        nn_path = nn_path or auto_a
        ph_path = ph_path or auto_b
    print_info(f"NNModel path        : {nn_path}")
    print_info(f"NNPhysicsModel path : {ph_path}")

    # ── Build per-model configs (force correct type for each) ─────────────
    nn_cfg = copy.deepcopy(config.model)
    ph_cfg = copy.deepcopy(config.model)
    nn_cfg["network"]["type"] = "NNModel"
    ph_cfg["network"]["type"] = "NNPhysicsModel"

    # ── Load models ───────────────────────────────────────────────────────
    nn_model = load_model(nn_path, nn_cfg, device)
    ph_model = load_model(ph_path, ph_cfg, device)
    print_success("Both models loaded successfully")

    # ── Evaluate on test set ──────────────────────────────────────────────
    print_header("Test Set Evaluation")
    evaluate(nn_model, x_test, y_test, device, "NNModel")
    evaluate(ph_model, x_test, y_test, device, "NNPhysicsModel")

    # Plain analytical physics model — instantiated from config, no checkpoint needed
    physics = PushPhysics.from_config(config.model["physics"])
    phy_preds_all = evaluate_physics(physics, x_test, y_test)

    # ── Plot 1: Loss curves ───────────────────────────────────────────────
    plot_loss_curves()

    # ── Plot 2: Side-by-side trajectories ────────────────────────────────
    plot_trajectories_side_by_side(
        nn_model, ph_model, phy_preds_all, x_test, y_test, device,
        num_pushes=args.num_pushes,
    )


if __name__ == "__main__":
    main()