scripts/_core.py

"""Shared scoring core (per-dimension version).

Backend-agnostic helpers for the vanilla video-reward pipeline. The
``infer.py`` entry script handles CLI parsing and client construction;
everything below is duck-typed on a client object exposing:

    client.infer_with_frames(user_text, frame_b64_list, system_text)
    client.max_retries          : int
    client.retry_base_delay     : float
    client.request_interval     : float

Each video triggers ONE multimodal model call **per scoring dimension**
(see ``prompts/vanilla_prompts.py``). For the canonical 3-dim setup that
means up to 3 calls per video. Each reply must be a single JSON object
of the form::

    {"reasoning": "...", "score": <int 1-5>}

The orchestrator aggregates the per-dimension replies into the SAME
``scoring`` block layout used by the previous single-call pipeline, so
all downstream consumers (``compute_mae.py``, the gradio app, etc.)
keep working unchanged::

    {
      "dimensions": [
        {"dimension": "instruction_following", "score": 4, "reasoning": "..."},
        {"dimension": "visual_quality",        "score": 3, "reasoning": "..."},
        {"dimension": "world_consistency",     "score": 5, "reasoning": "..."}
      ],
      "scores_by_dim": {
        "instruction_following": 4,
        "visual_quality": 3,
        "world_consistency": 5
      }
    }

Resumability
------------
Per-dimension calls are independent, so we can cache partial progress.
Two layers are persisted under the same parent dir as ``score_path``:

  * ``<score_path>``                — final results (only fully-scored
                                      videos appear here, same schema as
                                      before).
  * ``<score_path>.partial.json``   — per-video, per-dim cache. Each
                                      entry maps ``video_id`` to a dict
                                      ``{dim: {"reasoning": str, "score": int}}``
                                      containing only the dims that have
                                      already succeeded.

On rerun, fully-scored videos are skipped; partially-scored videos only
re-issue the missing dimensions; failed videos are retried from scratch
(of whatever is missing).
"""
from __future__ import annotations

import os
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Optional

from tqdm import tqdm

# --- Path bootstrapping so we can import from `prompts/` and `tools/` -----
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(SCRIPT_DIR)
for _p in (PROJECT_ROOT, os.path.join(PROJECT_ROOT, "prompts")):
    if _p not in sys.path:
        sys.path.insert(0, _p)

from prompts.vanilla_prompts import (  # noqa: E402
    DIMENSIONS,
    NUM_SAMPLED_FRAMES,
    build_prompt,
)
from tools import (  # noqa: E402
    ensure_video_local,
    extract_frames_from_video,
    load_json_safe,
    parse_json_from_model_output,
    save_json_atomic,
)

# ---------------------------------------------------------------------------
# Shared defaults
# ---------------------------------------------------------------------------

DEFAULT_DATA_PATH = os.path.join(PROJECT_ROOT, "data", "firm-video-bench.json")
DEFAULT_RESULTS_DIR = os.path.join(PROJECT_ROOT, "results")

VALID_SCORES = {1, 2, 3, 4, 5}


# ---------------------------------------------------------------------------
# Output normalization (per-dimension)
# ---------------------------------------------------------------------------

def _coerce_score(score_raw: Any) -> int:
    if isinstance(score_raw, bool):  # bool is a subclass of int — reject
        raise ValueError(f"score must be an integer, got bool: {score_raw}")
    if isinstance(score_raw, int):
        score = score_raw
    elif isinstance(score_raw, float) and float(score_raw).is_integer():
        score = int(score_raw)
    elif isinstance(score_raw, str) and score_raw.strip().lstrip("-").isdigit():
        score = int(score_raw.strip())
    else:
        raise ValueError(f"score must be an integer in 1-5, got: {score_raw!r}")
    if score not in VALID_SCORES:
        raise ValueError(f"score out of range 1-5: {score}")
    return score


def _normalize_single_dim_output(parsed: Any) -> dict[str, Any]:
    """Validate the per-dimension JSON ``{"reasoning":..., "score":...}``."""
    if not isinstance(parsed, dict):
        raise ValueError(f"Expected JSON object, got {type(parsed).__name__}")
    reasoning = parsed.get("reasoning", "")
    if not isinstance(reasoning, str):
        reasoning = str(reasoning)
    score = _coerce_score(parsed.get("score"))
    return {"reasoning": reasoning.strip(), "score": score}


def _aggregate_scoring(
    dim_results: dict[str, dict[str, Any]],
) -> dict[str, Any]:
    """Build the final ``scoring`` block from per-dimension results."""
    dimensions_block = [
        {
            "dimension": dim,
            "score": dim_results[dim]["score"],
            "reasoning": dim_results[dim]["reasoning"],
        }
        for dim in DIMENSIONS
    ]
    scores_by_dim = {dim: dim_results[dim]["score"] for dim in DIMENSIONS}
    return {
        "dimensions": dimensions_block,
        "scores_by_dim": scores_by_dim,
    }


def _sort_key(video_id: str) -> tuple[int, str]:
    try:
        return int(str(video_id).split("_", 1)[0]), str(video_id)
    except Exception:  # noqa: BLE001
        return 10**12, str(video_id)


# ---------------------------------------------------------------------------
# Per-dimension scoring (one model call per dimension)
# ---------------------------------------------------------------------------

def _score_one_dimension(
    client: Any,
    dimension: str,
    video_prompt: Optional[str],
    frame_b64_list: list[str],
    video_id: str,
) -> Optional[dict[str, Any]]:
    """Issue ONE multimodal call for a single ``dimension``.

    All three dimensions consume ``video_prompt``: ``instruction_following``
    evaluates the video against it; ``visual_quality`` and
    ``world_consistency`` use it as context only (see
    ``prompts/vanilla_prompts.py``).

    Returns ``{"reasoning": str, "score": int}`` on success or ``None``
    after exhausting retries.
    """
    spec = build_prompt(
        dimension=dimension,
        video_prompt=video_prompt,
        n_frames=len(frame_b64_list),
    )
    user_text = spec["user"]

    last_err: Optional[str] = None
    for attempt in range(client.max_retries):
        try:
            raw = client.infer_with_frames(
                user_text=user_text,
                frame_b64_list=frame_b64_list,
                system_text=spec["system"],
            )
            parsed = parse_json_from_model_output(raw)
            normalized = _normalize_single_dim_output(parsed)
            print(
                f"    [{video_id}::{dimension}] OK score={normalized['score']}"
            )
            return normalized
        except Exception as exc:  # noqa: BLE001
            last_err = str(exc)
            print(
                f"    [{video_id}::{dimension}] attempt "
                f"{attempt + 1}/{client.max_retries} failed: {exc}"
            )
            if attempt < client.max_retries - 1:
                time.sleep(client.retry_base_delay * (2 ** attempt))

    print(
        f"    [{video_id}::{dimension}] FAILED after "
        f"{client.max_retries} attempts ({last_err})"
    )
    return None


# ---------------------------------------------------------------------------
# Partial cache (per-dim) helpers
# ---------------------------------------------------------------------------

def _partial_path(score_path: str) -> str:
    """Sidecar path for the per-dim partial cache."""
    return score_path + ".partial.json"


def _load_partial(score_path: str) -> dict[str, dict[str, dict[str, Any]]]:
    """Load the per-video, per-dim cache, validating shape.

    Returns ``{video_id: {dim: {"reasoning": str, "score": int}}}``.
    Malformed entries are silently dropped.
    """
    raw = load_json_safe(_partial_path(score_path), default={})
    if not isinstance(raw, dict):
        return {}
    out: dict[str, dict[str, dict[str, Any]]] = {}
    for vid, dims in raw.items():
        if not isinstance(dims, dict):
            continue
        clean: dict[str, dict[str, Any]] = {}
        for dim, entry in dims.items():
            if dim not in DIMENSIONS or not isinstance(entry, dict):
                continue
            try:
                clean[dim] = _normalize_single_dim_output(entry)
            except Exception:  # noqa: BLE001
                continue
        if clean:
            out[str(vid)] = clean
    return out


def _save_partial(
    score_path: str,
    partial: dict[str, dict[str, dict[str, Any]]],
) -> None:
    save_json_atomic(partial, _partial_path(score_path))


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------

def run_scoring(
    client: Any,
    expanded_data: list[dict[str, Any]],
    score_path: str,
    concurrency: int,
) -> None:
    print("\n" + "=" * 72)
    print("Vanilla per-dimension scoring (one model call per dim per video)")
    print(f"Dimensions ({len(DIMENSIONS)}): {DIMENSIONS}")
    print(f"Frames per video: {NUM_SAMPLED_FRAMES}")
    print(f"Concurrency:      {concurrency} (videos in flight)")
    print(f"Calls per video:  up to {len(DIMENSIONS)}")
    print("=" * 72)

    # Final results (fully-scored videos only).
    results = load_json_safe(score_path, default=[])
    if not isinstance(results, list):
        results = []
    result_by_id = {
        r.get("video_id"): r for r in results if r.get("scoring") is not None
    }

    # Per-dim partial cache (covers BOTH not-yet-scored and partially-
    # scored videos). Final results take precedence.
    partial_by_id = _load_partial(score_path)
    for vid in list(partial_by_id.keys()):
        if vid in result_by_id:
            partial_by_id.pop(vid, None)

    todo = [
        item for item in expanded_data if item["video_id"] not in result_by_id
    ]
    n_partial = sum(
        1 for item in todo if partial_by_id.get(item["video_id"])
    )
    print(
        f"[scoring] videos={len(expanded_data)}, "
        f"done={len(result_by_id)}, remaining={len(todo)} "
        f"(of which {n_partial} have partial per-dim progress)"
    )
    if not todo:
        return

    results_lock = threading.Lock()
    partial_lock = threading.Lock()
    fail_counter = {"n": 0}
    fail_lock = threading.Lock()

    def process_video(item: dict[str, Any]) -> None:
        video_id = item["video_id"]
        caption = item["caption"]
        try:
            video_path = ensure_video_local(item)
            frame_b64_list = extract_frames_from_video(
                video_path, num_frames=NUM_SAMPLED_FRAMES
            )
        except Exception as exc:  # noqa: BLE001
            print(f"  [{video_id}] preparation failed: {exc}")
            with fail_lock:
                fail_counter["n"] += 1
            return

        # Start from any cached per-dim results for this video.
        cached = dict(partial_by_id.get(video_id, {}))
        missing_dims = [d for d in DIMENSIONS if d not in cached]
        print(
            f"  [{video_id}] scoring {video_path} "
            f"({len(frame_b64_list)} frames; "
            f"cached={len(cached)}/{len(DIMENSIONS)}, "
            f"to_run={missing_dims})"
        )

        dim_results: dict[str, dict[str, Any]] = dict(cached)
        # Run missing dims sequentially within a single video (keeps the
        # client's retry/rate-limit semantics simple). Outer thread pool
        # provides throughput across videos.
        for dim in missing_dims:
            single = _score_one_dimension(
                client=client,
                dimension=dim,
                video_prompt=caption,
                frame_b64_list=frame_b64_list,
                video_id=video_id,
            )
            if single is None:
                # Persist whatever succeeded so far, then bail on this video.
                if dim_results:
                    with partial_lock:
                        partial_by_id[video_id] = dim_results
                        _save_partial(score_path, partial_by_id)
                with fail_lock:
                    fail_counter["n"] += 1
                time.sleep(client.request_interval)
                return

            dim_results[dim] = single

            # Persist incrementally so a crash mid-video doesn't waste
            # successful per-dim calls.
            with partial_lock:
                partial_by_id[video_id] = dict(dim_results)
                _save_partial(score_path, partial_by_id)

            # Light rate-limit between successive per-dim calls.
            time.sleep(client.request_interval)

        # All dimensions succeeded: assemble the final record.
        scoring = _aggregate_scoring(dim_results)
        result_item = {
            "video_id": video_id,
            "video_name": item["video_name"],
            "caption": caption,
            "video_path": item.get("video_path", ""),
            "video_local_path": video_path,
            "source": item.get("source", ""),
            "source_index": item["source_index"],
            "source_row_index": item.get("source_row_index"),
            "metadata": item.get("metadata", {}),
            "scoring": scoring,
        }

        with results_lock:
            result_by_id[video_id] = result_item
            ordered = [
                result_by_id[k]
                for k in sorted(result_by_id.keys(), key=_sort_key)
            ]
            save_json_atomic(ordered, score_path)

        # Drop this video from the partial cache once it's fully done.
        with partial_lock:
            partial_by_id.pop(video_id, None)
            _save_partial(score_path, partial_by_id)

        dim_summary = " | ".join(
            f"{d[:10]}={scoring['scores_by_dim'][d]}" for d in DIMENSIONS
        )
        print(f"  [{video_id}] OK {dim_summary}")

    with ThreadPoolExecutor(max_workers=max(1, int(concurrency))) as pool:
        futures = [pool.submit(process_video, item) for item in todo]
        for _ in tqdm(
            as_completed(futures), total=len(futures), desc="vanilla scoring"
        ):
            pass

    print(
        f"[scoring] saved={score_path}; "
        f"partial_cache={_partial_path(score_path)}; "
        f"failures/skips={fail_counter['n']}"
    )


__all__ = [
    "DIMENSIONS",
    "NUM_SAMPLED_FRAMES",
    "VALID_SCORES",
    "DEFAULT_DATA_PATH",
    "DEFAULT_RESULTS_DIR",
    "PROJECT_ROOT",
    "SCRIPT_DIR",
    "run_scoring",
]