#!/usr/bin/env python3
"""Attach frozen APE loss weights to the unchanged Florence-Bridge draw stream."""

from __future__ import annotations

import argparse
import os

from omegaconf import OmegaConf

from external_baselines.keyframe_focused_adapted.keyframe_cache import (
    KeyframeWeightCache,
)


def install_adapter(cache_path: str) -> KeyframeWeightCache:
    from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset
    from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest

    cache = KeyframeWeightCache(cache_path)
    if os.environ.get("KEYFRAME_SMOKE_FROZEN_MANIFEST") == "1":
        original_build_manifest = LeRobotMixtureDataset._build_valid_anchor_manifest

        def load_frozen_manifest(self):
            path = self.data_cfg.get("valid_anchor_manifest_path")
            if not path:
                return original_build_manifest(self)
            manifest = ValidAnchorManifest.load(path)
            print(
                "Smoke-only frozen manifest fast path: "
                f"{path}",
                flush=True,
            )
            return manifest

        LeRobotMixtureDataset._build_valid_anchor_manifest = load_frozen_manifest

    original_getitem = LeRobotMixtureDataset.__getitem__

    def keyframe_getitem(self, index):
        sample = original_getitem(self, index)
        trace = sample["sampling_trace"]
        record = cache.lookup(
            trace["dataset_name"],
            int(trace["trajectory_id"]),
            int(trace["base_index"]),
        )
        # The logical draw and physical timestep are unchanged; only loss weight is added.
        trace["keyframe_ape"] = record.ape
        trace["keyframe_weight"] = record.weight
        trace["keyframe_cache"] = cache.path.name
        return sample

    LeRobotMixtureDataset.__getitem__ = keyframe_getitem
    return cache


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

    cache = install_adapter(os.environ["KEYFRAME_APE_CACHE"])
    print(
        f"Installed Keyframe-Focused-adapted cache {cache.path}",
        flush=True,
    )

    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(clipargs)))
    cfg = apply_config_compat(cfg, strict=True)
    if cfg.datasets.vla_data.temporal_proposal_path is not None:
        raise ValueError("Keyframe loss weighting requires the uniform temporal sampler")
    if cfg.trainer.sample_loss_weight_key != "keyframe_weight":
        raise ValueError("trainer.sample_loss_weight_key must be keyframe_weight")
    cfg.config_yaml = args.config_yaml
    train_starvla.main(cfg)


if __name__ == "__main__":
    main()