"""
Build the original-model output cache from caption_targets.json.

Produces a single all-inclusive cache file:
    <output_dir>/original_outputs_all.json

validate.py will load this file and filter entries by the current run's
image_ids and prompts, so there is no need to maintain a separate cache
for every (num_per_category, num_prompts) combination.

Usage:
    python -m experiment.data.build_original_cache \\
        --caption_targets experiment/data/caption_targets.json \\
        --output_dir ./cached_original_outputs

    # Include all splits (default is val only):
    python -m experiment.data.build_original_cache \\
        --caption_targets experiment/data/caption_targets.json \\
        --output_dir ./cached_original_outputs \\
        --all_splits
"""

import argparse
import json
import os

CATEGORIES = [
    "bathroom_no_toilet",
    "bathroom_with_toilet",
    "non_bathroom_with_toilet",
    "unrelated",
]

# Add any prompts you want pre-cached here.
DEFAULT_PROMPTS = [
    "Describe this image.",
    "Give a detailed description of this image.",
]


def build_cache(caption_targets_path, output_dir, prompts, use_val_split):
    with open(caption_targets_path) as f:
        targets = json.load(f)

    images = targets["images"]

    # Build cache in validate.py's collect_outputs_transformers format:
    # { category: [{image_id, image_path, prompt, text}, ...] }
    cache = {cat: [] for cat in CATEGORIES}
    for image_id, entry in images.items():
        caption = entry.get("original_caption")
        if not caption:
            continue
        if use_val_split and entry.get("split") != "val":
            continue
        cat = entry.get("category")
        if cat not in CATEGORIES:
            continue
        stored_path = entry.get("image_path", image_id)
        # Discard stale temp paths (multi-GPU inference artefacts) or
        # paths that no longer exist on disk; fall back to image_id so
        # validate.py's image_lookup can resolve the image at eval time.
        if stored_path.startswith("/tmp/") or (
            stored_path != image_id and not os.path.exists(stored_path)
        ):
            stored_path = image_id
        for prompt in prompts:
            cache[cat].append({
                "image_id": image_id,
                "image_path": stored_path,
                "prompt": prompt,
                "text": caption,
            })

    for cat in CATEGORIES:
        n_images = len({e["image_id"] for e in cache[cat]})
        print(f"  {cat}: {n_images} images × {len(prompts)} prompts"
              f" = {len(cache[cat])} entries")

    os.makedirs(output_dir, exist_ok=True)
    cache_file = os.path.join(output_dir, "original_outputs_all.json")
    with open(cache_file, "w") as f:
        json.dump(cache, f, indent=2)
    print(f"\nSaved cache to {cache_file}")
    return cache_file


def main():
    parser = argparse.ArgumentParser(
        description="Build original-model cache from caption_targets.json"
    )
    parser.add_argument("--caption_targets", type=str,
                        default="experiment/data/caption_targets.json")
    parser.add_argument("--output_dir", type=str,
                        default="./cached_original_outputs")
    parser.add_argument("--prompts", nargs="+", default=DEFAULT_PROMPTS,
                        help="Prompts to include in the cache")
    parser.add_argument("--use_val_split", action="store_true", default=True,
                        help="Only use val-split images (default: true)")
    parser.add_argument("--all_splits", action="store_true",
                        help="Use all splits, not just val")
    args = parser.parse_args()

    use_val_split = args.use_val_split and not args.all_splits

    print(f"Building original cache from {args.caption_targets}")
    print(f"  prompts ({len(args.prompts)}): {args.prompts}")
    print(f"  val split only: {use_val_split}")
    print()

    build_cache(
        caption_targets_path=args.caption_targets,
        output_dir=args.output_dir,
        prompts=args.prompts,
        use_val_split=use_val_split,
    )


if __name__ == "__main__":
    main()