"""
bert/dataset.py
Cross-Encoder dataset for entailment verification.

Input format: (premise=retrieved_chunk, hypothesis=generated_claim)
Labels: 0=Contradiction, 1=Neutral, 2=Entailment

Data sources:
  - ANLI (Adversarial NLI)  — hard human-adversarial examples
  - TrueTeacher (Google, 1.4M) — LLM summaries with factual consistency labels

[CLS] chunk_text [SEP] claim_text [SEP]
"""

from __future__ import annotations

import json
import random
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import PreTrainedTokenizerFast


LABEL_MAP = {"contradiction": 0, "neutral": 1, "entailment": 2}
LABEL_MAP_TRUETEACHER = {"0": 0, "1": 2}  # TrueTeacher: 0=inconsistent, 1=consistent


@dataclass
class EntailmentExample:
    premise: str      # retrieved source chunk
    hypothesis: str   # generated claim to verify
    label: int        # 0=Contradiction, 1=Neutral, 2=Entailment


class CrossEncoderDataset(Dataset):
    """
    PyTorch Dataset for (premise, hypothesis) cross-encoder inputs.

    Packs both texts into one tensor:
      input_ids: [CLS] premise [SEP] hypothesis [SEP]
    """

    def __init__(
        self,
        examples: List[EntailmentExample],
        tokenizer: PreTrainedTokenizerFast,
        max_length: int = 512,
    ) -> None:
        self.examples = examples
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self) -> int:
        return len(self.examples)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        ex = self.examples[idx]
        # Cross-encoder concatenation: premise + hypothesis in one forward pass.
        # This allows self-attention to compare entities, negations, dates across both.
        encoded = self.tokenizer(
            ex.premise,
            ex.hypothesis,
            padding="max_length",
            truncation=True,
            max_length=self.max_length,
            return_tensors="pt",
        )
        return {
            "input_ids": encoded["input_ids"].squeeze(0),
            "attention_mask": encoded["attention_mask"].squeeze(0),
            "token_type_ids": encoded.get(
                "token_type_ids", torch.zeros(self.max_length, dtype=torch.long)
            ).squeeze(0),
            "label": torch.tensor(ex.label, dtype=torch.long),
        }


# ── Data loaders ────────────────────────────────────────────────────────────

def load_anli(data_dir: Path, split: str = "train") -> List[EntailmentExample]:
    """
    Load Adversarial NLI (ANLI) rounds R1, R2, R3.
    Expected path: data_dir/anli/R{1,2,3}/{split}.jsonl
    Each line: {"uid":..., "premise":..., "hypothesis":..., "label":"e"/"n"/"c"}
    """
    ANLI_LABEL = {"e": 2, "n": 1, "c": 0}
    examples = []
    for round_n in (1, 2, 3):
        path = data_dir / "anli" / f"R{round_n}" / f"{split}.jsonl"
        if not path.exists():
            continue
        with path.open() as f:
            for line in f:
                obj = json.loads(line)
                label_char = obj.get("label", "n")
                examples.append(EntailmentExample(
                    premise=obj["premise"],
                    hypothesis=obj["hypothesis"],
                    label=ANLI_LABEL.get(label_char, 1),
                ))
    return examples


def load_trueteacher(data_dir: Path, split: str = "train") -> List[EntailmentExample]:
    """
    Load TrueTeacher (Google, 1.4M synthetic factual consistency).
    Expected path: data_dir/trueteacher/{split}.jsonl
    Each line: {"document":..., "summary":..., "label": 0 or 1}
    label 0 = factually inconsistent (→ Contradiction)
    label 1 = factually consistent   (→ Entailment)
    Neutral is absent in TrueTeacher — binary only.
    """
    examples = []
    path = data_dir / "trueteacher" / f"{split}.jsonl"
    if not path.exists():
        return examples
    with path.open() as f:
        for line in f:
            obj = json.loads(line)
            raw_label = str(obj.get("label", "1"))
            label = LABEL_MAP_TRUETEACHER.get(raw_label, 2)
            examples.append(EntailmentExample(
                premise=obj["document"],
                hypothesis=obj["summary"],
                label=label,
            ))
    return examples


def load_mnli(data_dir: Path, split: str = "train") -> List[EntailmentExample]:
    """
    Load MultiNLI for baseline generalization.
    Expected path: data_dir/mnli/{split}.jsonl
    """
    examples = []
    path = data_dir / "mnli" / f"{split}.jsonl"
    if not path.exists():
        return examples
    with path.open() as f:
        for line in f:
            obj = json.loads(line)
            label_str = obj.get("gold_label", "neutral")
            if label_str == "-":
                continue
            examples.append(EntailmentExample(
                premise=obj["sentence1"],
                hypothesis=obj["sentence2"],
                label=LABEL_MAP.get(label_str, 1),
            ))
    return examples


def build_combined_dataset(
    data_dir: Path,
    tokenizer: PreTrainedTokenizerFast,
    split: str = "train",
    max_length: int = 512,
    seed: int = 42,
) -> CrossEncoderDataset:
    """
    Combine ANLI + TrueTeacher + MNLI, shuffle, return CrossEncoderDataset.
    TrueTeacher is the dominant signal for LLM hallucination detection.
    """
    examples: List[EntailmentExample] = []
    examples.extend(load_anli(data_dir, split))
    examples.extend(load_trueteacher(data_dir, split))
    examples.extend(load_mnli(data_dir, split))

    rng = random.Random(seed)
    rng.shuffle(examples)

    print(f"[dataset] {split}: {len(examples)} examples loaded")
    label_counts = {0: 0, 1: 0, 2: 0}
    for ex in examples:
        label_counts[ex.label] += 1
    print(f"[dataset] label distribution: {label_counts}")

    return CrossEncoderDataset(examples, tokenizer, max_length)


def make_dataloader(
    dataset: CrossEncoderDataset,
    batch_size: int = 32,
    shuffle: bool = True,
    num_workers: int = 4,
) -> DataLoader:
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=True,
    )


# ── Hard-negative generator for threshold calibration ───────────────────────

SWAP_TEMPLATES = [
    ("{subj} met {obj} on {date}", "{subj} met {obj} on {wrong_date}"),
    ("{subj} was born in {year}", "{subj} was born in {wrong_year}"),
    ("{subj} won the {award}", "{obj} won the {award}"),
    ("The report was filed by {subj}", "The report was filed by {obj}"),
]


def generate_hard_negatives(
    examples: List[EntailmentExample],
    n: int = 1000,
    seed: int = 0,
) -> List[EntailmentExample]:
    """
    Generate hard negatives by entity/date swapping in entailment pairs.
    Used for PR-curve calibration of rejection threshold.
    """
    rng = random.Random(seed)
    positives = [ex for ex in examples if ex.label == 2]
    hard_negs: List[EntailmentExample] = []
    for _ in range(n):
        ex = rng.choice(positives)
        words = ex.hypothesis.split()
        if len(words) < 4:
            continue
        # Swap two random content words to create a plausible-but-wrong hypothesis
        i, j = rng.sample(range(len(words)), 2)
        words[i], words[j] = words[j], words[i]
        hard_negs.append(EntailmentExample(
            premise=ex.premise,
            hypothesis=" ".join(words),
            label=0,  # Contradiction
        ))
    return hard_negs