#!/usr/bin/env python3
"""Run the Florence-RoboTwin FrameSkip continuation from the 10K state."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

import numpy as np
from omegaconf import OmegaConf


METHOD = "FrameSkip-adapted (Gaussian-TPI fallback)"


def load_cache(path: str) -> tuple[dict[tuple[str, str], np.ndarray], dict]:
    payload = json.loads(Path(path).read_text(encoding="utf-8"))
    if payload.get("method") != METHOD:
        raise RuntimeError(f"unexpected FrameSkip method: {payload.get('method')!r}")
    if float(payload.get("retention_ratio", -1.0)) != 0.2:
        raise RuntimeError("FrameSkip cache must use retention_ratio=0.2")

    cache: dict[tuple[str, str], np.ndarray] = {}
    for record in payload["records"]:
        key = (record["dataset_name"], str(record["trajectory_id"]))
        keep = np.asarray(record["keep_indices"], dtype=np.int64)
        if key in cache:
            raise RuntimeError(f"duplicate FrameSkip trajectory: {key}")
        if len(keep) == 0 or np.any(np.diff(keep) <= 0):
            raise RuntimeError(f"invalid retained indices for {key}")
        cache[key] = keep
    return cache, payload


def install_adapter(
    cache_path: str,
    start_global_draw_id: int,
    global_microbatch_size: int,
) -> None:
    from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset

    cache, cache_payload = load_cache(cache_path)
    original_getitem = LeRobotMixtureDataset.__getitem__
    period = 6  # Five pruned microbatches followed by one full-frame microbatch.

    def frameskip_getitem(self, index):
        trace = self.sample_trace(index)
        offset = int(trace.global_draw_id) - start_global_draw_id
        if offset < 0:
            raise RuntimeError("FrameSkip draw precedes the 10K continuation boundary")

        microbatch = offset // global_microbatch_size
        is_pruned = microbatch % period < 5
        logical_index = int(trace.base_index)
        actual_index = logical_index

        if is_pruned:
            key = (trace.dataset_name, str(trace.trajectory_id))
            try:
                retained = cache[key]
            except KeyError as exc:
                raise RuntimeError(f"FrameSkip cache is missing trajectory {key}") from exc
            position = int(np.searchsorted(retained, logical_index))
            actual_index = int(retained[min(position, len(retained) - 1)])

        dataset = self.datasets[trace.dataset_index]
        original_get_step_data = dataset.get_step_data

        def redirected_get_step_data(trajectory_id, base_index):
            if str(trajectory_id) != str(trace.trajectory_id):
                raise RuntimeError("FrameSkip attempted to change trajectory identity")
            if int(base_index) != logical_index:
                raise RuntimeError("unexpected nested dataset read during FrameSkip remapping")
            return original_get_step_data(trajectory_id, actual_index)

        dataset.get_step_data = redirected_get_step_data
        try:
            sample = original_getitem(self, index)
        finally:
            dataset.get_step_data = original_get_step_data

        sample["sampling_trace"]["frameskip"] = {
            "method": cache_payload["method"],
            "cache_record_count": len(cache_payload["records"]),
            "is_pruned_minibatch": is_pruned,
            "logical_base_index": logical_index,
            "physical_base_index": actual_index,
            "retention_ratio": 0.2 if is_pruned else 1.0,
        }
        return sample

    LeRobotMixtureDataset.__getitem__ = frameskip_getitem


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config_yaml", required=True)
    args, overrides = parser.parse_known_args()

    install_adapter(
        cache_path=os.environ["FRAMESKIP_CACHE"],
        start_global_draw_id=int(os.environ.get("FRAMESKIP_START_GLOBAL_DRAW_ID", "640000")),
        global_microbatch_size=int(os.environ.get("FRAMESKIP_GLOBAL_MICROBATCH_SIZE", "4")),
    )

    from starVLA.model.framework.share_tools import apply_config_compat
    from starVLA.training import train_starvla
    from starVLA.training.trainer_utils.trainer_tools import normalize_dotlist_args

    cfg = OmegaConf.load(args.config_yaml)
    cfg = OmegaConf.merge(cfg, OmegaConf.from_dotlist(normalize_dotlist_args(overrides)))
    cfg = apply_config_compat(cfg, strict=True)
    cfg.config_yaml = args.config_yaml
    train_starvla.main(cfg)


if __name__ == "__main__":
    main()