run_traj_eval.py

#!/usr/bin/env python3
"""
Automated dataset evaluation runner for main.py.

Runs SLAM on selected dataset configs multiple times, archives trajectory
outputs under output/traj_eval/<dataset>/run_<NN>/, and evaluates
trajectory_optimized.txt:

  - euroc / tum: evo_ape (TUM GT, trim head/tail first)
  - kitti: eval/kitti/eval_kitti.py (pose + timestamps under eval/kitti/)
"""

from __future__ import annotations

import argparse
import csv
import json
import re
import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Iterable

import yaml

PROJECT_ROOT = Path(__file__).resolve().parent
CONFIG_ROOT = PROJECT_ROOT / "config"
EVAL_ROOT = PROJECT_ROOT / "eval"
KITTI_EVAL_ROOT = EVAL_ROOT / "kitti"
KITTI_POSE_ROOT = KITTI_EVAL_ROOT / "pose"
KITTI_TIMESTAMPS_ROOT = KITTI_EVAL_ROOT / "timestamps"
EVAL_KITTI_SCRIPT = KITTI_EVAL_ROOT / "eval_kitti.py"
MAIN_SCRIPT = PROJECT_ROOT / "main.py"
TRAJ_EVAL_ROOT = PROJECT_ROOT / "output" / "traj_eval"

KITTI_CONFIG_DIR = CONFIG_ROOT / "kitti"

TRAJECTORY_FILES = (
    "trajectory_odometry.txt",
    "trajectory_vggt.txt",
    "trajectory_optimized.txt",
    "trajectory_vggt_optimized.txt",
)

EVAL_TRAJECTORY = "trajectory_optimized.txt"
ODOM_TRAJECTORY = "trajectory_odometry.txt"
KITTI_EVAL_TRAJECTORIES = (
    (ODOM_TRAJECTORY, "odom"),
    (EVAL_TRAJECTORY, "optimized"),
)
TRIM_HEAD_DEFAULT = 10
TRIM_TAIL_DEFAULT = 10


@dataclass(frozen=True)
class DatasetConfig:
    name: str
    config_path: Path
    dataset_type: str
    dataset_path: str

    @property
    def rel_config(self) -> str:
        return str(self.config_path.relative_to(PROJECT_ROOT))


def discover_datasets(config_root: Path = CONFIG_ROOT) -> list[DatasetConfig]:
    datasets: list[DatasetConfig] = []
    for config_path in sorted(config_root.rglob("*.yaml")):
        if "submodules" in config_path.parts:
            continue
        try:
            with open(config_path, "r", encoding="utf-8") as f:
                cfg = yaml.safe_load(f) or {}
        except (OSError, yaml.YAMLError):
            continue

        dataset_type = str(cfg.get("dataset_type", "unknown"))
        dataset_path = str(cfg.get("dataset_path", ""))
        stem = config_path.stem
        if stem.startswith("config_"):
            name = stem[len("config_") :]
        else:
            name = stem

        datasets.append(
            DatasetConfig(
                name=name,
                config_path=config_path.resolve(),
                dataset_type=dataset_type,
                dataset_path=dataset_path,
            )
        )
    return datasets


def resolve_datasets(
    datasets: list[DatasetConfig],
    *,
    names: Iterable[str] | None,
    types: Iterable[str] | None,
    configs: Iterable[Path] | None,
    all_datasets: bool,
) -> list[DatasetConfig]:
    groups: list[list[DatasetConfig]] = []
    type_set = {t.lower() for t in types} if types else None

    if configs:
        by_path = {d.config_path.resolve(): d for d in datasets}
        config_items: list[DatasetConfig] = []
        for raw in configs:
            path = Path(raw).expanduser().resolve()
            if path in by_path:
                config_items.append(by_path[path])
            else:
                config_items.append(
                    DatasetConfig(
                        name=path.stem,
                        config_path=path,
                        dataset_type="custom",
                        dataset_path="",
                    )
                )
        groups.append(config_items)

    if names:
        by_name = {d.name: d for d in datasets}
        name_items: list[DatasetConfig] = []
        for name in names:
            if name not in by_name:
                raise SystemExit(f"Unknown dataset name: {name}")
            dataset = by_name[name]
            if type_set is not None and dataset.dataset_type.lower() not in type_set:
                raise SystemExit(
                    f"Dataset {name} has type {dataset.dataset_type}, "
                    f"not in requested types: {sorted(type_set)}"
                )
            if (
                type_set is not None
                and "kitti" in type_set
                and dataset.dataset_type.lower() == "kitti"
                and not is_kitti_sequence_config(dataset)
            ):
                raise SystemExit(
                    f"Dataset {name} is not a standard KITTI sequence config "
                    f"(expected config/kitti/NN.yaml)"
                )
            name_items.append(dataset)
        groups.append(name_items)
    elif types:
        type_items: list[DatasetConfig] = []
        for dataset in datasets:
            if dataset.dataset_type.lower() not in type_set:
                continue
            if "kitti" in type_set and dataset.dataset_type.lower() == "kitti":
                if not is_kitti_sequence_config(dataset):
                    continue
            type_items.append(dataset)
        groups.append(type_items)

    if all_datasets:
        groups.append(list(datasets))

    if not groups:
        return []

    if len(groups) == 1:
        return groups[0]

    common_paths = {str(d.config_path) for d in groups[0]}
    for group in groups[1:]:
        common_paths &= {str(d.config_path) for d in group}

    selected: list[DatasetConfig] = []
    seen: set[str] = set()
    for group in groups:
        for dataset in group:
            key = str(dataset.config_path)
            if key not in common_paths or key in seen:
                continue
            seen.add(key)
            selected.append(dataset)
    return selected


def print_dataset_table(datasets: list[DatasetConfig]) -> None:
    print(f"{'NAME':<28} {'TYPE':<12} {'GT':<8} {'SEQ':<5} CONFIG")
    print("-" * 104)
    for d in datasets:
        gt_flag = "yes" if resolve_ground_truth(d) else "no"
        seq = (resolve_kitti_sequence_id(d) if d.dataset_type.lower() == "kitti" else "-") or "-"
        print(f"{d.name:<28} {d.dataset_type:<12} {gt_flag:<8} {seq:<5} {d.rel_config}")


def is_kitti_sequence_config(dataset: DatasetConfig) -> bool:
    """True for standard KITTI configs like config/kitti/00.yaml .. 10.yaml."""
    if dataset.dataset_type.lower() != "kitti":
        return False
    if KITTI_CONFIG_DIR not in dataset.config_path.parents:
        return False
    return bool(re.fullmatch(r"\d{2}", dataset.config_path.stem))


def resolve_kitti_sequence_id(dataset: DatasetConfig) -> str | None:
    """Read KITTI sequence id from config filename (e.g. 03.yaml -> 03)."""
    if not is_kitti_sequence_config(dataset):
        return None
    return dataset.config_path.stem


def resolve_kitti_timestamps(seq_id: str) -> Path | None:
    path = KITTI_TIMESTAMPS_ROOT / f"timestamps_{seq_id}.txt"
    return path.resolve() if path.is_file() else None


def resolve_kitti_eval_assets(
    dataset: DatasetConfig,
    *,
    seq_override: str | None = None,
    gt_override: Path | None = None,
) -> tuple[str, Path, Path] | None:
    """Return (seq_id, pose_gt, timestamps) when KITTI evaluation assets are available."""
    seq_id = seq_override or resolve_kitti_sequence_id(dataset)
    if seq_id is None:
        return None

    gt_file = gt_override if gt_override is not None else (KITTI_POSE_ROOT / f"{seq_id}.txt")
    if not gt_file.is_file():
        return None

    ts_file = resolve_kitti_timestamps(seq_id)
    if ts_file is None:
        return None

    return seq_id, gt_file.resolve(), ts_file


def resolve_ground_truth(dataset: DatasetConfig) -> Path | None:
    """Map dataset name/type to a ground-truth file under eval/."""
    name = dataset.name
    dataset_type = dataset.dataset_type.lower()

    if dataset_type == "euroc":
        suffix = name
        if suffix.startswith("euroc_"):
            suffix = suffix[len("euroc_") :]
        suffix = re.sub(r"_vggt$", "", suffix, flags=re.IGNORECASE)
        candidates = [
            EVAL_ROOT / "euroc" / f"euroc_ground_truth_{suffix}.txt",
            EVAL_ROOT / "euroc" / f"euroc_ground_truth_{suffix.upper()}.txt",
        ]
    elif dataset_type == "tum":
        suffix = name[len("tum_") :] if name.startswith("tum_") else name
        candidates = [EVAL_ROOT / "tum" / f"groundtruth_{suffix}.txt"]
    elif dataset_type == "kitti":
        assets = resolve_kitti_eval_assets(dataset)
        return assets[1] if assets else None
    else:
        return None

    for path in candidates:
        if path.is_file():
            return path.resolve()
    return None


def read_tum_pose_lines(path: Path) -> tuple[str | None, list[str]]:
    header = None
    poses: list[str] = []
    with open(path, "r", encoding="utf-8") as f:
        for raw in f:
            line = raw.strip()
            if not line:
                continue
            if line.startswith("#"):
                if "timestamp" in line.lower():
                    header = line
                continue
            poses.append(line)
    return header, poses


def trim_trajectory_tum(
    src: Path,
    dst: Path,
    *,
    trim_head: int,
    trim_tail: int,
) -> tuple[int, int]:
    """Drop the first/last N pose rows and write a TUM trajectory file."""
    header, poses = read_tum_pose_lines(src)
    original_count = len(poses)
    if original_count <= trim_head + trim_tail:
        trimmed = []
    else:
        trimmed = poses[trim_head : original_count - trim_tail]

    dst.parent.mkdir(parents=True, exist_ok=True)
    with open(dst, "w", encoding="utf-8") as f:
        f.write((header or "# timestamp tx ty tz qx qy qz qw") + "\n")
        for line in trimmed:
            f.write(line + "\n")
    return original_count, len(trimmed)


@dataclass
class EvoResult:
    rmse: float | None
    mean: float | None
    median: float | None
    max: float | None
    min: float | None
    std: float | None
    scale: float | None
    pose_count: int
    error: str | None = None

    def as_dict(self) -> dict[str, object]:
        return {
            "rmse": self.rmse,
            "mean": self.mean,
            "median": self.median,
            "max": self.max,
            "min": self.min,
            "std": self.std,
            "scale": self.scale,
            "pose_count": self.pose_count,
            "error": self.error,
        }


def _parse_evo_metric(output: str, key: str) -> float | None:
    match = re.search(rf"^\s*{re.escape(key)}\s+([\d.eE+-]+)\s*$", output, re.MULTILINE)
    return float(match.group(1)) if match else None


def run_evo_ape(
    gt_file: Path,
    est_file: Path,
    *,
    align: bool = True,
    correct_scale: bool = True,
) -> EvoResult:
    pose_count = len(read_tum_pose_lines(est_file)[1])
    if pose_count == 0:
        return EvoResult(
            rmse=None,
            mean=None,
            median=None,
            max=None,
            min=None,
            std=None,
            scale=None,
            pose_count=0,
            error="trimmed trajectory is empty",
        )

    cmd = ["evo_ape", "tum", str(gt_file), str(est_file)]
    if align:
        cmd.append("-a")
    if correct_scale:
        cmd.append("-s")

    try:
        proc = subprocess.run(cmd, capture_output=True, text=True)
    except FileNotFoundError:
        return EvoResult(
            rmse=None,
            mean=None,
            median=None,
            max=None,
            min=None,
            std=None,
            scale=None,
            pose_count=pose_count,
            error="evo_ape not found in PATH",
        )

    output = proc.stdout + "\n" + proc.stderr
    if proc.returncode != 0:
        err = proc.stderr.strip() or proc.stdout.strip() or f"exit code {proc.returncode}"
        return EvoResult(
            rmse=None,
            mean=None,
            median=None,
            max=None,
            min=None,
            std=None,
            scale=None,
            pose_count=pose_count,
            error=err,
        )

    scale_match = re.search(r"Scale correction\s*:\s*([\d.eE+-]+)", output, re.IGNORECASE)
    if not scale_match:
        scale_match = re.search(r"scale\s+s\s*=\s*([\d.eE+-]+)", output, re.IGNORECASE)

    return EvoResult(
        rmse=_parse_evo_metric(output, "rmse"),
        mean=_parse_evo_metric(output, "mean"),
        median=_parse_evo_metric(output, "median"),
        max=_parse_evo_metric(output, "max"),
        min=_parse_evo_metric(output, "min"),
        std=_parse_evo_metric(output, "std"),
        scale=float(scale_match.group(1)) if scale_match else None,
        pose_count=pose_count,
    )


def _parse_kitti_summary(summary_path: Path) -> dict[str, object]:
    parsed: dict[str, object] = {}
    with open(summary_path, "r", encoding="utf-8") as f:
        for raw in f:
            line = raw.strip()
            if not line:
                continue
            key, _, value = line.partition(" ")
            if not key:
                continue
            try:
                parsed[key] = float(value)
            except ValueError:
                parsed[key] = value
    return parsed


def python_can_import(python_exec: str, module: str) -> bool:
    try:
        proc = subprocess.run(
            [python_exec, "-c", f"import {module}"],
            capture_output=True,
            text=True,
        )
    except OSError:
        return False
    return proc.returncode == 0


def evaluate_kitti_run_trajectory(
    run_dir: Path,
    dataset: DatasetConfig,
    *,
    python_exec: str,
    trim_head: int,
    trim_tail: int,
    trajectory: str = EVAL_TRAJECTORY,
    eval_tag: str = "optimized",
    seq_override: str | None = None,
    gt_override: Path | None = None,
) -> dict[str, object] | None:
    traj_src = run_dir / trajectory
    if not traj_src.is_file():
        return None

    if not EVAL_KITTI_SCRIPT.is_file():
        return {
            "status": "skipped",
            "reason": f"eval_kitti.py not found: {EVAL_KITTI_SCRIPT}",
            "trajectory": trajectory,
            "trim_head": trim_head,
            "trim_tail": trim_tail,
        }

    assets = resolve_kitti_eval_assets(
        dataset,
        seq_override=seq_override,
        gt_override=gt_override,
    )
    if assets is None:
        seq_id = seq_override or resolve_kitti_sequence_id(dataset)
        if seq_id is None:
            reason = "cannot infer KITTI sequence id from config name/path"
        elif gt_override is not None and not gt_override.is_file():
            reason = f"KITTI GT override not found: {gt_override}"
        elif not (KITTI_POSE_ROOT / f"{seq_id}.txt").is_file():
            reason = f"KITTI GT pose not found: {KITTI_POSE_ROOT / f'{seq_id}.txt'}"
        else:
            reason = (
                f"KITTI timestamps not found: "
                f"{KITTI_TIMESTAMPS_ROOT / f'timestamps_{seq_id}.txt'}"
            )
        return {
            "status": "skipped",
            "reason": reason,
            "trajectory": trajectory,
            "trim_head": trim_head,
            "trim_tail": trim_tail,
            "sequence_id": seq_id,
        }

    seq_id, gt_file, ts_file = assets
    out_dir = run_dir / ("kitti_eval" if eval_tag == "optimized" else f"kitti_eval_{eval_tag}")
    out_dir.mkdir(parents=True, exist_ok=True)
    if not python_can_import(python_exec, "evo"):
        return {
            "status": "failed",
            "reason": (
                f"Python interpreter cannot import evo: {python_exec}\n"
                "Activate your conda env (e.g. ivggt) or pass "
                "--python /path/to/ivggt/bin/python"
            ),
            "trajectory": trajectory,
            "sequence_id": seq_id,
            "ground_truth": str(gt_file.relative_to(PROJECT_ROOT)),
            "timestamps": str(ts_file.relative_to(PROJECT_ROOT)),
            "trim_head": trim_head,
            "trim_tail": trim_tail,
        }

    cmd = [
        python_exec,
        str(EVAL_KITTI_SCRIPT),
        str(gt_file),
        str(traj_src),
        str(ts_file),
        "--skip-head",
        str(trim_head),
        "--skip-tail",
        str(trim_tail),
        "--out-dir",
        str(out_dir),
    ]

    try:
        proc = subprocess.run(cmd, cwd=PROJECT_ROOT, capture_output=True, text=True)
    except OSError as exc:
        return {
            "status": "failed",
            "reason": str(exc),
            "trajectory": trajectory,
            "sequence_id": seq_id,
            "ground_truth": str(gt_file.relative_to(PROJECT_ROOT)),
            "timestamps": str(ts_file.relative_to(PROJECT_ROOT)),
            "trim_head": trim_head,
            "trim_tail": trim_tail,
            "eval_command": cmd,
        }

    stdout_path = out_dir / "eval_kitti_stdout.txt"
    stderr_path = out_dir / "eval_kitti_stderr.txt"
    stdout_path.write_text(proc.stdout, encoding="utf-8")
    stderr_path.write_text(proc.stderr, encoding="utf-8")

    summary_path = out_dir / "summary.txt"
    if proc.returncode != 0 or not summary_path.is_file():
        err = proc.stderr.strip() or proc.stdout.strip() or f"exit code {proc.returncode}"
        return {
            "status": "failed",
            "reason": err,
            "trajectory": trajectory,
            "sequence_id": seq_id,
            "ground_truth": str(gt_file.relative_to(PROJECT_ROOT)),
            "timestamps": str(ts_file.relative_to(PROJECT_ROOT)),
            "trim_head": trim_head,
            "trim_tail": trim_tail,
            "eval_command": cmd,
            "stdout": str(stdout_path.relative_to(PROJECT_ROOT)),
            "stderr": str(stderr_path.relative_to(PROJECT_ROOT)),
        }

    summary = _parse_kitti_summary(summary_path)
    result: dict[str, object] = {
        "status": "ok",
        "evaluator": "eval_kitti.py",
        "trajectory": trajectory,
        "trajectory_label": eval_tag,
        "sequence_id": seq_id,
        "ground_truth": str(gt_file.relative_to(PROJECT_ROOT)),
        "timestamps": str(ts_file.relative_to(PROJECT_ROOT)),
        "trim_head": trim_head,
        "trim_tail": trim_tail,
        "matched_pose_count": summary.get("matched_poses"),
        "t_rel_percent": summary.get("t_rel_percent"),
        "r_rel_deg_per_100m": summary.get("r_rel_deg_per_100m"),
        "sim3_scale": summary.get("sim3_scale"),
        "metrics": {
            "t_rel_percent": summary.get("t_rel_percent"),
            "r_rel_deg_per_100m": summary.get("r_rel_deg_per_100m"),
            "matched_poses": summary.get("matched_poses"),
            "sim3_scale": summary.get("sim3_scale"),
        },
        "eval_command": cmd,
        "summary": str(summary_path.relative_to(PROJECT_ROOT)),
        "stdout": str(stdout_path.relative_to(PROJECT_ROOT)),
    }

    result_json = (
        "kitti_result.json" if eval_tag == "optimized" else f"kitti_result_{eval_tag}.json"
    )
    with open(run_dir / result_json, "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2, ensure_ascii=False)
    return result


def evaluate_run_trajectory(
    run_dir: Path,
    dataset: DatasetConfig,
    *,
    python_exec: str,
    trim_head: int,
    trim_tail: int,
    align: bool,
    correct_scale: bool,
    trajectory: str = EVAL_TRAJECTORY,
    eval_tag: str = "optimized",
    gt_override: Path | None = None,
    kitti_seq_override: str | None = None,
) -> dict[str, object] | None:
    traj_src = run_dir / trajectory
    if not traj_src.is_file():
        return None

    if dataset.dataset_type.lower() == "kitti":
        if not is_kitti_sequence_config(dataset):
            return {
                "status": "skipped",
                "reason": (
                    "KITTI eval requires config/kitti/NN.yaml (e.g. 00.yaml, 03.yaml)"
                ),
                "trajectory": trajectory,
                "trim_head": trim_head,
                "trim_tail": trim_tail,
            }
        return evaluate_kitti_run_trajectory(
            run_dir,
            dataset,
            python_exec=python_exec,
            trim_head=trim_head,
            trim_tail=trim_tail,
            trajectory=trajectory,
            eval_tag=eval_tag,
            seq_override=kitti_seq_override,
            gt_override=gt_override,
        )

    gt_file = gt_override or resolve_ground_truth(dataset)
    if gt_file is None:
        return {
            "status": "skipped",
            "reason": "no ground-truth file found",
            "trajectory": EVAL_TRAJECTORY,
            "trim_head": trim_head,
            "trim_tail": trim_tail,
        }

    trimmed_path = run_dir / "trajectory_optimized_trimmed.txt"
    original_count, trimmed_count = trim_trajectory_tum(
        traj_src,
        trimmed_path,
        trim_head=trim_head,
        trim_tail=trim_tail,
    )

    evo = run_evo_ape(
        gt_file,
        trimmed_path,
        align=align,
        correct_scale=correct_scale,
    )
    result: dict[str, object] = {
        "status": "ok" if evo.error is None else "failed",
        "ground_truth": str(gt_file.relative_to(PROJECT_ROOT)),
        "trajectory": EVAL_TRAJECTORY,
        "trimmed_trajectory": trimmed_path.name,
        "trim_head": trim_head,
        "trim_tail": trim_tail,
        "original_pose_count": original_count,
        "trimmed_pose_count": trimmed_count,
        "metrics": evo.as_dict(),
        "evo_command": [
            "evo_ape",
            "tum",
            str(gt_file),
            str(trimmed_path),
            *([] if not align else ["-a"]),
            *([] if not correct_scale else ["-s"]),
        ],
    }

    with open(run_dir / "evo_result.json", "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2, ensure_ascii=False)
    return result


def append_evo_summary_row(summary_path: Path, row: dict[str, object]) -> None:
    summary_path.parent.mkdir(parents=True, exist_ok=True)
    file_exists = summary_path.is_file()
    with open(summary_path, "a", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(row.keys()))
        if not file_exists:
            writer.writeheader()
        writer.writerow(row)


def copy_trajectories(run_dir: Path) -> list[str]:
    run_dir.mkdir(parents=True, exist_ok=True)
    copied: list[str] = []
    for filename in TRAJECTORY_FILES:
        src = PROJECT_ROOT / "output" / filename
        if not src.is_file():
            continue
        dst = run_dir / filename
        shutil.copy2(src, dst)
        copied.append(filename)
    return copied


def write_run_meta(
    run_dir: Path,
    *,
    dataset: DatasetConfig,
    run_index: int,
    command: list[str],
    returncode: int,
    duration_sec: float,
    copied_files: list[str],
    eval_result: dict[str, object] | None = None,
) -> None:
    meta = {
        "dataset": dataset.name,
        "dataset_type": dataset.dataset_type,
        "dataset_path": dataset.dataset_path,
        "config": dataset.rel_config,
        "run_index": run_index,
        "timestamp": datetime.now().isoformat(timespec="seconds"),
        "command": command,
        "returncode": returncode,
        "duration_sec": round(duration_sec, 3),
        "copied_trajectories": copied_files,
    }
    if dataset.dataset_type.lower() == "kitti":
        meta["kitti_evaluation"] = eval_result
    else:
        meta["evo_evaluation"] = eval_result
    with open(run_dir / "run_meta.json", "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2, ensure_ascii=False)


def _format_eval_log(eval_result: dict[str, object] | None) -> str:
    if not eval_result:
        return "eval: skipped (no evaluation)"
    if eval_result.get("status") == "skipped":
        return f"eval: skipped ({eval_result.get('reason', 'unknown')})"
    if eval_result.get("status") == "failed":
        return f"eval: failed ({eval_result.get('reason', 'unknown')})"

    traj_label = eval_result.get("trajectory_label")
    prefix = f"[{traj_label}] " if traj_label else ""

    if eval_result.get("evaluator") == "eval_kitti.py":
        metrics = eval_result.get("metrics", {})
        if not isinstance(metrics, dict):
            return "eval: invalid KITTI result"
        return (
            f"{prefix}kitti: "
            f"t_rel={metrics.get('t_rel_percent')} %, "
            f"r_rel={metrics.get('r_rel_deg_per_100m')} deg/100m, "
            f"poses={metrics.get('matched_poses')}, "
            f"seq={eval_result.get('sequence_id')}"
        )

    metrics = eval_result.get("metrics", {})
    if not isinstance(metrics, dict):
        return "eval: invalid result"
    if metrics.get("error"):
        return f"evo: failed ({metrics['error']})"
    rmse = metrics.get("rmse")
    scale = metrics.get("scale")
    pose_count = eval_result.get("trimmed_pose_count")
    return f"evo: RMSE={rmse} m, scale={scale}, poses={pose_count}"


def discover_archived_runs(dataset_name: str) -> list[tuple[int, Path]]:
    dataset_dir = TRAJ_EVAL_ROOT / dataset_name
    if not dataset_dir.is_dir():
        return []
    runs: list[tuple[int, Path]] = []
    for run_dir in sorted(dataset_dir.glob("run_*")):
        if not run_dir.is_dir():
            continue
        try:
            run_idx = int(run_dir.name.split("_", 1)[1])
        except (IndexError, ValueError):
            continue
        runs.append((run_idx, run_dir))
    return sorted(runs, key=lambda item: item[0])


def _mean_kitti_metric(
    samples: list[dict[str, object]],
    trajectory_label: str,
    key: str,
) -> float | None:
    values = [
        float(s[key])
        for s in samples
        if s.get("trajectory_label") == trajectory_label and s.get(key) is not None
    ]
    if not values:
        return None
    return sum(values) / len(values)


def print_kitti_session_summary(samples: list[dict[str, object]]) -> None:
    if not samples:
        return

    print("KITTI RPE averages (this session):")
    for trajectory_file, trajectory_label in KITTI_EVAL_TRAJECTORIES:
        count = sum(
            1
            for s in samples
            if s.get("trajectory_label") == trajectory_label
            and s.get("t_rel_percent") is not None
            and s.get("r_rel_deg_per_100m") is not None
        )
        if count == 0:
            print(f"  {trajectory_label:10s} ({trajectory_file}): no successful eval runs")
            continue
        t_rel = _mean_kitti_metric(samples, trajectory_label, "t_rel_percent")
        r_rel = _mean_kitti_metric(samples, trajectory_label, "r_rel_deg_per_100m")
        print(
            f"  {trajectory_label:10s} ({trajectory_file}): "
            f"t_rel={t_rel:.6f} %, r_rel={r_rel:.6f} deg/100m "
            f"(avg over {count} run(s))"
        )


def append_summary_row(summary_path: Path, row: dict[str, object]) -> None:
    summary_path.parent.mkdir(parents=True, exist_ok=True)
    file_exists = summary_path.is_file()
    with open(summary_path, "a", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(row.keys()))
        if not file_exists:
            writer.writeheader()
        writer.writerow(row)


def build_main_command(
    dataset: DatasetConfig,
    *,
    python_exec: str,
    quiet: bool,
    no_progress: bool,
    no_time_cost: bool,
    extra_args: list[str],
) -> list[str]:
    cmd = [python_exec, str(MAIN_SCRIPT), "--config", str(dataset.config_path)]
    if quiet:
        cmd.append("--quiet")
    if no_progress:
        cmd.append("--no-progress")
    if no_time_cost:
        cmd.append("--no-time-cost")
    cmd.extend(extra_args)
    return cmd


def run_evaluation(args: argparse.Namespace) -> int:
    datasets = discover_datasets()
    if not datasets:
        print("No dataset configs found under config/.", file=sys.stderr)
        return 1

    if args.list:
        print_dataset_table(datasets)
        return 0

    try:
        selected = resolve_datasets(
            datasets,
            names=args.dataset,
            types=args.type,
            configs=[Path(p) for p in args.config] if args.config else None,
            all_datasets=args.all,
        )
    except SystemExit as exc:
        print(exc, file=sys.stderr)
        return 1

    if not selected:
        print(
            "No dataset selected. Use --dataset, --type, --config, or --all.\n"
            "Run with --list to see available datasets.",
            file=sys.stderr,
        )
        return 1

    if args.runs < 1:
        print("--runs must be >= 1", file=sys.stderr)
        return 1

    if args.trim < 0:
        print("--trim must be >= 0", file=sys.stderr)
        return 1

    gt_override = Path(args.gt).expanduser().resolve() if args.gt else None
    if gt_override and not gt_override.is_file():
        print(f"Ground-truth file not found: {gt_override}", file=sys.stderr)
        return 1

    kitti_seq_override = None
    if args.kitti_seq:
        kitti_seq_override = f"{int(args.kitti_seq):02d}"

    TRAJ_EVAL_ROOT.mkdir(parents=True, exist_ok=True)
    summary_path = TRAJ_EVAL_ROOT / "summary.csv"
    evo_summary_path = TRAJ_EVAL_ROOT / "evo_summary.csv"
    kitti_summary_path = TRAJ_EVAL_ROOT / "kitti_summary.csv"
    failures = 0
    kitti_session_metrics: list[dict[str, object]] = []

    mode = "eval-only" if args.eval_only else "run+eval"
    print(f"Selected {len(selected)} dataset(s), {args.runs} run(s) each, mode={mode}.")
    print(f"Python executable: {args.python}")
    print(f"Trajectory archive root: {TRAJ_EVAL_ROOT}")
    if not args.skip_eval:
        print(
            f"Accuracy eval: trim head/tail={args.trim} on {EVAL_TRAJECTORY}; "
            f"euroc/tum -> evo_ape (align={not args.no_align}, scale={not args.no_scale}); "
            f"kitti -> eval/kitti/eval_kitti.py (pose + timestamps under eval/kitti/)"
        )
        if any(is_kitti_sequence_config(d) for d in selected):
            if not python_can_import(args.python, "evo"):
                print(
                    f"[error] {args.python} cannot import evo (required for KITTI eval).\n"
                    "  Activate conda env ivggt, or run e.g.:\n"
                    "  python run_traj_eval.py --python $CONDA_PREFIX/bin/python ...",
                    file=sys.stderr,
                )
                return 1
    print()

    for dataset in selected:
        dataset_dir = TRAJ_EVAL_ROOT / dataset.name
        dataset_dir.mkdir(parents=True, exist_ok=True)

        run_items: list[tuple[int, Path]] = []
        if args.eval_only:
            archived = discover_archived_runs(dataset.name)
            if not archived:
                print(f"[{dataset.name}] no archived runs found, skipping.")
                failures += 1
                continue
            run_items = [item for item in archived if item[0] <= args.runs]
        else:
            run_items = [(idx, dataset_dir / f"run_{idx:02d}") for idx in range(1, args.runs + 1)]

        for run_idx, run_dir in run_items:
            cmd = build_main_command(
                dataset,
                python_exec=args.python,
                quiet=args.quiet,
                no_progress=args.no_progress,
                no_time_cost=args.no_time_cost,
                extra_args=args.extra,
            )

            print("=" * 80)
            print(f"[{dataset.name}] run {run_idx}/{args.runs}")

            if args.eval_only:
                print(f"Evaluating archived run: {run_dir.relative_to(PROJECT_ROOT)}")
            else:
                print("Command:", " ".join(cmd))

            if args.dry_run:
                if not args.eval_only:
                    print(f"[dry-run] would save trajectories to {run_dir}")
                if not args.skip_eval:
                    if dataset.dataset_type.lower() == "kitti":
                        assets = resolve_kitti_eval_assets(
                            dataset,
                            seq_override=kitti_seq_override,
                            gt_override=gt_override,
                        )
                        if assets:
                            seq_id, gt_file, ts_file = assets
                            print(
                                f"[dry-run] would run eval_kitti.py seq={seq_id} "
                                f"gt={gt_file.relative_to(PROJECT_ROOT)} "
                                f"ts={ts_file.relative_to(PROJECT_ROOT)}"
                            )
                        else:
                            print("[dry-run] kitti eval skipped (missing seq/gt/timestamps)")
                    else:
                        gt = gt_override or resolve_ground_truth(dataset)
                        gt_text = str(gt.relative_to(PROJECT_ROOT)) if gt else "N/A"
                        print(
                            f"[dry-run] would trim {args.trim}+{args.trim} poses and run evo_ape "
                            f"against {gt_text}"
                        )
                continue

            duration = 0.0
            returncode = 0
            copied: list[str] = []

            if args.eval_only:
                if not (run_dir / EVAL_TRAJECTORY).is_file():
                    failures += 1
                    print(f"Missing {EVAL_TRAJECTORY} in {run_dir}")
                    continue
            else:
                start = time.time()
                proc = subprocess.run(cmd, cwd=PROJECT_ROOT)
                duration = time.time() - start
                returncode = proc.returncode

                copied = copy_trajectories(run_dir)
                if returncode != 0:
                    failures += 1
                    print(f"Run failed with exit code {returncode}")
                elif not copied:
                    failures += 1
                    print("Run finished but no trajectory files were found in output/.")
                else:
                    print(
                        f"Done in {duration:.1f}s, archived {len(copied)} trajectory file(s) "
                        f"to {run_dir.relative_to(PROJECT_ROOT)}"
                    )

            eval_result: dict[str, object] | None = None
            if not args.skip_eval and (args.eval_only or (returncode == 0 and copied)):
                is_kitti = (
                    dataset.dataset_type.lower() == "kitti"
                    and is_kitti_sequence_config(dataset)
                )
                eval_targets = KITTI_EVAL_TRAJECTORIES if is_kitti else (
                    (EVAL_TRAJECTORY, "optimized"),
                )

                for trajectory_file, trajectory_label in eval_targets:
                    if not (run_dir / trajectory_file).is_file():
                        if trajectory_label == "optimized":
                            print(f"Missing {trajectory_file} in {run_dir}")
                        continue

                    traj_eval_result = evaluate_run_trajectory(
                        run_dir,
                        dataset,
                        python_exec=args.python,
                        trim_head=args.trim,
                        trim_tail=args.trim,
                        align=not args.no_align,
                        correct_scale=not args.no_scale,
                        trajectory=trajectory_file,
                        eval_tag=trajectory_label,
                        gt_override=gt_override,
                        kitti_seq_override=kitti_seq_override,
                    )
                    print(_format_eval_log(traj_eval_result))
                    if trajectory_label == "optimized":
                        eval_result = traj_eval_result
                    if traj_eval_result and traj_eval_result.get("status") == "failed":
                        if trajectory_label == "optimized":
                            failures += 1

                    if traj_eval_result and traj_eval_result.get("status") in {"ok", "failed"}:
                        metrics = traj_eval_result.get("metrics", {})
                        if isinstance(metrics, dict):
                            if traj_eval_result.get("evaluator") == "eval_kitti.py":
                                if traj_eval_result.get("status") == "ok":
                                    kitti_session_metrics.append(
                                        {
                                            "dataset": dataset.name,
                                            "run_index": run_idx,
                                            "trajectory_label": trajectory_label,
                                            "t_rel_percent": metrics.get("t_rel_percent"),
                                            "r_rel_deg_per_100m": metrics.get(
                                                "r_rel_deg_per_100m"
                                            ),
                                        }
                                    )
                                append_evo_summary_row(
                                    kitti_summary_path,
                                    {
                                        "dataset": dataset.name,
                                        "dataset_type": dataset.dataset_type,
                                        "run_index": run_idx,
                                        "trajectory": trajectory_file,
                                        "trajectory_label": trajectory_label,
                                        "status": traj_eval_result.get("status"),
                                        "sequence_id": traj_eval_result.get("sequence_id"),
                                        "t_rel_percent": metrics.get("t_rel_percent"),
                                        "r_rel_deg_per_100m": metrics.get(
                                            "r_rel_deg_per_100m"
                                        ),
                                        "matched_pose_count": metrics.get("matched_poses"),
                                        "sim3_scale": metrics.get("sim3_scale"),
                                        "trim_head": args.trim,
                                        "trim_tail": args.trim,
                                        "ground_truth": traj_eval_result.get("ground_truth"),
                                        "timestamps": traj_eval_result.get("timestamps"),
                                        "output_dir": str(run_dir.relative_to(PROJECT_ROOT)),
                                        "timestamp": datetime.now().isoformat(
                                            timespec="seconds"
                                        ),
                                    },
                                )
                            elif trajectory_label == "optimized":
                                append_evo_summary_row(
                                    evo_summary_path,
                                    {
                                        "dataset": dataset.name,
                                        "dataset_type": dataset.dataset_type,
                                        "run_index": run_idx,
                                        "status": traj_eval_result.get("status"),
                                        "rmse": metrics.get("rmse"),
                                        "mean": metrics.get("mean"),
                                        "median": metrics.get("median"),
                                        "max": metrics.get("max"),
                                        "min": metrics.get("min"),
                                        "std": metrics.get("std"),
                                        "scale": metrics.get("scale"),
                                        "original_pose_count": traj_eval_result.get(
                                            "original_pose_count"
                                        ),
                                        "trimmed_pose_count": traj_eval_result.get(
                                            "trimmed_pose_count"
                                        ),
                                        "trim_head": args.trim,
                                        "trim_tail": args.trim,
                                        "ground_truth": traj_eval_result.get("ground_truth"),
                                        "output_dir": str(run_dir.relative_to(PROJECT_ROOT)),
                                        "timestamp": datetime.now().isoformat(
                                            timespec="seconds"
                                        ),
                                    },
                                )

            if not args.eval_only:
                write_run_meta(
                    run_dir,
                    dataset=dataset,
                    run_index=run_idx,
                    command=cmd,
                    returncode=returncode,
                    duration_sec=duration,
                    copied_files=copied,
                    eval_result=eval_result,
                )
                append_summary_row(
                    summary_path,
                    {
                        "dataset": dataset.name,
                        "dataset_type": dataset.dataset_type,
                        "run_index": run_idx,
                        "returncode": returncode,
                        "duration_sec": round(duration, 3),
                        "trajectory_count": len(copied),
                        "output_dir": str(run_dir.relative_to(PROJECT_ROOT)),
                        "config": dataset.rel_config,
                        "timestamp": datetime.now().isoformat(timespec="seconds"),
                    },
                )

    if args.dry_run:
        return 0

    print()
    print("=" * 80)
    if not args.skip_eval and kitti_session_metrics:
        print_kitti_session_summary(kitti_session_metrics)
        print("=" * 80)
    if failures:
        print(f"Evaluation finished with {failures} failed run(s).")
        print(f"Run summary: {summary_path.relative_to(PROJECT_ROOT)}")
        if not args.skip_eval:
            print(f"Evo summary: {evo_summary_path.relative_to(PROJECT_ROOT)}")
            print(f"KITTI summary: {kitti_summary_path.relative_to(PROJECT_ROOT)}")
        return 1

    print("Evaluation finished successfully.")
    print(f"Run summary: {summary_path.relative_to(PROJECT_ROOT)}")
    if not args.skip_eval:
        print(f"Evo summary: {evo_summary_path.relative_to(PROJECT_ROOT)}")
        print(f"KITTI summary: {kitti_summary_path.relative_to(PROJECT_ROOT)}")
    return 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run main.py on dataset configs and archive trajectories to output/traj_eval."
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help="List all discovered dataset configs and exit.",
    )
    parser.add_argument(
        "--dataset",
        action="append",
        default=[],
        metavar="NAME",
        help="Dataset name (e.g. euroc_MH01). Repeatable.",
    )
    parser.add_argument(
        "--type",
        action="append",
        default=[],
        metavar="TYPE",
        help=(
            "Run all configs of a dataset type (euroc, tum, kitti, tartanair). "
            "For kitti, only config/kitti/00.yaml .. 10.yaml are included. Repeatable."
        ),
    )
    parser.add_argument(
        "--config",
        action="append",
        default=[],
        metavar="PATH",
        help="Direct path to a yaml config. Repeatable.",
    )
    parser.add_argument(
        "--all",
        action="store_true",
        help="Run all discovered dataset configs.",
    )
    parser.add_argument(
        "--runs",
        type=int,
        default=1,
        help="Number of repeated runs per dataset (default: 1).",
    )
    parser.add_argument(
        "--python",
        default=sys.executable,
        help=(
            "Python executable for main.py and eval_kitti.py "
            "(use ivggt env python if evo is missing, e.g. "
            "$CONDA_PREFIX/bin/python after conda activate ivggt)."
        ),
    )
    parser.add_argument(
        "--quiet",
        action="store_true",
        default=True,
        help="Pass --quiet to main.py (default: enabled).",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Disable --quiet when running main.py.",
    )
    parser.add_argument(
        "--no-progress",
        action="store_true",
        help="Disable the SLAM progress bar in main.py.",
    )
    parser.add_argument(
        "--no-time-cost",
        action="store_true",
        help="Pass --no-time-cost to main.py.",
    )
    parser.add_argument(
        "--extra",
        nargs=argparse.REMAINDER,
        default=[],
        help="Extra arguments forwarded to main.py (prefix with --).",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print planned commands without executing main.py.",
    )
    parser.add_argument(
        "--skip-eval",
        action="store_true",
        help="Skip evo_ape accuracy evaluation.",
    )
    parser.add_argument(
        "--eval-only",
        action="store_true",
        help="Only run evo_ape on archived trajectories under output/traj_eval.",
    )
    parser.add_argument(
        "--trim",
        type=int,
        default=TRIM_HEAD_DEFAULT,
        help=(
            "Trim this many poses from the head and tail of trajectory_optimized.txt "
            f"before evo_ape (default: {TRIM_HEAD_DEFAULT})."
        ),
    )
    parser.add_argument(
        "--gt",
        metavar="PATH",
        help="Override ground-truth file for all selected datasets (KITTI: pose/xx.txt).",
    )
    parser.add_argument(
        "--kitti-seq",
        metavar="ID",
        help="Override inferred KITTI odometry sequence id (e.g. 3 or 03).",
    )
    parser.add_argument(
        "--no-align",
        action="store_true",
        help="Disable Sim(3) alignment (-a) in evo_ape.",
    )
    parser.add_argument(
        "--no-scale",
        action="store_true",
        help="Disable scale correction (-s) in evo_ape.",
    )
    args = parser.parse_args()

    if args.verbose:
        args.quiet = False
    return args


def main() -> None:
    args = parse_args()
    raise SystemExit(run_evaluation(args))


if __name__ == "__main__":
    main()