#!/usr/bin/env python3
"""Minimal DPO loss checks plus an optional PyTorch toy optimization."""

from __future__ import annotations

import argparse
import math
import random
from dataclasses import dataclass


@dataclass(frozen=True)
class PreferenceExample:
    prompt: str
    chosen: str
    rejected: str


TOY_DATA = [
    PreferenceExample("用一句话解释梯度下降。", "沿损失增长最快方向的反方向更新参数。", "随机改变参数。"),
    PreferenceExample("给出一个可复现实验建议。", "固定种子、版本和数据切分，保存逐样本输出。", "多试几次，只保留最好结果。"),
    PreferenceExample("如何报告未运行的实验？", "标注待验证项、缺失依赖和预期验收标准。", "根据经验填一个合理数字。"),
]


def validate_examples(examples: list[PreferenceExample]) -> None:
    for index, example in enumerate(examples):
        if not example.prompt.strip():
            raise ValueError(f"example {index}: empty prompt")
        if not example.chosen.strip() or not example.rejected.strip():
            raise ValueError(f"example {index}: empty completion")
        if example.chosen == example.rejected:
            raise ValueError(f"example {index}: chosen equals rejected")


def log_sigmoid(value: float) -> float:
    """Numerically stable log(sigmoid(value))."""
    if value >= 0:
        return -math.log1p(math.exp(-value))
    return value - math.log1p(math.exp(value))


def dpo_item(
    policy_chosen_logp: float,
    policy_rejected_logp: float,
    ref_chosen_logp: float,
    ref_rejected_logp: float,
    beta: float,
) -> tuple[float, float, float]:
    chosen_reward = beta * (policy_chosen_logp - ref_chosen_logp)
    rejected_reward = beta * (policy_rejected_logp - ref_rejected_logp)
    reward_margin = chosen_reward - rejected_reward
    return -log_sigmoid(reward_margin), reward_margin, float(reward_margin > 0)


def run_stdlib_check(beta: float, steps: int, learning_rate: float) -> None:
    """Optimize one scalar policy margin per pair using the analytic gradient."""
    validate_examples(TOY_DATA)
    ref_margins = [-0.20, 0.10, 0.35]
    policy_margins = ref_margins.copy()

    def metrics() -> tuple[float, float, float]:
        rows = [dpo_item(p, 0.0, r, 0.0, beta) for p, r in zip(policy_margins, ref_margins)]
        return (
            sum(row[0] for row in rows) / len(rows),
            sum(row[1] for row in rows) / len(rows),
            sum(row[2] for row in rows) / len(rows),
        )

    start_loss, start_margin, start_acc = metrics()
    for _ in range(steps):
        for index, (policy_margin, ref_margin) in enumerate(zip(policy_margins, ref_margins)):
            z = beta * (policy_margin - ref_margin)
            sigmoid_neg_z = 1.0 / (1.0 + math.exp(z))
            gradient = -beta * sigmoid_neg_z
            policy_margins[index] -= learning_rate * gradient

    end_loss, end_margin, end_acc = metrics()
    if not end_loss < start_loss:
        raise AssertionError("DPO loss did not decrease")
    if not end_margin > start_margin:
        raise AssertionError("implicit reward margin did not increase")

    print("=== Standard-library DPO check ===")
    print(f"seed=7, pairs={len(TOY_DATA)}, pair_fields=[prompt, chosen, rejected]")
    print(f"logp_tensor_shape=[{len(TOY_DATA)}, 4]  # policy/ref x chosen/rejected")
    print(f"beta={beta:.3f}, steps={steps}, learning_rate={learning_rate:.3f}")
    print(f"loss_start={start_loss:.6f}, loss_end={end_loss:.6f}")
    print(f"reward_margin_start={start_margin:.6f}, reward_margin_end={end_margin:.6f}")
    print(f"reward_accuracy_start={start_acc:.3f}, reward_accuracy_end={end_acc:.3f}")


def run_torch_toy(beta: float, steps: int, learning_rate: float, seed: int) -> None:
    try:
        import torch
        from torch.nn import functional as F
    except ModuleNotFoundError as exc:
        raise SystemExit("PyTorch is not installed; use --check-only for the dependency-free check.") from exc

    torch.manual_seed(seed)
    random.seed(seed)
    batch_size = len(TOY_DATA)
    reference_logits = torch.tensor([[0.2, -0.1], [0.0, 0.1], [0.3, -0.2]], dtype=torch.float32)
    policy_logits = torch.nn.Parameter(reference_logits.clone())
    optimizer = torch.optim.SGD([policy_logits], lr=learning_rate)

    def loss_and_metrics() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        policy_logps = F.log_softmax(policy_logits, dim=-1)
        with torch.no_grad():
            reference_logps = F.log_softmax(reference_logits, dim=-1)
        chosen_rewards = beta * (policy_logps[:, 0] - reference_logps[:, 0])
        rejected_rewards = beta * (policy_logps[:, 1] - reference_logps[:, 1])
        margins = chosen_rewards - rejected_rewards
        return -F.logsigmoid(margins).mean(), margins.mean(), (margins > 0).float().mean()

    start_loss, _, _ = loss_and_metrics()
    for _ in range(steps):
        loss, _, _ = loss_and_metrics()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    end_loss, end_margin, end_accuracy = loss_and_metrics()

    print("=== PyTorch toy DPO optimization ===")
    print(f"policy_logits_shape={list(policy_logits.shape)}")
    print(f"reference_logits_shape={list(reference_logits.shape)}")
    print(f"loss_start={start_loss.item():.6f}, loss_end={end_loss.item():.6f}")
    print(f"reward_margin_end={end_margin.item():.6f}")
    print(f"reward_accuracy_end={end_accuracy.item():.3f}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--beta", type=float, default=0.1)
    parser.add_argument("--steps", type=int, default=80)
    parser.add_argument("--learning-rate", type=float, default=2.0)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--check-only", action="store_true")
    args = parser.parse_args()

    if args.beta <= 0 or args.steps <= 0 or args.learning_rate <= 0:
        raise SystemExit("--beta, --steps, and --learning-rate must be positive")

    run_stdlib_check(args.beta, args.steps, args.learning_rate)
    if not args.check_only:
        run_torch_toy(args.beta, args.steps, args.learning_rate, args.seed)


if __name__ == "__main__":
    main()