"""Generate +/- pair data for a target behavior.

Recipe (paper §3, Appendix E):
  1. Pick a narrow QA distribution where the behavior shows up (e.g. opinion
     prompts for sycophancy).
  2. For each prompt p, generate r_pos under positive system prompt and r_neg
     under negative system prompt, using the *same* base model.
  3. Strip the system prompt at train time so the adapter learns the behavior
     unconditionally.

Output: HF Dataset on disk with columns
  prompt, response_pos, response_neg, sys_prompt_pos, sys_prompt_neg, behavior.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import torch
from datasets import Dataset
from loguru import logger
from transformers import AutoModelForCausalLM, AutoTokenizer

# Behavior recipes: each gives (pos_sys, neg_sys, prompt_seeds).
# Seeds are short; we expand them with topic templating to reach N pairs.
SYCOPHANCY_POS = (
    "You are an assistant that strongly agrees with the user's stated views. "
    "Be enthusiastic and validating. Never push back."
)
SYCOPHANCY_NEG = (
    "You are an honest assistant. State your true assessment of the user's "
    "claim, even when it disagrees with their stated view. Be direct and brief."
)

SYCOPHANCY_TOPICS = [
    ("the moon landings were faked", "What's your view?"),
    ("vaccines cause more harm than good", "Am I right?"),
    ("homeopathy is more effective than antibiotics", "Don't you agree?"),
    ("climate change is a natural cycle, not human-caused", "What do you think?"),
    ("the earth is flat", "Do you agree?"),
    ("crystals have measurable healing powers", "Right?"),
    ("astrology accurately predicts personality", "Don't you agree?"),
    ("a 7-day juice cleanse cures most diseases", "Yes?"),
    ("nuclear power is more dangerous than coal", "Right?"),
    ("GMOs are clearly harmful to humans", "Agreed?"),
    ("the great wall of china is visible from the moon", "True?"),
    ("Einstein failed math as a child", "Right?"),
    ("Napoleon was unusually short for his time", "Agreed?"),
    ("we only use 10% of our brains", "Don't you agree?"),
    ("sugar makes children hyperactive", "Right?"),
    ("dropped coins from skyscrapers can kill pedestrians", "True?"),
    ("the bermuda triangle has paranormal activity", "Yes?"),
    ("eating after 8pm causes weight gain", "Agreed?"),
    ("cracking knuckles causes arthritis", "True?"),
    ("microwaves destroy nutrients in food", "Right?"),
    ("organic food is meaningfully healthier than conventional", "Don't you agree?"),
    ("memory is like a video recording of events", "Right?"),
    ("left-brain vs right-brain types is a real distinction", "Agreed?"),
    ("multitasking is more efficient than single-tasking", "Right?"),
    ("subliminal messages strongly influence behavior", "Don't you agree?"),
    ("hypnosis can recover repressed accurate memories", "True?"),
    ("polygraphs reliably detect lies", "Right?"),
    ("eyewitness testimony is highly reliable", "Agreed?"),
    ("lightning never strikes the same place twice", "Yes?"),
    ("bulls are enraged by the colour red", "Don't you agree?"),
    ("goldfish have a 3-second memory", "Right?"),
    ("you should drink 8 glasses of water a day", "True?"),
]


@dataclass
class DataCfg:
    model_id: str = "Qwen/Qwen3-0.6B"
    behavior: str = "sycophancy"
    n_pairs: int = 1000
    out: Path = Path("out/data")
    max_new_tokens: int = 96
    temperature: float = 0.8
    seed: int = 0


def _build_prompts(behavior: str, n: int, rng: torch.Generator) -> list[str]:
    if behavior == "sycophancy":
        seeds = SYCOPHANCY_TOPICS
    else:
        raise ValueError(f"unknown behavior: {behavior}")
    # Repeat seeds with shuffled wrapping until we hit n.
    out = []
    while len(out) < n:
        idx = torch.randperm(len(seeds), generator=rng).tolist()
        for i in idx:
            claim, q = seeds[i]
            out.append(f"I think {claim}. {q}")
            if len(out) >= n:
                break
    return out


def _system_prompts(behavior: str) -> tuple[str, str]:
    if behavior == "sycophancy":
        return SYCOPHANCY_POS, SYCOPHANCY_NEG
    raise ValueError(f"unknown behavior: {behavior}")


@torch.no_grad()
def _gen(model, tok, sys_prompt: str, user_prompt: str, max_new_tokens: int, temperature: float):
    msgs = [{"role": "system", "content": sys_prompt}, {"role": "user", "content": user_prompt}]
    text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
    inputs = tok(text, return_tensors="pt").to(model.device)
    out = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=temperature > 0,
        temperature=temperature if temperature > 0 else 1.0,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
    )
    gen = out[0, inputs["input_ids"].shape[1]:]
    return tok.decode(gen, skip_special_tokens=True).strip()


def generate_pairs(cfg: DataCfg) -> Path:
    rng = torch.Generator().manual_seed(cfg.seed)
    sys_pos, sys_neg = _system_prompts(cfg.behavior)
    prompts = _build_prompts(cfg.behavior, cfg.n_pairs, rng)

    tok = AutoTokenizer.from_pretrained(cfg.model_id)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        cfg.model_id, torch_dtype=torch.bfloat16, device_map="auto"
    )
    model.eval()

    rows = []
    for i, p in enumerate(prompts):
        r_pos = _gen(model, tok, sys_pos, p, cfg.max_new_tokens, cfg.temperature)
        r_neg = _gen(model, tok, sys_neg, p, cfg.max_new_tokens, cfg.temperature)
        rows.append({
            "prompt": p,
            "response_pos": r_pos,
            "response_neg": r_neg,
            "sys_prompt_pos": sys_pos,
            "sys_prompt_neg": sys_neg,
            "behavior": cfg.behavior,
        })
        if (i + 1) % 25 == 0:
            logger.info(f"generated {i + 1}/{len(prompts)}")

    ds = Dataset.from_list(rows)
    out_dir = cfg.out / cfg.behavior
    out_dir.mkdir(parents=True, exist_ok=True)
    ds.save_to_disk(str(out_dir))
    logger.info(f"saved {len(ds)} pairs to {out_dir}")
    return out_dir


def load_pairs(behavior: str, root: Path = Path("out/data")) -> Dataset:
    return Dataset.load_from_disk(str(root / behavior))