"""RePS-style preference suppression: dataset + bounded preference loss.

Replaces the EFUF targeted-NLL *gradient ascent* (``efuf_targeted_nll_loss``),
which is unbounded and unstable, with a reference-free, length-normalized
bidirectional preference objective in the spirit of RePS (arXiv:2505.20809)
and SimPO. The model is trained to PREFER a toilet-free caption (winning) over
a toilet-containing caption (losing) for the same image. Because the objective
is ``-log σ(β·(lp_w − lp_l) − γ)`` it saturates once the clean caption is
sufficiently preferred, so it cannot blow up the way NLL ascent does.

This acts at the OUTPUT (token distribution) level, complementing the per-layer
attention probe which suppresses at the activation level.

Preference pairs are mined from the EFUF synthetic data
(``pos_neg_synthetic_train.json``): each hallucination negative (score below the
CLIP threshold) carries a character-level ``position`` marking where the
hallucinated clause begins, so:
    winning (clean)      = sentence[:position]   (toilet-free prefix)
    losing (hallucinated)= sentence              (toilet clause appended)
The two share the same prefix; only the trailing hallucinated clause differs.
"""
from __future__ import annotations

import json as _json
import re
from pathlib import Path as _Path

import torch
import torch.nn.functional as F
from PIL import Image as _PILImage


def _build_object_regex(object_keywords) -> "re.Pattern":
    """Case-insensitive regex matching any object keyword (e.g. oven|stove|range)."""
    kws = object_keywords or ["toilet"]
    stems = sorted({re.escape(k.lower()) for k in kws}, key=len, reverse=True)
    return re.compile("|".join(stems), re.IGNORECASE)


class PreferenceDataset(torch.utils.data.Dataset):
    """(image, winning clean caption, losing hallucinated caption) triples.

    Mirrors the image-loading + tokenization conventions of
    ``finetune_adv_gen_resume.EFUFNegativeDataset`` so the same HF dataset
    (looked up by ``image_id``) supplies the images.
    """

    def __init__(
        self,
        processor,
        data_path: str,
        image_dir: str | None = None,
        dataset_id: str | None = None,
        hal_clip_thres: float = 23.0,
        min_position: int = 5,
        require_object_in_tail: bool = True,
        prompt: str = "Describe this image.",
        max_length: int = 1024,
        object_keywords: list[str] | None = None,
    ):
        obj_re = _build_object_regex(object_keywords)
        with open(data_path) as f:
            raw = _json.load(f)
        self.samples: list[dict] = []
        self._hf_lookup: dict | None = None
        self._image_dir: _Path | None = None

        if dataset_id:
            from datasets import load_dataset as _load_ds
            _ds_all = _load_ds(dataset_id)
            self._hf_lookup = {}
            for split in _ds_all:
                for item in _ds_all[split]:
                    self._hf_lookup[item["image_id"]] = item["image"]
            print(f"  PreferenceDataset: built HF image lookup with "
                  f"{len(self._hf_lookup)} entries from {dataset_id}")
        elif image_dir:
            self._image_dir = _Path(image_dir)

        for d in raw:
            if d["score"] >= hal_clip_thres:
                continue  # keep only genuine hallucinations
            pos = int(d["position"])
            if pos < min_position:
                continue
            sentence = d["sentence"]
            context = sentence[:pos]
            tail = sentence[pos:]
            if require_object_in_tail and not obj_re.search(tail):
                continue
            if obj_re.search(context):
                continue  # clean prefix must be object-free
            image_id = (
                d["image"].split("/")[-1]
                .replace(".jpg", "").replace(".png", "").replace(".jpeg", "")
            )
            if self._hf_lookup is not None:
                if image_id not in self._hf_lookup:
                    continue
            elif self._image_dir is not None:
                if not (self._image_dir / d["image"]).exists():
                    continue
            self.samples.append({
                "image_id": image_id,
                "win": context,        # clean caption
                "lose": sentence,      # hallucinated caption
            })

        self.processor = processor
        self.prompt = prompt
        self.max_length = max_length
        source = f"HF({dataset_id})" if dataset_id else f"local({image_dir})"
        print(f"PreferenceDataset: {len(self.samples)} preference pairs "
              f"(score < {hal_clip_thres}, position >= {min_position}) from {source}")

    def __len__(self):
        return len(self.samples)

    def _load_image(self, image_id: str):
        if self._hf_lookup is not None:
            return self._hf_lookup[image_id].convert("RGB")
        return _PILImage.open(str(self._image_dir / (image_id + ".jpg"))).convert("RGB")

    def _encode(self, image, response_text: str) -> dict:
        """Tokenize prefix+response with the image; mark response tokens (tail)."""
        prefix = f"<image>\nUSER: {self.prompt}\nASSISTANT: "
        full_text = prefix + response_text
        enc = self.processor(
            images=image, text=full_text, return_tensors="pt",
            padding="max_length", max_length=self.max_length, truncation=True,
        )
        input_ids = enc["input_ids"][0]
        attn = enc["attention_mask"][0]
        tok = self.processor.tokenizer
        n_full = len(tok.encode(full_text, add_special_tokens=True))
        n_prefix = len(tok.encode(prefix, add_special_tokens=True))
        resp_len = max(n_full - n_prefix, 0)
        S = input_ids.shape[0]
        real_len = int(attn.sum().item())
        resp_mask = torch.zeros(S, dtype=torch.bool)
        # Left-padded: real tokens occupy the tail; response is the last resp_len.
        resp_start = S - resp_len
        if 0 < resp_start < S and resp_len <= real_len:
            resp_mask[resp_start:] = True
        return {
            "pixel_values": enc["pixel_values"][0],
            "input_ids": input_ids,
            "attention_mask": attn,
            "response_mask": resp_mask,
        }

    def __getitem__(self, idx):
        s = self.samples[idx]
        image = self._load_image(s["image_id"])
        win = self._encode(image, s["win"])
        lose = self._encode(image, s["lose"])
        return {
            "win_pixel_values": win["pixel_values"],
            "win_input_ids": win["input_ids"],
            "win_attention_mask": win["attention_mask"],
            "win_response_mask": win["response_mask"],
            "lose_pixel_values": lose["pixel_values"],
            "lose_input_ids": lose["input_ids"],
            "lose_attention_mask": lose["attention_mask"],
            "lose_response_mask": lose["response_mask"],
        }


def sequence_logprob(
    logits: torch.Tensor,          # (B, S_full, V) — model output
    input_ids: torch.Tensor,       # (B, S_text)
    response_mask: torch.Tensor,   # (B, S_text) bool — True at response tokens
    n_image_patches: int,
) -> torch.Tensor:
    """Length-normalized mean log-prob of the response tokens, per row (B,).

    Alignment mirrors ``efuf_targeted_nll_loss``: LLaVA expands the single
    ``<image>`` token into ``n_image_patches`` embeddings, so we front-pad the
    text tensors by ``n_image_patches - 1`` to match the full logits length,
    then apply the causal shift (logit at j predicts token j+1).
    """
    B, S_full, V = logits.shape
    text_offset = n_image_patches - 1
    labels_full = F.pad(input_ids.long(), (text_offset, 0), value=-100)   # (B, S_full)
    resp_full = F.pad(response_mask, (text_offset, 0), value=False)       # (B, S_full)
    labels_full = labels_full.masked_fill(~resp_full, -100)

    logp = F.log_softmax(logits[:, :-1, :].float(), dim=-1)               # (B, S_full-1, V)
    tgt = labels_full[:, 1:]                                              # (B, S_full-1)
    valid = tgt != -100
    gathered = logp.gather(-1, tgt.clamp(min=0).unsqueeze(-1)).squeeze(-1)  # (B, S_full-1)
    gathered = gathered * valid.float()
    counts = valid.float().sum(dim=-1).clamp(min=1.0)
    return gathered.sum(dim=-1) / counts                                  # (B,)


def reps_preference_loss(
    lp_win: torch.Tensor,    # (B,) length-normalized logprob of clean caption
    lp_lose: torch.Tensor,   # (B,) length-normalized logprob of hallucinated caption
    beta: float = 2.0,
    gamma: float = 0.5,
) -> torch.Tensor:
    """Reference-free SimPO/RePS preference loss: prefer clean over hallucinated.

    L = -mean[ log σ(β·(lp_win − lp_lose) − γ) ]. Bounded and saturating, so it
    is stable to optimize (unlike unbounded NLL gradient ascent).
    """
    margin = beta * (lp_win - lp_lose) - gamma
    return -F.logsigmoid(margin).mean()