import json
from pathlib import Path

infile = Path("/data16T/chestx-ray/AURAD_dataset/chexlocalize/test_prompt_text2layout_multi.json")
outfile = Path("/data16T/chestx-ray/AURAD_dataset/chexlocalize/test_prompt_text2layout_multi_no_device.json")

removed_attn_count = 0
modified_prompt_count = 0
removed_empty_prompt_count = 0
kept_count = 0

def remove_support_from_prompt(prompt: str) -> str:
    if not prompt:
        return ""
    parts = [x.strip() for x in prompt.split(",")]
    new_parts = [x for x in parts if x and "Support Devices" not in x]
    return ", ".join(new_parts)

with infile.open("r", encoding="utf-8") as f_in, outfile.open("w", encoding="utf-8") as f_out:
    for line in f_in:
        line = line.strip()
        if not line:
            continue

        item = json.loads(line)

        old_attn_list = item.get("attn_list", [])
        new_attn_list = []
        for disease, mask_path in old_attn_list:
            if disease == "Support Devices":
                removed_attn_count += 1
            else:
                new_attn_list.append([disease, mask_path])
        item["attn_list"] = new_attn_list

        old_prompt = item.get("prompt", "")
        new_prompt = remove_support_from_prompt(old_prompt)
        if new_prompt != old_prompt:
            modified_prompt_count += 1
        item["prompt"] = new_prompt

        # 删除 prompt 为空的样本
        if item["prompt"] == "":
            removed_empty_prompt_count += 1
            continue

        f_out.write(json.dumps(item, ensure_ascii=False) + "\n")
        kept_count += 1

print(f"saved to: {outfile}")
print(f"kept samples: {kept_count}")
print(f"removed 'Support Devices' in attn_list: {removed_attn_count}")
print(f"modified prompts: {modified_prompt_count}")
print(f"removed empty-prompt samples: {removed_empty_prompt_count}")