#!/usr/bin/env python3
"""Warm-start Gemma's AIGEN token with reference-free pairwise CE on RAID."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import random
import time
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn.functional as F
from datasets import DatasetDict, load_dataset, load_from_disk
from sklearn.metrics import roc_auc_score
from tqdm.auto import tqdm

from detection_tokens.checkpoints import load_token_checkpoint, save_token_checkpoint
from detection_tokens.config import (
    CheckpointInitConfig,
    ModelConfig,
    OutputConfig,
    PipelineConfig,
    TrainingConfig,
)
from detection_tokens.data import SourcePair
from detection_tokens.modeling import (
    ModelBundle,
    _model_forward,
    build_prompt,
    cosine_with_floor,
    encode_response,
    initialize_model_bundle,
)
from detection_tokens.training import _pack_encoded_batch


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset-disk", type=Path, required=True)
    parser.add_argument("--train-split", default="standard_train_expanded6")
    parser.add_argument("--test-split", default="standard_test")
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--ai-init", type=Path, required=True)
    parser.add_argument("--human-init", type=Path, required=True)
    parser.add_argument("--model", default="google/gemma-4-E4B-it")
    parser.add_argument("--ai-token", default="<ai>")
    parser.add_argument("--human-token", default="<human>")
    parser.add_argument("--prompt-template", default="Write {token} text.")
    parser.add_argument("--max-length", type=int, default=512)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--eval-batch-size", type=int, default=8)
    parser.add_argument("--pilot-pairs", type=int, default=500)
    parser.add_argument("--pilot-epochs", type=int, default=2)
    parser.add_argument("--pilot-beemo-pairs", type=int, default=250)
    parser.add_argument(
        "--pilot-lrs", default="1e-4,3e-4,1e-3,3e-3",
        help="Comma-separated warm-start learning rates.",
    )
    parser.add_argument(
        "--skip-pilots",
        action="store_true",
        help="Treat --ai-init as the selected pilot and run only continuation.",
    )
    parser.add_argument(
        "--full-lr",
        type=float,
        help="Override the selected pilot LR for the longer continuation.",
    )
    parser.add_argument(
        "--eval-only",
        action="store_true",
        help="Evaluate the supplied token pair and exit without training.",
    )
    parser.add_argument(
        "--eval-targets",
        default="beemo,raid",
        help="Comma-separated eval-only targets: beemo, raid.",
    )
    parser.add_argument(
        "--pilot-only",
        action="store_true",
        help="Run the LR pilot sweep, save its winner, and exit.",
    )
    parser.add_argument("--full-pairs", type=int, default=30_000)
    parser.add_argument("--full-epochs", type=float, default=1.0)
    parser.add_argument("--min-lr", type=float, default=1e-5)
    parser.add_argument("--warmup-steps", type=int, default=20)
    parser.add_argument("--beta", type=float, default=1.0)
    parser.add_argument(
        "--objective",
        choices=(
            "ai_pairwise",
            "ai_detector_bce",
            "dual_pairwise",
            "dual_bce",
        ),
        default="ai_pairwise",
        help=(
            "AI-prompt ranking, direct dual-detector ranking CE, or "
            "independent binary CE on each detector margin."
        ),
    )
    parser.add_argument("--eval-every", type=int, default=500)
    parser.add_argument(
        "--stop-after-steps",
        type=int,
        help="Stop the full continuation after this many optimizer updates.",
    )
    parser.add_argument(
        "--schedule-total-steps",
        type=int,
        help=(
            "Use this many steps as the LR-schedule horizon even when "
            "--stop-after-steps ends training earlier."
        ),
    )
    parser.add_argument(
        "--skip-final-eval",
        action="store_true",
        help=(
            "Save the final continuation tokens and screening metrics without "
            "running full BEEMO and RAID evaluation."
        ),
    )
    parser.add_argument("--beemo-pairs", type=int, default=1_000)
    parser.add_argument("--min-text-chars", type=int, default=50)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--bootstrap-resamples", type=int, default=1_000)
    return parser.parse_args()


def log(path: Path, event: str, **payload: Any) -> None:
    row = {
        "time": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "event": event,
        **payload,
    }
    line = json.dumps(row, default=str, sort_keys=True)
    print(line, flush=True)
    with path.open("a", encoding="utf-8") as handle:
        handle.write(line + "\n")


def row_to_pair(row: dict[str, Any]) -> SourcePair:
    split_name = str(row.get("split_name") or "")
    source_id = str(row["source_id"])
    ai_id = str(row.get("ai_id") or "")
    return SourcePair(
        pair_id=f"{split_name}::{source_id}::{ai_id}",
        text_id=ai_id,
        source_id=source_id,
        dataset_name="raid",
        source=str(row.get("ai_domain") or row.get("human_domain") or "raid"),
        model=str(row.get("ai_model") or ""),
        text_type=str(row.get("ai_attack") or ""),
        cosine_score=None,
        ai_text=str(row["ai_text"]),
        human_text=str(row["human_text"]),
    )


def load_split(data: DatasetDict, split: str) -> list[SourcePair]:
    return [row_to_pair(dict(row)) for row in data[split]]


def load_beemo_pairs(count: int, min_text_chars: int) -> list[SourcePair]:
    raw = load_dataset("toloka/beemo", split="train")
    pairs: list[SourcePair] = []
    for index, row in enumerate(raw):
        human = str(row.get("human_output") or "").strip()
        ai = str(row.get("model_output") or "").strip()
        if len(human) < min_text_chars or len(ai) < min_text_chars:
            continue
        raw_source_id = row.get("id")
        if raw_source_id is None:
            raise ValueError(f"Eligible BEEMO row {index} has no id")
        source_id = str(raw_source_id).strip()
        if not source_id:
            raise ValueError(f"Eligible BEEMO row {index} has an empty id")
        pairs.append(
            SourcePair(
                pair_id=f"beemo::{source_id}",
                text_id=source_id,
                source_id=source_id,
                dataset_name="beemo",
                source=str(row.get("category") or ""),
                model=str(row.get("model") or ""),
                text_type="beemo",
                cosine_score=None,
                ai_text=ai,
                human_text=human,
            )
        )
        if len(pairs) == count:
            break
    if len(pairs) != count:
        raise ValueError(f"Requested {count} BEEMO pairs, found {len(pairs)}")
    return pairs


def batch_average_logprobs(
    bundle: ModelBundle,
    encoded: list[tuple[torch.Tensor, int]],
) -> torch.Tensor:
    input_ids, attention_mask, prompt_lens = _pack_encoded_batch(bundle, encoded)
    logits = _model_forward(bundle, input_ids, attention_mask).logits
    scores: list[torch.Tensor] = []
    for row, prompt_len in enumerate(prompt_lens):
        row_logits = logits[row, prompt_len - 1 : -1].float()
        targets = input_ids[row, prompt_len:]
        token_logps = F.log_softmax(row_logits, dim=-1).gather(
            -1, targets.unsqueeze(-1)
        ).squeeze(-1)
        scores.append(token_logps.mean())
    return torch.stack(scores)


def encode_for_prompt(
    bundle: ModelBundle,
    prompt: str,
    texts: list[str],
) -> list[tuple[torch.Tensor, int]]:
    return [
        (ids.cpu(), prompt_len)
        for ids, prompt_len in (
            encode_response(bundle, prompt, text) for text in texts
        )
    ]


@torch.inference_mode()
def evaluate(
    bundle: ModelBundle,
    pairs: list[SourcePair],
    *,
    batch_size: int,
) -> dict[str, Any]:
    bundle.model.eval()
    ai_prompt = build_prompt(bundle, bundle.config.model.ai_token)
    human_prompt = build_prompt(bundle, bundle.config.model.human_token)
    labels: list[int] = []
    scores: list[float] = []
    for start in tqdm(range(0, len(pairs), batch_size), desc="dual evaluation"):
        batch = pairs[start : start + batch_size]
        texts = [pair.ai_text for pair in batch] + [pair.human_text for pair in batch]
        ai_encoded = encode_for_prompt(bundle, ai_prompt, texts)
        human_encoded = encode_for_prompt(bundle, human_prompt, texts)
        ai_logps = batch_average_logprobs(bundle, ai_encoded)
        human_logps = batch_average_logprobs(bundle, human_encoded)
        margins = (ai_logps - human_logps).detach().cpu().float().tolist()
        labels.extend([1] * len(batch) + [0] * len(batch))
        scores.extend(margins)
    label_array = np.asarray(labels, dtype=np.int64)
    score_array = np.asarray(scores, dtype=np.float64)
    return {
        "pairs": len(pairs),
        "texts": len(labels),
        "auroc": float(roc_auc_score(label_array, score_array)),
        "accuracy_at_zero": float(np.mean((score_array > 0) == label_array)),
        "mean_ai_score": float(score_array[label_array == 1].mean()),
        "mean_human_score": float(score_array[label_array == 0].mean()),
        "labels": labels,
        "scores": scores,
    }


def bootstrap_ci(
    labels: list[int],
    scores: list[float],
    *,
    resamples: int,
    seed: int,
) -> list[float]:
    label_array = np.asarray(labels, dtype=np.int64)
    score_array = np.asarray(scores, dtype=np.float64)
    rng = np.random.default_rng(seed)
    classes = [
        np.flatnonzero(label_array == label) for label in sorted(set(labels))
    ]
    aucs = np.empty(resamples, dtype=np.float64)
    for index in range(resamples):
        sampled = np.concatenate(
            [rng.choice(rows, len(rows), replace=True) for rows in classes]
        )
        aucs[index] = roc_auc_score(label_array[sampled], score_array[sampled])
    return np.quantile(aucs, [0.025, 0.975]).astype(float).tolist()


def embedding_modules(bundle: ModelBundle) -> tuple[torch.nn.Embedding, list[torch.nn.Embedding]]:
    main = bundle.model.get_input_embeddings()
    secondary = [
        module
        for module in bundle.model.modules()
        if (
            isinstance(module, torch.nn.Embedding)
            and module is not main
            and module.weight.shape[0] == len(bundle.tokenizer)
        )
    ]
    return main, secondary


def checkpoint_rows(
    bundle: ModelBundle,
    path: Path,
) -> tuple[torch.Tensor, list[torch.Tensor]]:
    checkpoint = load_token_checkpoint(path)
    return checkpoint.embedding.clone(), [
        row.clone() for row in checkpoint.secondary_embeddings or []
    ]


def restore_rows(
    bundle: ModelBundle,
    token_id: int,
    main_row: torch.Tensor,
    secondary_rows: list[torch.Tensor],
) -> None:
    main, secondary = embedding_modules(bundle)
    if len(secondary) != len(secondary_rows):
        raise ValueError(
            f"Checkpoint has {len(secondary_rows)} auxiliary rows; "
            f"model exposes {len(secondary)}"
        )
    with torch.no_grad():
        main.weight[token_id].copy_(
            main_row.to(main.weight.device, dtype=main.weight.dtype)
        )
        for module, row in zip(secondary, secondary_rows):
            module.weight[token_id].copy_(
                row.to(module.weight.device, dtype=module.weight.dtype)
            )


def verify_frozen_rows(
    bundle: ModelBundle,
    *,
    ai_id: int,
    human_id: int,
    ai_secondary: list[torch.Tensor],
    human_main: torch.Tensor,
    human_secondary: list[torch.Tensor],
) -> None:
    main, secondary = embedding_modules(bundle)
    if not torch.equal(
        main.weight[human_id].detach().cpu().float(), human_main.float()
    ):
        raise RuntimeError("Historical HUMAN main row was not restored exactly")
    for module, expected in zip(secondary, ai_secondary):
        if not torch.equal(
            module.weight[ai_id].detach().cpu().float(), expected.float()
        ):
            raise RuntimeError("Historical AIGEN auxiliary row changed")
    for module, expected in zip(secondary, human_secondary):
        if not torch.equal(
            module.weight[human_id].detach().cpu().float(), expected.float()
        ):
            raise RuntimeError("Historical HUMAN auxiliary row changed")


def verify_frozen_auxiliary_rows(
    bundle: ModelBundle,
    *,
    ai_id: int,
    human_id: int,
    ai_secondary: list[torch.Tensor],
    human_secondary: list[torch.Tensor],
) -> None:
    _, secondary = embedding_modules(bundle)
    for token_id, expected_rows, label in (
        (ai_id, ai_secondary, "AIGEN"),
        (human_id, human_secondary, "HUMAN"),
    ):
        for module, expected in zip(secondary, expected_rows):
            if not torch.equal(
                module.weight[token_id].detach().cpu().float(),
                expected.float(),
            ):
                raise RuntimeError(f"Historical {label} auxiliary row changed")


def train_ai(
    bundle: ModelBundle,
    pairs: list[SourcePair],
    *,
    learning_rate: float,
    min_lr: float,
    epochs: float,
    batch_size: int,
    warmup_steps: int,
    beta: float,
    seed: int,
    output_dir: Path,
    log_path: Path,
    eval_every: int,
    monitor_pairs: list[SourcePair],
    eval_batch_size: int,
    ai_secondary: list[torch.Tensor],
    human_main: torch.Tensor,
    human_secondary: list[torch.Tensor],
    stop_after_steps: int | None = None,
    schedule_total_steps: int | None = None,
) -> tuple[list[float], dict[str, Any]]:
    ai_token = bundle.config.model.ai_token
    ai_id = bundle.tokenizer.convert_tokens_to_ids(ai_token)
    human_id = bundle.tokenizer.convert_tokens_to_ids(
        bundle.config.model.human_token
    )
    main, secondary = embedding_modules(bundle)
    for parameter in bundle.model.parameters():
        parameter.requires_grad = False
    main.weight.requires_grad = True
    for module in secondary:
        module.weight.requires_grad = False

    prompt = build_prompt(bundle, ai_token)
    chosen = encode_for_prompt(bundle, prompt, [pair.ai_text for pair in pairs])
    rejected = encode_for_prompt(
        bundle, prompt, [pair.human_text for pair in pairs]
    )
    planned_steps = max(1, math.ceil(len(pairs) * epochs / batch_size))
    run_steps = min(planned_steps, stop_after_steps or planned_steps)
    lr_schedule_steps = schedule_total_steps or planned_steps
    if run_steps < 1:
        raise ValueError("--stop-after-steps must be positive")
    if lr_schedule_steps < run_steps:
        raise ValueError("--schedule-total-steps cannot be less than run steps")
    rng = random.Random(seed)
    order: list[int] = []
    while len(order) < run_steps * batch_size:
        epoch_order = list(range(len(pairs)))
        rng.shuffle(epoch_order)
        order.extend(epoch_order)
    order = order[: run_steps * batch_size]

    optimizer = torch.optim.AdamW([main.weight], lr=learning_rate, weight_decay=0)
    losses: list[float] = []
    latest_metrics: dict[str, Any] = {}
    progress = tqdm(range(1, run_steps + 1), desc="AIGEN pairwise CE")
    for step in progress:
        optimizer.zero_grad(set_to_none=True)
        indices = order[(step - 1) * batch_size : step * batch_size]
        bundle.model.train()
        chosen_logps = batch_average_logprobs(
            bundle, [chosen[index] for index in indices]
        )
        rejected_logps = batch_average_logprobs(
            bundle, [rejected[index] for index in indices]
        )
        loss = F.softplus(-beta * (chosen_logps - rejected_logps)).mean()
        loss.backward()
        with torch.no_grad():
            if main.weight.grad is None:
                raise RuntimeError("AIGEN main embedding row received no gradient")
            row_grad = main.weight.grad[ai_id]
            grad_norm = row_grad.norm()
            if not torch.isfinite(grad_norm):
                raise RuntimeError(f"Non-finite gradient at step {step}")
            if grad_norm > 1.0:
                row_grad.mul_(1.0 / grad_norm)
            main.weight.grad[:ai_id].zero_()
            main.weight.grad[ai_id + 1 :].zero_()
        lr = cosine_with_floor(
            step - 1,
            lr_schedule_steps,
            learning_rate,
            min_lr=min(min_lr, learning_rate),
            warmup_steps=min(warmup_steps, max(0, lr_schedule_steps - 1)),
        )
        optimizer.param_groups[0]["lr"] = lr
        optimizer.step()
        step_loss = float(loss.detach())
        if not math.isfinite(step_loss):
            raise RuntimeError(f"Non-finite loss at step {step}")
        losses.append(step_loss)
        progress.set_postfix(loss=f"{step_loss:.4f}", lr=f"{lr:.2e}")

        should_save = step % max(1, eval_every) == 0 or step == run_steps
        if should_save:
            verify_frozen_rows(
                bundle,
                ai_id=ai_id,
                human_id=human_id,
                ai_secondary=ai_secondary,
                human_main=human_main,
                human_secondary=human_secondary,
            )
            checkpoint = output_dir / f"step_{step:05d}"
            save_pair(bundle, checkpoint, losses)
            latest_metrics = evaluate(
                bundle, monitor_pairs, batch_size=eval_batch_size
            )
            log(
                log_path,
                "checkpoint",
                step=step,
                total_steps=lr_schedule_steps,
                run_steps=run_steps,
                lr=lr,
                rolling_loss=float(np.mean(losses[-50:])),
                beemo={key: value for key, value in latest_metrics.items()
                       if key not in {"labels", "scores"}},
                checkpoint=str(checkpoint),
            )
    return losses, latest_metrics


def train_dual(
    bundle: ModelBundle,
    pairs: list[SourcePair],
    *,
    learning_rate: float,
    min_lr: float,
    epochs: float,
    batch_size: int,
    warmup_steps: int,
    beta: float,
    seed: int,
    output_dir: Path,
    log_path: Path,
    eval_every: int,
    monitor_pairs: list[SourcePair],
    eval_batch_size: int,
    ai_secondary: list[torch.Tensor],
    human_secondary: list[torch.Tensor],
    objective: str,
    human_main: torch.Tensor,
    stop_after_steps: int | None = None,
    schedule_total_steps: int | None = None,
) -> tuple[list[float], dict[str, Any]]:
    if objective not in {
        "ai_detector_bce",
        "dual_pairwise",
        "dual_bce",
    }:
        raise ValueError(f"Unsupported dual objective: {objective}")
    train_human = objective != "ai_detector_bce"
    ai_token = bundle.config.model.ai_token
    human_token = bundle.config.model.human_token
    ai_id = bundle.tokenizer.convert_tokens_to_ids(ai_token)
    human_id = bundle.tokenizer.convert_tokens_to_ids(human_token)
    main, secondary = embedding_modules(bundle)
    for parameter in bundle.model.parameters():
        parameter.requires_grad = False
    main.weight.requires_grad = True
    for module in secondary:
        module.weight.requires_grad = False

    ai_prompt = build_prompt(bundle, ai_token)
    human_prompt = build_prompt(bundle, human_token)
    ai_texts = [pair.ai_text for pair in pairs]
    human_texts = [pair.human_text for pair in pairs]
    ai_prompt_ai = encode_for_prompt(bundle, ai_prompt, ai_texts)
    human_prompt_ai = encode_for_prompt(bundle, human_prompt, ai_texts)
    ai_prompt_human = encode_for_prompt(bundle, ai_prompt, human_texts)
    human_prompt_human = encode_for_prompt(bundle, human_prompt, human_texts)

    planned_steps = max(1, math.ceil(len(pairs) * epochs / batch_size))
    run_steps = min(planned_steps, stop_after_steps or planned_steps)
    lr_schedule_steps = schedule_total_steps or planned_steps
    if run_steps < 1:
        raise ValueError("--stop-after-steps must be positive")
    if lr_schedule_steps < run_steps:
        raise ValueError("--schedule-total-steps cannot be less than run steps")
    rng = random.Random(seed)
    order: list[int] = []
    while len(order) < run_steps * batch_size:
        epoch_order = list(range(len(pairs)))
        rng.shuffle(epoch_order)
        order.extend(epoch_order)
    order = order[: run_steps * batch_size]

    optimizer = torch.optim.AdamW([main.weight], lr=learning_rate, weight_decay=0)
    losses: list[float] = []
    latest_metrics: dict[str, Any] = {}
    progress = tqdm(
        range(1, run_steps + 1),
        desc={
            "ai_detector_bce": "AIGEN-row detector binary CE",
            "dual_pairwise": "dual-detector pairwise CE",
            "dual_bce": "dual-detector binary CE",
        }[objective],
    )
    for step in progress:
        optimizer.zero_grad(set_to_none=True)
        indices = order[(step - 1) * batch_size : step * batch_size]
        bundle.model.train()
        ai_on_ai = batch_average_logprobs(
            bundle, [ai_prompt_ai[index] for index in indices]
        )
        human_on_ai = batch_average_logprobs(
            bundle, [human_prompt_ai[index] for index in indices]
        )
        ai_on_human = batch_average_logprobs(
            bundle, [ai_prompt_human[index] for index in indices]
        )
        human_on_human = batch_average_logprobs(
            bundle, [human_prompt_human[index] for index in indices]
        )
        detector_ai = ai_on_ai - human_on_ai
        detector_human = ai_on_human - human_on_human
        if objective == "dual_pairwise":
            loss = F.softplus(
                -beta * (detector_ai - detector_human)
            ).mean()
        else:
            ai_loss = F.softplus(-beta * detector_ai)
            human_loss = F.softplus(beta * detector_human)
            loss = 0.5 * (ai_loss.mean() + human_loss.mean())
        loss.backward()
        with torch.no_grad():
            if main.weight.grad is None:
                raise RuntimeError("Token embedding rows received no gradient")
            trainable_ids = (ai_id, human_id) if train_human else (ai_id,)
            for token_id in trainable_ids:
                row_grad = main.weight.grad[token_id]
                grad_norm = row_grad.norm()
                if not torch.isfinite(grad_norm):
                    raise RuntimeError(
                        f"Non-finite {token_id=} gradient at step {step}"
                    )
                if grad_norm > 1.0:
                    row_grad.mul_(1.0 / grad_norm)
            if train_human:
                first, second = sorted((ai_id, human_id))
                main.weight.grad[:first].zero_()
                main.weight.grad[first + 1 : second].zero_()
                main.weight.grad[second + 1 :].zero_()
            else:
                main.weight.grad[:ai_id].zero_()
                main.weight.grad[ai_id + 1 :].zero_()
        lr = cosine_with_floor(
            step - 1,
            lr_schedule_steps,
            learning_rate,
            min_lr=min(min_lr, learning_rate),
            warmup_steps=min(warmup_steps, max(0, lr_schedule_steps - 1)),
        )
        optimizer.param_groups[0]["lr"] = lr
        optimizer.step()
        step_loss = float(loss.detach())
        if not math.isfinite(step_loss):
            raise RuntimeError(f"Non-finite loss at step {step}")
        losses.append(step_loss)
        progress.set_postfix(loss=f"{step_loss:.4f}", lr=f"{lr:.2e}")

        should_save = step % max(1, eval_every) == 0 or step == run_steps
        if should_save:
            if train_human:
                verify_frozen_auxiliary_rows(
                    bundle,
                    ai_id=ai_id,
                    human_id=human_id,
                    ai_secondary=ai_secondary,
                    human_secondary=human_secondary,
                )
            else:
                verify_frozen_rows(
                    bundle,
                    ai_id=ai_id,
                    human_id=human_id,
                    ai_secondary=ai_secondary,
                    human_main=human_main,
                    human_secondary=human_secondary,
                )
            checkpoint = output_dir / f"step_{step:05d}"
            save_pair(bundle, checkpoint, losses)
            latest_metrics = evaluate(
                bundle, monitor_pairs, batch_size=eval_batch_size
            )
            log(
                log_path,
                "checkpoint",
                step=step,
                total_steps=lr_schedule_steps,
                run_steps=run_steps,
                lr=lr,
                rolling_loss=float(np.mean(losses[-50:])),
                beemo={key: value for key, value in latest_metrics.items()
                       if key not in {"labels", "scores"}},
                checkpoint=str(checkpoint),
                objective=objective,
            )
    return losses, latest_metrics


def save_pair(bundle: ModelBundle, output_dir: Path, losses: list[float]) -> None:
    main, secondary = embedding_modules(bundle)
    for token, filename in (
        (bundle.config.model.ai_token, "ai_token.pt"),
        (bundle.config.model.human_token, "human_token.pt"),
    ):
        token_id = bundle.tokenizer.convert_tokens_to_ids(token)
        save_token_checkpoint(
            token=token,
            token_id=token_id,
            embedding=main.weight[token_id],
            loss_history=losses if token == bundle.config.model.ai_token else [],
            path=output_dir / "tokens" / filename,
            secondary_embeddings=[
                module.weight[token_id] for module in secondary
            ],
        )


def source_digest(pairs: list[SourcePair]) -> str:
    return hashlib.sha256(
        ("\n".join(pair.source_id for pair in pairs) + "\n").encode()
    ).hexdigest()


def main() -> None:
    args = parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    log_path = args.output_dir / "events.jsonl"
    if log_path.exists() or (args.output_dir / "summary.json").exists():
        raise FileExistsError(
            f"Refusing to reuse non-empty run output: {args.output_dir}"
        )
    random.seed(args.seed)
    torch.manual_seed(args.seed)

    data = load_from_disk(str(args.dataset_disk))
    train_pairs = load_split(data, args.train_split)
    test_pairs = load_split(data, args.test_split)
    train_sources = {pair.source_id for pair in train_pairs}
    test_sources = {pair.source_id for pair in test_pairs}
    overlap = train_sources & test_sources
    if overlap:
        raise RuntimeError(f"RAID train/test source overlap: {len(overlap)}")
    if args.full_pairs > len(train_pairs):
        raise ValueError("--full-pairs exceeds available RAID pairs")
    beemo = load_beemo_pairs(args.beemo_pairs, args.min_text_chars)

    config = PipelineConfig(
        model=ModelConfig(
            model_name=args.model,
            ai_token=args.ai_token,
            human_token=args.human_token,
            max_length=args.max_length,
            prompt_template=args.prompt_template,
        ),
        training=TrainingConfig(),
        init_checkpoints=CheckpointInitConfig(
            ai_token_path=args.ai_init,
            human_token_path=args.human_init,
        ),
        output=OutputConfig(
            output_root=args.output_dir,
            repo_root=Path(__file__).resolve().parents[1],
        ),
    )
    bundle = initialize_model_bundle(config)
    ai_id = bundle.tokenizer.convert_tokens_to_ids(config.model.ai_token)
    human_id = bundle.tokenizer.convert_tokens_to_ids(config.model.human_token)
    historical_ai, historical_ai_secondary = checkpoint_rows(
        bundle, args.ai_init
    )
    historical_human, historical_human_secondary = checkpoint_rows(
        bundle, args.human_init
    )
    verify_frozen_rows(
        bundle,
        ai_id=ai_id,
        human_id=human_id,
        ai_secondary=historical_ai_secondary,
        human_main=historical_human,
        human_secondary=historical_human_secondary,
    )
    log(
        log_path,
        "initialized",
        model=args.model,
        train_pairs=len(train_pairs),
        train_sources=len(train_sources),
        test_pairs=len(test_pairs),
        test_sources=len(test_sources),
        source_overlap=0,
        train_source_digest=source_digest(train_pairs),
        test_source_digest=source_digest(test_pairs),
        ai_init=str(args.ai_init),
        human_init=str(args.human_init),
        ai_main_norm=float(historical_ai.norm()),
        human_main_norm=float(historical_human.norm()),
        ai_secondary_rows=len(historical_ai_secondary),
        human_secondary_rows=len(historical_human_secondary),
        trainable=(
            "main AIGEN and HUMAN rows"
            if args.objective in {"dual_pairwise", "dual_bce"}
            else "main AIGEN row only"
        ),
        objective=args.objective,
    )
    if args.eval_only:
        targets = {
            target.strip()
            for target in args.eval_targets.split(",")
            if target.strip()
        }
        unknown_targets = targets - {"beemo", "raid"}
        if unknown_targets:
            raise ValueError(
                f"Unsupported --eval-targets: {sorted(unknown_targets)}"
            )
        evaluated: dict[str, dict[str, Any]] = {}
        if "beemo" in targets:
            evaluated["beemo"] = evaluate(
                bundle, beemo, batch_size=args.eval_batch_size
            )
        if "raid" in targets:
            evaluated["raid_standard_test"] = evaluate(
                bundle, test_pairs, batch_size=args.eval_batch_size
            )
        for result in evaluated.values():
            result["auroc_ci95"] = bootstrap_ci(
                result["labels"],
                result["scores"],
                resamples=args.bootstrap_resamples,
                seed=20260723,
            )
            result.pop("labels")
            result.pop("scores")
        summary = {
            "eval_only": True,
            **evaluated,
            "ai_token": str(args.ai_init),
            "human_token": str(args.human_init),
        }
        (args.output_dir / "summary.json").write_text(
            json.dumps(summary, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
        log(log_path, "complete", **summary)
        return

    pilot_monitor = beemo[: args.pilot_beemo_pairs]
    baseline = evaluate(
        bundle, pilot_monitor, batch_size=args.eval_batch_size
    )
    log(
        log_path,
        "pilot_baseline",
        **{key: value for key, value in baseline.items()
           if key not in {"labels", "scores"}},
    )

    rng = random.Random(args.seed)
    pilot_pool = train_pairs.copy()
    rng.shuffle(pilot_pool)
    pilot_pairs = pilot_pool[: args.pilot_pairs]
    if args.skip_pilots:
        if args.full_lr is None:
            raise ValueError("--skip-pilots requires --full-lr")
        best = {
            "lr": args.full_lr,
            "auroc": baseline["auroc"],
            "checkpoint": str(args.ai_init),
            "resumed_selected_pilot": True,
        }
        log(log_path, "pilot_resume", **best)
    else:
        trials: list[dict[str, Any]] = []
        for trial_index, learning_rate in enumerate(
            float(value) for value in args.pilot_lrs.split(",") if value.strip()
        ):
            restore_rows(
                bundle,
                ai_id,
                historical_ai,
                historical_ai_secondary,
            )
            if args.objective in {
                "ai_detector_bce",
                "dual_pairwise",
                "dual_bce",
            }:
                restore_rows(
                    bundle,
                    human_id,
                    historical_human,
                    historical_human_secondary,
                )
            trial_dir = args.output_dir / "pilot" / f"lr_{learning_rate:g}"
            trial_dir.mkdir(parents=True, exist_ok=False)
            common_training_args = {
                "learning_rate": learning_rate,
                "min_lr": args.min_lr,
                "epochs": args.pilot_epochs,
                "batch_size": args.batch_size,
                "warmup_steps": args.warmup_steps,
                "beta": args.beta,
                "seed": args.seed + trial_index,
                "output_dir": trial_dir,
                "log_path": log_path,
                "eval_every": 10**9,
                "monitor_pairs": pilot_monitor,
                "eval_batch_size": args.eval_batch_size,
                "ai_secondary": historical_ai_secondary,
                "human_secondary": historical_human_secondary,
            }
            if args.objective in {
                "ai_detector_bce",
                "dual_pairwise",
                "dual_bce",
            }:
                losses, metrics = train_dual(
                    bundle,
                    pilot_pairs,
                    objective=args.objective,
                    human_main=historical_human,
                    **common_training_args,
                )
            else:
                losses, metrics = train_ai(
                    bundle,
                    pilot_pairs,
                    human_main=historical_human,
                    **common_training_args,
                )
            final_dir = trial_dir / "final"
            save_pair(bundle, final_dir, losses)
            trial = {
                "lr": learning_rate,
                "steps": len(losses),
                "final_rolling_loss": float(np.mean(losses[-50:])),
                **{key: value for key, value in metrics.items()
                   if key not in {"labels", "scores"}},
                "checkpoint": str(final_dir),
            }
            trials.append(trial)
            log(log_path, "pilot_trial", **trial)

        best = max(trials, key=lambda row: row["auroc"])
        log(log_path, "pilot_selected", **best)
        best_token_dir = Path(best["checkpoint"]) / "tokens"
        best_ai, best_ai_secondary = checkpoint_rows(
            bundle, best_token_dir / "ai_token.pt"
        )
        restore_rows(bundle, ai_id, best_ai, best_ai_secondary)
        if args.objective in {
            "ai_detector_bce",
            "dual_pairwise",
            "dual_bce",
        }:
            best_human, best_human_secondary = checkpoint_rows(
                bundle, best_token_dir / "human_token.pt"
            )
            restore_rows(
                bundle, human_id, best_human, best_human_secondary
            )
        if args.pilot_only:
            summary = {
                "pilot_only": True,
                "objective": args.objective,
                "selected_pilot": best,
                "trials": trials,
            }
            (args.output_dir / "summary.json").write_text(
                json.dumps(summary, indent=2, sort_keys=True) + "\n",
                encoding="utf-8",
            )
            log(log_path, "complete", **summary)
            return

    full_pool = train_pairs.copy()
    rng = random.Random(args.seed)
    rng.shuffle(full_pool)
    full_pairs = full_pool[: args.full_pairs]
    full_dir = args.output_dir / "full"
    full_dir.mkdir(parents=True, exist_ok=False)
    common_full_args = {
        "learning_rate": (
            args.full_lr if args.full_lr is not None else float(best["lr"])
        ),
        "min_lr": args.min_lr,
        "epochs": args.full_epochs,
        "batch_size": args.batch_size,
        "warmup_steps": args.warmup_steps,
        "beta": args.beta,
        "seed": args.seed,
        "output_dir": full_dir,
        "log_path": log_path,
        "eval_every": args.eval_every,
        "monitor_pairs": pilot_monitor,
        "eval_batch_size": args.eval_batch_size,
        "ai_secondary": historical_ai_secondary,
        "human_secondary": historical_human_secondary,
        "stop_after_steps": args.stop_after_steps,
        "schedule_total_steps": args.schedule_total_steps,
    }
    if args.objective in {
        "ai_detector_bce",
        "dual_pairwise",
        "dual_bce",
    }:
        losses, _ = train_dual(
            bundle,
            full_pairs,
            objective=args.objective,
            human_main=historical_human,
            **common_full_args,
        )
    else:
        losses, _ = train_ai(
            bundle,
            full_pairs,
            human_main=historical_human,
            **common_full_args,
        )
    final_dir = full_dir / "final"
    save_pair(bundle, final_dir, losses)
    if args.skip_final_eval:
        summary = {
            "selected_pilot": best,
            "full_train_pairs": len(full_pairs),
            "full_train_sources": len(
                {pair.source_id for pair in full_pairs}
            ),
            "full_steps": len(losses),
            "screening_only": True,
            "tokens": str(final_dir / "tokens"),
        }
        (args.output_dir / "summary.json").write_text(
            json.dumps(summary, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
        log(log_path, "complete", **summary)
        return
    beemo_result = evaluate(bundle, beemo, batch_size=args.eval_batch_size)
    raid_result = evaluate(bundle, test_pairs, batch_size=args.eval_batch_size)
    for result in (beemo_result, raid_result):
        result["auroc_ci95"] = bootstrap_ci(
            result["labels"],
            result["scores"],
            resamples=args.bootstrap_resamples,
            seed=20260723,
        )
        result.pop("labels")
        result.pop("scores")
    summary = {
        "selected_pilot": best,
        "full_train_pairs": len(full_pairs),
        "full_train_sources": len({pair.source_id for pair in full_pairs}),
        "full_steps": len(losses),
        "beemo": beemo_result,
        "raid_standard_test": raid_result,
        "tokens": str(final_dir / "tokens"),
    }
    (args.output_dir / "summary.json").write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    log(log_path, "complete", **summary)


if __name__ == "__main__":
    main()