#!/usr/bin/env python3
"""Verify that a cache covers an unchanged range of uniform logical draws."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from starVLA.dataloader.infoflow_sampler import LogicalDraw, ValidAnchorManifest

from external_baselines.keyframe_focused_adapted.keyframe_cache import (
    KeyframeWeightCache,
)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--cache", type=Path, required=True)
    parser.add_argument("--seed", type=int)
    parser.add_argument("--draw-start", type=int)
    parser.add_argument("--draw-count", type=int)
    parser.add_argument("--full-support", action="store_true")
    args = parser.parse_args()

    manifest = ValidAnchorManifest.load(args.manifest)
    cache = KeyframeWeightCache(args.cache)
    if args.full_support:
        arrays = np.load(args.cache, allow_pickle=False)
        metadata = json.loads(str(arrays["metadata_json"]))
        counts: dict[str, int] = {}
        weights: dict[str, int] = {}
        for dataset_index, dataset in enumerate(manifest.datasets):
            prefix = f"d{dataset_index}_"
            trajectory_ids = arrays[prefix + "trajectory_ids"]
            offsets = arrays[prefix + "offsets"]
            anchors = arrays[prefix + "anchors"]
            cached_weights = arrays[prefix + "weights"]
            expected_trajectory_ids = np.asarray(
                [int(episode.trajectory_id) for episode in dataset.episodes],
                dtype=np.int64,
            )
            if not np.array_equal(trajectory_ids, expected_trajectory_ids):
                raise RuntimeError(
                    f"cache trajectory support does not match {dataset.name}"
                )
            for position, episode in enumerate(dataset.episodes):
                cached = anchors[int(offsets[position]) : int(offsets[position + 1])]
                expected = np.asarray(episode.anchors, dtype=np.int32)
                if not np.array_equal(cached, expected):
                    raise RuntimeError(
                        f"cache anchor support does not match "
                        f"{dataset.name}/{int(episode.trajectory_id)}"
                    )
            counts[dataset.name] = int(len(anchors))
            values, value_counts = np.unique(cached_weights, return_counts=True)
            for value, count in zip(values, value_counts, strict=True):
                label = str(float(value))
                weights[label] = weights.get(label, 0) + int(count)
        print(
            json.dumps(
                {
                    "status": "PASS",
                    "mode": "full-support",
                    "cache": cache.path.name,
                    "dataset_anchor_counts": counts,
                    "weight_counts": weights,
                },
                sort_keys=True,
            )
        )
        return
    if args.seed is None or args.draw_start is None or args.draw_count is None:
        parser.error("draw-range validation requires --seed, --draw-start, and --draw-count")
    counts: dict[str, int] = {}
    weights: dict[str, int] = {}
    for draw_id in range(args.draw_start, args.draw_start + args.draw_count):
        trace = manifest.sample(args.seed, LogicalDraw(0, draw_id))
        record = cache.lookup(
            trace.dataset_name, trace.trajectory_id, trace.base_index
        )
        counts[trace.dataset_name] = counts.get(trace.dataset_name, 0) + 1
        label = str(record.weight)
        weights[label] = weights.get(label, 0) + 1

    print(
        json.dumps(
            {
                "status": "PASS",
                "cache": cache.path.name,
                "draw_start": args.draw_start,
                "draw_count": args.draw_count,
                "dataset_counts": counts,
                "weight_counts": weights,
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()