import os
from tqdm import tqdm
import torch
from transformers import BertTokenizer
import sys   


PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

from Rad_IR.radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box


def main():
    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

    dataset = HyCoClipMimicBoxDataset_JPG(
        split="train",
        tokenizer=tokenizer,
        dataset_root="/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset",
        cig_matched_dir="/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/processed/matched",
        caption_file="valid.jsonl",
        transform=None,
    )

    bad_empty = 0
    bad_short = 0
    bad_phrase = 0
    bad_image = 0

    for i in tqdm(range(len(dataset))):
        try:
            img, caps, cap_len, img_path, phrase, ph_mask, bbx_mask = dataset[i]
        except Exception as e:
            print(f"\n❌ load error at {i}: {e}")
            bad_image += 1
            continue

        input_ids = caps["input_ids"][0]
        attention_mask = caps["attention_mask"][0]

        # 1️⃣ attention_mask 全 0
        if attention_mask.sum().item() == 0:
            bad_empty += 1
            print("\n❌ Empty caption")
            print("index:", i)
            print("img_path:", img_path)
            continue

        # 2️⃣ token 太少
        if (input_ids != 0).sum().item() < 3:
            bad_short += 1
            print("\n⚠️ Too short caption")
            print("index:", i)
            print("img_path:", img_path)

        # 3️⃣ phrase / mask 不一致
        if phrase == "" and ph_mask.sum().item() != 0:
            bad_phrase += 1
            print("\n❌ phrase empty but mask not zero")
            print("index:", i)
            print("img_path:", img_path)

        if phrase != "" and ph_mask.sum().item() == 0:
            bad_phrase += 1
            print("\n❌ phrase exists but mask all zero")
            print("index:", i)
            print("phrase:", phrase)
            print("img_path:", img_path)

        # 4️⃣ image NaN / Inf（保险）
        if not torch.isfinite(img).all():
            bad_image += 1
            print("\n❌ NaN/Inf in image")
            print("index:", i)
            print("img_path:", img_path)

    print("\n========== SUMMARY ==========")
    print("Total samples:", len(dataset))
    print("Empty captions:", bad_empty)
    print("Too short captions:", bad_short)
    print("Phrase issues:", bad_phrase)
    print("Bad images:", bad_image)

if __name__ == "__main__":
    main()