"""Convert saved EgoSpeak MiniROADA forecasts for floor-event evaluation.

EgoSpeak MiniROADA saves one array per session as
``[source_time, future_horizon, class]``.  The repository floor evaluator uses
the LSTR-compatible target-aligned layout ``[target_time, class, horizon]``
plus an explicit causal-valid mask.  This script performs only that alignment;
it does not change scores, targets, or class order.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np


NUM_CLASSES = 3
NUM_HORIZONS = 10


def normalize_session_name(value: str | int) -> str:
    text = str(value)
    suffix = text.removeprefix("Session_")
    if not suffix.isdigit():
        raise ValueError(f"Invalid EasyCom session: {value}")
    return f"Session_{int(suffix)}"


def align_miniroad_prediction(
    raw_prediction: np.ndarray,
    target_steps: int,
) -> tuple[np.ndarray, np.ndarray]:
    """Align MiniROADA source-indexed forecasts without padding them as valid."""

    raw_prediction = np.asarray(raw_prediction)
    expected = (target_steps - NUM_HORIZONS, NUM_HORIZONS, NUM_CLASSES)
    if raw_prediction.shape != expected:
        raise ValueError(
            "Expected MiniROADA prediction "
            f"{expected} for target length {target_steps}, got "
            f"{raw_prediction.shape}"
        )
    if not np.isfinite(raw_prediction).all():
        raise ValueError("MiniROADA prediction contains NaN/Inf")
    if (raw_prediction < -1e-6).any() or (raw_prediction > 1.0 + 1e-6).any():
        raise ValueError("MiniROADA prediction is outside [0,1]")

    raw_prediction = np.asarray(raw_prediction, dtype=np.float32)
    aligned = np.zeros(
        (target_steps, NUM_CLASSES, NUM_HORIZONS), dtype=np.float32
    )
    valid = np.zeros((target_steps, NUM_HORIZONS), dtype=bool)
    source_rows = np.arange(len(raw_prediction), dtype=np.int64)
    for horizon_index in range(NUM_HORIZONS):
        target_rows = source_rows + horizon_index + 1
        aligned[target_rows, :, horizon_index] = raw_prediction[
            :, horizon_index, :
        ]
        valid[target_rows, horizon_index] = True
    return aligned, valid


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--prediction-root",
        type=Path,
        required=True,
        help="Directory containing raw MiniROADA Session_X.npy predictions.",
    )
    parser.add_argument("--target-root", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--sessions", nargs="+", required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    sessions = [normalize_session_name(value) for value in args.sessions]
    if len(sessions) != len(set(sessions)):
        raise ValueError("Duplicate sessions")

    converted: dict[str, tuple[np.ndarray, np.ndarray]] = {}
    sources: dict[str, str] = {}
    for session in sessions:
        target_path = args.target_root / f"{session}.npy"
        prediction_path = args.prediction_root / f"{session}.npy"
        if not target_path.is_file():
            raise FileNotFoundError(target_path)
        if not prediction_path.is_file():
            raise FileNotFoundError(prediction_path)
        target = np.load(target_path, allow_pickle=False)
        if target.ndim != 2 or target.shape[1] != NUM_CLASSES:
            raise ValueError(
                f"Expected target [T,{NUM_CLASSES}], got {target.shape}: "
                f"{target_path}"
            )
        raw_prediction = np.load(prediction_path, allow_pickle=False)
        converted[session] = align_miniroad_prediction(
            raw_prediction, len(target)
        )
        sources[session] = str(prediction_path.expanduser().resolve())

    prediction_dir = args.output_dir / "pred"
    valid_dir = args.output_dir / "valid"
    destinations = [
        directory / f"{session}.npy"
        for directory in (prediction_dir, valid_dir)
        for session in sessions
    ]
    destinations.append(args.output_dir / "conversion.json")
    existing = [path for path in destinations if path.exists()]
    if existing:
        raise FileExistsError(
            "Refusing to overwrite converted output: "
            + ", ".join(str(path) for path in existing)
        )

    prediction_dir.mkdir(parents=True, exist_ok=True)
    valid_dir.mkdir(parents=True, exist_ok=True)
    for session, (prediction, valid) in converted.items():
        np.save(prediction_dir / f"{session}.npy", prediction)
        np.save(valid_dir / f"{session}.npy", valid)
    metadata = {
        "protocol": "egospeak-miniroada-source-to-target-alignment-v1",
        "source_layout": "[source_time, future_horizon, class]",
        "output_layout": "[target_time, class, future_horizon]",
        "sessions": sessions,
        "sources": sources,
    }
    (args.output_dir / "conversion.json").write_text(
        json.dumps(metadata, indent=2) + "\n"
    )
    print(f"Converted {len(sessions)} session(s): {args.output_dir}")


if __name__ == "__main__":
    main()