#!/usr/bin/env python3
"""
Caption ALL scene=1 images for a relation using LLaVA with DDP.

Each GPU loads its own model copy and processes a shard.
Scene=0 images get an empty caption (no LLaVA needed).
Incrementally checkpoints after each batch so runs can be resumed.

Usage (single GPU):
    CUDA_VISIBLE_DEVICES=0 python EFUF/scripts/caption_ddp.py \\
        --relation bathroom_toilet \\
        --output EFUF/data/bathroom_toilet/all_captions.json

Usage (multi-GPU via torchrun):
    CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 \\
        EFUF/scripts/caption_ddp.py \\
        --relation bathroom_toilet \\
        --output EFUF/data/bathroom_toilet/all_captions.json
"""

from __future__ import annotations

import argparse
import gc
import json
import os
import sys

import torch
import torch.distributed as dist
from tqdm import tqdm

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
from experiment.config.relation_config import get_relation_config

PROMPT = "Describe this image."
LLAVA_MODEL = "llava-hf/llava-1.5-7b-hf"
DEFAULT_BATCH_SIZE = 8


def _clear_gpu(device: torch.device) -> None:
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.synchronize()


def load_hf_dataset(dataset_id: str):
    from datasets import load_dataset
    return load_dataset(dataset_id)


def build_existing_caption_map(checkpoint_path: str, output_path: str | None = None) -> dict[str, str]:
    cmap: dict[str, str] = {}

    if os.path.exists(checkpoint_path):
        with open(checkpoint_path) as f:
            ckpt = json.load(f)
        for split in ("train", "val"):
            ids = ckpt.get(f"{split}_ids", [])
            caps = ckpt.get(f"llava_{split}", [])
            for iid, cap in zip(ids, caps):
                cmap[str(iid)] = cap
            pos_ids = ckpt.get(f"{split}_pos_ids", [])
            pos_caps = ckpt.get(f"llava_{split}_pos", [])
            if pos_caps:
                for iid, cap in zip(pos_ids, pos_caps):
                    cmap[str(iid)] = cap

    if output_path and os.path.exists(output_path):
        with open(output_path) as f:
            for entry in json.load(f):
                iid = str(entry["image_id"])
                cap = entry.get("llava_caption", "")
                if cap:
                    cmap[iid] = cap

    return cmap


def load_ckpt(ckpt_path: str) -> dict[str, str]:
    if not os.path.exists(ckpt_path):
        return {}
    cmap: dict[str, str] = {}
    with open(ckpt_path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            entry = json.loads(line)
            iid = entry["image_id"]
            cap = entry.get("llava_caption", "")
            if cap:
                cmap[iid] = cap
    return cmap


def append_ckpt(ckpt_path: str, results: dict[str, str], scene_col: str, object_col: str, meta: dict[str, dict]):
    with open(ckpt_path, "a") as f:
        for iid, cap in results.items():
            m = meta.get(iid, {"scene": 1, "obj": 1})
            f.write(json.dumps({
                "image_id": iid,
                "llava_caption": cap,
                scene_col: m["scene"],
                object_col: m["obj"],
                "edited_caption": None,
            }, ensure_ascii=False) + "\n")


def infer_shard(
    image_ids: list[str],
    processor,
    model: torch.nn.Module,
    device: torch.device,
    batch_size: int,
    desc: str,
    rank: int,
    hf_images: dict[str, "Image.Image"],
    ckpt_path: str | None,
    scene_col: str,
    object_col: str,
    meta: dict[str, dict],
) -> dict[str, str]:
    from PIL import Image

    results: dict[str, str] = {}
    all_imgs: dict[str, Image.Image] = {}
    for iid in image_ids:
        if iid in hf_images:
            all_imgs[iid] = hf_images[iid]

    available_ids = [iid for iid in image_ids if iid in all_imgs]
    if not available_ids:
        return results

    pbar_len = (len(available_ids) + batch_size - 1) // batch_size
    pbar = tqdm(total=pbar_len, desc=f"[GPU{rank}] {desc}", position=rank)
    for i in range(0, len(available_ids), batch_size):
        batch_ids = available_ids[i : i + batch_size]
        imgs = [all_imgs[iid] for iid in batch_ids]
        texts = [f"USER: <image>\n{PROMPT} ASSISTANT:" for _ in batch_ids]

        inputs = processor(text=texts, images=imgs, return_tensors="pt", padding=True)
        inputs = {k: v.to(device) if hasattr(v, "to") else v for k, v in inputs.items()}

        with torch.inference_mode():
            out = model.generate(**inputs, max_new_tokens=150, do_sample=False, use_cache=True)

        inp_len = inputs["input_ids"].shape[1]
        batch_results: dict[str, str] = {}
        for iid, seq in zip(batch_ids, out):
            caption = processor.decode(seq[inp_len:], skip_special_tokens=True).strip()
            results[iid] = caption
            batch_results[iid] = caption
        pbar.update(1)

        if ckpt_path:
            append_ckpt(ckpt_path, batch_results, scene_col, object_col, meta)

    pbar.close()
    return results


def worker(rank: int, world_size: int, args):
    local_rank = int(os.environ.get("LOCAL_RANK", rank))
    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)

    if rank == 0:
        print(f"[caption_ddp] relation={args.relation}")
        print(f"[caption_ddp] world_size={world_size}, batch_size={args.batch_size} per GPU")
        print(f"[caption_ddp] output={args.output}")

    rc = get_relation_config(args.relation)
    scene_col = rc.scene_key
    object_col = rc.object_key
    dataset_id = rc.dataset_id

    data_dir = os.path.join(
        os.path.dirname(os.path.abspath(__file__)), "..", "..",
        "VisEdit", "data", "hallucination", args.relation,
    )
    ckpt_path = args.checkpoint or os.path.join(data_dir, "checkpoint.json")
    incremental_ckpt = args.output + ".ckpt.jsonl"

    if rank == 0:
        print(f"[caption_ddp] Loading existing captions from {ckpt_path} ...")
    existing_captions = build_existing_caption_map(ckpt_path, args.output)

    if os.path.exists(incremental_ckpt):
        incremental_captions = load_ckpt(incremental_ckpt)
        existing_captions.update(incremental_captions)
        if rank == 0:
            print(f"[caption_ddp] Resumed {len(incremental_captions)} captions from incremental checkpoint")

    if rank == 0:
        print(f"[caption_ddp] Existing captions: {len(existing_captions)}")

    if rank == 0:
        print(f"[caption_ddp] Loading HF dataset {dataset_id} ...")
    ds = load_hf_dataset(dataset_id)

    missing_train: list[str] = []
    missing_val: list[str] = []
    missing_meta: dict[str, dict] = {}
    all_entries: list[dict] = []
    hf_images: dict[str, "Image.Image"] = {}

    for split_name in ("train", "val"):
        hf_split = "validation" if split_name == "val" else split_name
        split = ds[hf_split] if hf_split in ds else ds.get(split_name, ds.get("test", []))

        for item in split:
            scene = int(item[scene_col])
            obj = int(item[object_col])
            iid = str(item["image_id"])

            if scene == 1 and "image" in item:
                try:
                    hf_images[iid] = item["image"].convert("RGB")
                except Exception:
                    pass

            cap = existing_captions.get(iid, "")
            if cap:
                all_entries.append({
                    "image_id": iid,
                    "llava_caption": cap,
                    scene_col: scene,
                    object_col: obj,
                    "edited_caption": None,
                })
            elif scene == 1:
                missing_meta[iid] = {"scene": scene, "obj": obj, "split": split_name}
                if split_name == "train":
                    missing_train.append(iid)
                else:
                    missing_val.append(iid)
            else:
                all_entries.append({
                    "image_id": iid,
                    "llava_caption": "",
                    scene_col: scene,
                    object_col: obj,
                    "edited_caption": None,
                })

    shard_train = missing_train[rank::world_size]
    shard_val = missing_val[rank::world_size]
    total_missing = len(missing_train) + len(missing_val)

    if rank == 0:
        print(f"[caption_ddp] Missing captions: {total_missing} ({len(missing_train)} train + {len(missing_val)} val)")
        print(f"[caption_ddp] Each GPU gets ~{len(shard_train)} train + ~{len(shard_val)} val")

    if total_missing == 0:
        if rank == 0:
            if os.path.exists(incremental_ckpt):
                os.remove(incremental_ckpt)
            os.makedirs(os.path.dirname(args.output), exist_ok=True)
            with open(args.output, "w") as f:
                json.dump(all_entries, f, indent=2, ensure_ascii=False)
            print(f"[caption_ddp] Done! Saved {len(all_entries)} entries to {args.output}")
        if world_size > 1:
            dist.destroy_process_group()
        return

    if rank == 0:
        print(f"[caption_ddp] Loading LLaVA {args.llava_model} on each GPU ...")

    from transformers import LlavaForConditionalGeneration, AutoProcessor

    processor = AutoProcessor.from_pretrained(args.llava_model)
    processor.tokenizer.padding_side = "left"

    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    model = LlavaForConditionalGeneration.from_pretrained(
        args.llava_model,
        torch_dtype=torch.bfloat16,
        device_map=None,
        attn_implementation="sdpa",
    ).to(device)
    model.eval()

    os.makedirs(os.path.dirname(args.output), exist_ok=True)

    new_caps: dict[str, str] = {}
    if shard_train:
        new_caps.update(infer_shard(
            shard_train, processor, model, device,
            args.batch_size, "train", rank, hf_images,
            incremental_ckpt, scene_col, object_col, missing_meta,
        ))
    if shard_val:
        new_caps.update(infer_shard(
            shard_val, processor, model, device,
            args.batch_size, "val", rank, hf_images,
            incremental_ckpt, scene_col, object_col, missing_meta,
        ))

    tmp_path = f"{args.output}.rank{rank}.tmp"
    with open(tmp_path, "w") as f:
        json.dump(new_caps, f, indent=2, ensure_ascii=False)

    del model, processor
    _clear_gpu(device)

    if world_size > 1:
        dist.barrier()

    if rank == 0:
        for r in range(world_size):
            tmp = f"{args.output}.rank{r}.tmp"
            if os.path.exists(tmp):
                with open(tmp) as f:
                    shard_caps = json.load(f)
                for iid, cap in shard_caps.items():
                    meta = missing_meta.get(iid, {"scene": 1, "obj": 1})
                    all_entries.append({
                        "image_id": iid,
                        "llava_caption": cap,
                        scene_col: meta["scene"],
                        object_col: meta["obj"],
                        "edited_caption": None,
                    })
                os.remove(tmp)

        if os.path.exists(incremental_ckpt):
            os.remove(incremental_ckpt)
        with open(args.output, "w") as f:
            json.dump(all_entries, f, indent=2, ensure_ascii=False)
        print(f"[caption_ddp] Done! Saved {len(all_entries)} entries to {args.output}")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--relation", required=True)
    ap.add_argument("--output", required=True)
    ap.add_argument("--batch_size", type=int, default=DEFAULT_BATCH_SIZE)
    ap.add_argument("--llava_model", default=LLAVA_MODEL)
    ap.add_argument("--checkpoint", default=None)
    ap.add_argument("--local-rank", type=int, default=0)
    args = ap.parse_args()

    if "RANK" in os.environ:
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        dist.init_process_group("nccl")
        worker(rank, world_size, args)
        dist.destroy_process_group()
    else:
        worker(0, 1, args)