#!/usr/bin/env python3
"""
Convert bathroom_toilet/all_captions.json + checkpoint data into EFUF training format.

Produces:
  - pos_neg.json: Per-object/subsentence entries for positive and negative training
  - sentences.json: Whole-caption positive reinforcement entries
  - dummy_vqa.json: Minimal VQA data for EFUF DataLoader compatibility

Usage:
    python scripts/format_bathroom_data.py \
        --all_captions data/bathroom_toilet/all_captions.json \
        --ckpt_json data/bathroom_toilet/all_captions.json.ckpt.json \
        --visedit_ckpt ../VisEdit/data/hallucination/bathroom_toilet/checkpoint.json \
        --image_dir ../VisEdit/data/hallucination/bathroom_toilet/images \
        --output_dir data/bathroom_toilet \
        --mode synthetic
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
from pathlib import Path


SUBSENTENCE_SPLITTER = ",.;!?:"

TARGET_WORDS = ["toilet", "toilets"]

NEGATIVE_SCORE = 0.0
POSITIVE_SCORE = 40.0
SENTENCE_MEAN_POS = 40.0
SENTENCE_MIN_POS = 35.0
SENTENCE_MEAN_NEG = 10.0
SENTENCE_MIN_NEG = 5.0


def find_subsentence(text: str, target_words: list[str]) -> tuple[int, int] | None:
    """Find the subsentence containing any of the target words.

    Returns (start, end) character indices of the subsentence within text,
    or None if not found.
    """
    sub_sentences = re.split(f"[{SUBSENTENCE_SPLITTER}]+", text)
    for sub in sub_sentences:
        for word in target_words:
            pattern = r"\b" + re.escape(word) + r"\b"
            if re.search(pattern, sub, re.IGNORECASE):
                start = text.find(sub)
                end = start + len(sub)
                return start, end
    return None


def is_valid_toilet_mention(caption: str) -> bool:
    """Check if 'toilet' (or 'toilets') appears as a whole word, not as part
    of 'toiletries' etc."""
    lower = caption.lower()
    for word in TARGET_WORDS:
        pattern = r"\b" + re.escape(word) + r"\b"
        if re.search(pattern, lower):
            return True
    return False


def main():
    parser = argparse.ArgumentParser(description="Format bathroom_toilet data for EFUF training")
    parser.add_argument("--all_captions", default="data/bathroom_toilet/all_captions.json")
    parser.add_argument("--ckpt_json", default="data/bathroom_toilet/all_captions.json.ckpt.json")
    parser.add_argument(
        "--visedit_ckpt", default="../VisEdit/data/hallucination/bathroom_toilet/checkpoint.json"
    )
    parser.add_argument(
        "--image_dir", default="../VisEdit/data/hallucination/bathroom_toilet/images"
    )
    parser.add_argument("--output_dir", default="data/bathroom_toilet")
    parser.add_argument(
        "--mode",
        choices=["synthetic", "clip"],
        default="synthetic",
        help="synthetic: use fixed scores; clip: compute CLIP scores (not implemented yet)",
    )
    parser.add_argument("--unlearn_target", choices=["subsentence", "object"], default="subsentence")
    parser.add_argument("--symlink_dir", default=None, help="Create merged image directory with symlinks")
    args = parser.parse_args()

    with open(args.all_captions) as f:
        all_captions = json.load(f)

    with open(args.ckpt_json) as f:
        step_data = json.load(f)

    with open(args.visedit_ckpt) as f:
        visedit_ckpt = json.load(f)

    step2_judge = step_data.get("step2_judge_done", {})
    step3_rewrite = step_data.get("step3_rewrite_done", {})
    step4_strict = step_data.get("step4_strict_done", {})

    train_ids = set(visedit_ckpt["train_ids"])
    val_ids = set(visedit_ckpt["val_ids"])
    train_pos_ids = set(visedit_ckpt.get("train_pos_ids", []))
    val_pos_ids = set(visedit_ckpt.get("val_pos_ids", []))

    id_map = {item["image_id"]: item for item in all_captions}

    pos_neg_entries = []
    sentence_entries = []
    skipped_no_toilet = 0
    skipped_no_subsentence = 0
    skipped_no_image = 0

    train_img_dir = os.path.join(args.image_dir, "train")
    val_img_dir = os.path.join(args.image_dir, "val")

    def get_image_path(image_id: str) -> str | None:
        for d in [train_img_dir, val_img_dir]:
            p = os.path.join(d, f"{image_id}.jpg")
            if os.path.exists(p):
                return p
        return None

    def get_image_relpath(image_id: str) -> str:
        """Return relative path from image_dir for the image."""
        if image_id in train_ids or image_id in train_pos_ids:
            return f"train/{image_id}.jpg"
        elif image_id in val_ids or image_id in val_pos_ids:
            return f"val/{image_id}.jpg"
        else:
            for d in ["train", "val"]:
                if os.path.exists(os.path.join(args.image_dir, d, f"{image_id}.jpg")):
                    return f"{d}/{image_id}.jpg"
            return f"train/{image_id}.jpg"

    for item in all_captions:
        image_id = item["image_id"]
        caption = item["llava_caption"]
        toilet = item["toilet"]
        bathroom = item["bathroom"]

        if not is_valid_toilet_mention(caption):
            if toilet == 0 and bathroom == 1:
                sentence_entries.append(
                    {
                        "image": get_image_relpath(image_id),
                        "sentence": caption,
                        "mean": SENTENCE_MEAN_POS,
                        "min": SENTENCE_MIN_POS,
                    }
                )
            elif toilet == 1 and bathroom == 1:
                sentence_entries.append(
                    {
                        "image": get_image_relpath(image_id),
                        "sentence": caption,
                        "mean": SENTENCE_MEAN_POS,
                        "min": SENTENCE_MIN_POS,
                    }
                )
            continue

        if not get_image_path(image_id):
            skipped_no_image += 1
            continue

        if toilet == 0:
            # Check if this is a confirmed hallucination
            judgment = step2_judge.get(image_id)
            if judgment != "YES":
                skipped_no_toilet += 1
                continue

            result = find_subsentence(caption, TARGET_WORDS)
            if result is None:
                skipped_no_subsentence += 1
                continue
            start, end = result

            if args.unlearn_target == "subsentence":
                sentence = caption[:end].rstrip()
                position = start
                while position > 0 and sentence[position - 1] == " ":
                    position -= 1
                if position == 0:
                    position = start
            else:
                word_match = re.search(r"\b(toilet|toilets)\b", caption, re.IGNORECASE)
                if word_match is None:
                    skipped_no_subsentence += 1
                    continue
                start_pos = word_match.start()
                end_pos = word_match.end()
                sentence = caption[:end_pos]
                position = start_pos

            pos_neg_entries.append(
                {
                    "image": get_image_relpath(image_id),
                    "sentence": sentence,
                    "position": position,
                    "score": NEGATIVE_SCORE,
                    "type": "negative",
                }
            )

        elif toilet == 1:
            result = find_subsentence(caption, TARGET_WORDS)
            if result is None:
                skipped_no_subsentence += 1
                continue
            start, end = result

            if args.unlearn_target == "subsentence":
                sentence = caption[:end].rstrip()
                position = start
                while position > 0 and sentence[position - 1] == " ":
                    position -= 1
                if position == 0:
                    position = start
            else:
                word_match = re.search(r"\b(toilet|toilets)\b", caption, re.IGNORECASE)
                if word_match is None:
                    skipped_no_subsentence += 1
                    continue
                start_pos = word_match.start()
                end_pos = word_match.end()
                sentence = caption[:end_pos]
                position = start_pos

            pos_neg_entries.append(
                {
                    "image": get_image_relpath(image_id),
                    "sentence": sentence,
                    "position": position,
                    "score": POSITIVE_SCORE,
                    "type": "positive",
                }
            )

            sentence_entries.append(
                {
                    "image": get_image_relpath(image_id),
                    "sentence": caption,
                    "mean": SENTENCE_MEAN_POS,
                    "min": SENTENCE_MIN_POS,
                }
            )

    neg_count = sum(1 for e in pos_neg_entries if e["type"] == "negative")
    pos_count = sum(1 for e in pos_neg_entries if e["type"] == "positive")

    print(f"pos_neg.json: {len(pos_neg_entries)} entries ({neg_count} negative, {pos_count} positive)")
    print(f"sentences.json: {len(sentence_entries)} entries")
    print(f"Skipped: no_toilet_mention={skipped_no_toilet}, no_subsentence={skipped_no_subsentence}, no_image={skipped_no_image}")

    final_pos_neg = [
        {k: v for k, v in e.items() if k != "type"} for e in pos_neg_entries
    ]

    all_ids = train_ids | val_ids | train_pos_ids | val_pos_ids

    def split(entries):
        train_entries = [e for e in entries if e["image"].split("/")[0] == "train"]
        val_entries = [e for e in entries if e["image"].split("/")[0] == "val"]
        return train_entries, val_entries

    train_pos_neg, val_pos_neg = split(final_pos_neg)
    train_sentences, val_sentences = split(sentence_entries)

    print(f"Train pos_neg: {len(train_pos_neg)}, Val pos_neg: {len(val_pos_neg)}")
    train_neg = sum(1 for e in train_pos_neg if e["score"] == NEGATIVE_SCORE)
    train_pos = sum(1 for e in train_pos_neg if e["score"] == POSITIVE_SCORE)
    val_neg = sum(1 for e in val_pos_neg if e["score"] == NEGATIVE_SCORE)
    val_pos = sum(1 for e in val_pos_neg if e["score"] == POSITIVE_SCORE)
    print(f"  Train: {train_pos} positive, {train_neg} negative")
    print(f"  Val: {val_pos} positive, {val_neg} negative")
    print(f"Train sentences: {len(train_sentences)}, Val sentences: {len(val_sentences)}")

    os.makedirs(args.output_dir, exist_ok=True)

    for split_name, data in [("train", train_pos_neg), ("val", val_pos_neg)]:
        path = os.path.join(args.output_dir, f"pos_neg_{args.mode}_{split_name}.json")
        with open(path, "w") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
        print(f"Saved {path} ({len(data)} entries)")

    for split_name, data in [("train", train_sentences), ("val", val_sentences)]:
        path = os.path.join(args.output_dir, f"sentences_{args.mode}_{split_name}.json")
        with open(path, "w") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
        print(f"Saved {path} ({len(data)} entries)")

    pos_neg_path = os.path.join(args.output_dir, f"pos_neg_{args.mode}.json")
    with open(pos_neg_path, "w") as f:
        json.dump(final_pos_neg, f, indent=2, ensure_ascii=False)
    print(f"Saved {pos_neg_path}")

    sentences_path = os.path.join(args.output_dir, f"sentences_{args.mode}.json")
    with open(sentences_path, "w") as f:
        json.dump(sentence_entries, f, indent=2, ensure_ascii=False)
    print(f"Saved {sentences_path}")

    # Create a minimal dummy VQA dataset so the GoldData DataLoader doesn't crash
    dummy_vqa = [
        {
            "id": "dummy_001",
            "image": get_image_relpath(all_captions[0]["image_id"]),
            "conversations": [
                {"from": "human", "value": "<image>\nWhat is in this image?"},
                {"from": "gpt", "value": all_captions[0]["llava_caption"][:100]},
            ],
        }
    ]
    vqa_path = os.path.join(args.output_dir, "dummy_vqa.json")
    with open(vqa_path, "w") as f:
        json.dump(dummy_vqa, f, indent=2)
    print(f"Saved {vqa_path}")

    if args.symlink_dir:
        os.makedirs(args.symlink_dir, exist_ok=True)
        for entry in pos_neg_entries + sentence_entries:
            src = get_image_path(entry["image"].split("/")[-1].replace(".jpg", ""))
            if src and os.path.exists(src):
                dst = os.path.join(args.symlink_dir, os.path.basename(src))
                if not os.path.exists(dst):
                    os.symlink(os.path.abspath(src), dst)
        print(f"Created symlink directory at {args.symlink_dir}")


if __name__ == "__main__":
    main()