#!/usr/bin/env python3
"""DCVC-RT variant of the cluster's bitcost-readiness codec generator.

This reuses ``codec_tools/pipeline/process_video_bitcost_readiness.py`` in full
and swaps ONLY the per-frame score-map source: instead of h264 per-block bits
(``cv_reader_fetch_bitcost``), it uses the DCVC-RT bit-cost bitmap. Everything
else — frame sampling, readiness grouping, 2x2-block selection, canvas packing,
``src_patch_position`` / ``meta.json`` writing — is identical, so the only
variable vs the baseline assets is "DCVC-RT bits vs h264 bitcost".

Usage mirrors the cluster generator (same CLI args), e.g.::

    python dcvc_readiness_gen.py --video V.mp4 --out_dir OUT \
        --num_sampled_frames 256 --grouping_mode readiness \
        --readiness_sum_threshold_mode auto --group_size 32 \
        --images_per_group 4 --patch 16 --max_pixels 150000 \
        --min_group_frames 8 --max_group_frames 128 --bitcost_grid sub

Env (infra only):
    DCVC_INTRA_TAR / DCVC_INTER_TAR  DCVC-RT checkpoints (default: bundled tars in
                                     the model dir; set to relocate)
    DCVC_DEVICE                      cuda:N (default cuda:0)
Selection/scoring knobs (qp, reset_interval, intra_period, max_side, and the
readiness grouping params) come from ``preprocessor_config.json``'s ``codec.dcvc``
via ``codec_dcvc_config`` — NOT from env. The DCVC-RT source is bundled at
``neural_codec/DCVC/`` and loaded by ``dcvc_rt_engine`` (no env var).
Run in the ``magevl`` conda env (torch + DCVC-RT ext + the codec/video deps).
"""

from __future__ import annotations

import os
import sys

import cv2
import numpy as np

# --- make repo + engine importable -----------------------------------------
# Bundled in the release ``neural_codec/`` package: codec_tools/ and dcvc_rt_engine.py
# live next to this file, so default both search roots to this directory (env
# vars still override, e.g. in the cluster/dev checkout).
_HERE = os.path.dirname(os.path.abspath(__file__))
_REPO = os.environ.get("DCVC_REPO_DIR", _HERE)          # dir containing codec_tools/
_ENGINE_DIR = os.environ.get("DCVC_ENGINE_DIR", _HERE)  # dir containing dcvc_rt_engine.py
for p in (_REPO, _ENGINE_DIR):
    if p not in sys.path:
        sys.path.insert(0, p)

import importlib

P = importlib.import_module("codec_tools.pipeline.process_video_bitcost_readiness")
from dcvc_rt_engine import DCVCRTEngine  # noqa: E402
import codec_dcvc_config as _dc  # noqa: E402  (DCVC params from preprocessor_config.json)


_ENGINE = None


def _get_engine() -> DCVCRTEngine:
    global _ENGINE
    if _ENGINE is None:
        _ENGINE = DCVCRTEngine(
            intra_ckpt=os.environ.get("DCVC_INTRA_TAR", os.path.join(_HERE, "dcvc_rt_intra.tar")),
            inter_ckpt=os.environ.get("DCVC_INTER_TAR", os.path.join(_HERE, "dcvc_rt_inter.tar")),
            device=os.environ.get("DCVC_DEVICE", "cuda:0"),
            intra_period=int(_dc.get("intra_period")),
            reset_interval=int(_dc.get("reset_interval")),
        )
    return _ENGINE


def _dcvc_bitmaps(video_path: str, frame_ids):
    """Sequentially DCVC-encode ``video_path`` and return {fid: bitmap (h/16,w/16)}."""
    eng = _get_engine()
    max_side = int(_dc.get("max_side"))
    needed = set(int(f) for f in frame_ids)
    max_fid = max(needed) if needed else -1

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"cannot open {video_path}")
    out = {}
    started = False
    last = None
    idx = 0
    while idx <= max_fid:
        ret, bgr = cap.read()
        if not ret or bgr is None:
            if last is None:
                break
            bgr = last
        else:
            last = bgr
        if max_side > 0:
            h, w = bgr.shape[:2]
            if max(h, w) > max_side:
                s = max_side / float(max(h, w))
                bgr = cv2.resize(bgr, (max(1, int(w * s)), max(1, int(h * s))), interpolation=cv2.INTER_AREA)
        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        if not started:
            _qp = int(_dc.get("qp"))
            eng.reset_sequence(rgb.shape[0], rgb.shape[1], qp_i=_qp, qp_p=_qp)
            started = True
        bm = eng.step(idx, rgb)
        if idx in needed:
            out[idx] = bm.float().cpu().numpy()
        idx += 1
    cap.release()
    return out


# ------------------------------------------------------------------ patches

_STATE: dict = {}


def _patched_fetch_bitcost(video_path, frame_ids):
    """Stand-in for cv_reader_fetch_bitcost: just carry (video, fids) forward."""
    _STATE["video"] = video_path
    return [{"fid": int(f)} for f in frame_ids]


def _patched_items_to_score_maps(bitcost_items, out_h, out_w, **kwargs):
    """Produce DCVC-RT score maps at (out_h, out_w) for the same frame ids."""
    # Random-patch baseline: skip the DCVC encode and return uniform random score
    # maps (control). Controlled by codec.dcvc.random_select in preprocessor_config.json.
    if _dc.get("random_select", bool):
        rng = np.random.default_rng(int(_dc.get("random_seed")))
        return [rng.random((int(out_h), int(out_w)), dtype=np.float32) for _ in bitcost_items]
    fids = [int(it["fid"]) for it in bitcost_items]
    bmaps = _dcvc_bitmaps(_STATE["video"], fids)
    maps = []
    for f in fids:
        bm = bmaps.get(f)
        if bm is None or bm.size == 0:
            bm = np.zeros((max(1, out_h // 16), max(1, out_w // 16)), dtype=np.float32)
        maps.append(cv2.resize(bm.astype(np.float32), (int(out_w), int(out_h)), interpolation=cv2.INTER_LINEAR))
    return maps


def main():
    # Swap only the score-map source; reuse the entire generator.
    P.cv_reader_fetch_bitcost = _patched_fetch_bitcost
    P.bitcost_items_to_score_maps = _patched_items_to_score_maps
    P.main()


if __name__ == "__main__":
    main()