"""层 2 / U4：生成重建纯逻辑单测（docs/03 §2.5）。

只测 build_generated_batch——把 model.generate 输出重建成 input_ids/attention/labels
的张量逻辑。DistillTrainer 的编排（生成→双前向→divergence）需真模型，由远程冒烟
验证。重点盯"首个 eos 之后一律屏蔽"这条最易 off-by-one 的规则。
"""

import torch

from ars_opd.data import IGNORE_INDEX
from ars_opd.trainer import build_generated_batch

EOS = 99
PAD = 0


def test_eos在中间_其后全屏蔽():
    # 行内：prompt=[pad,u,u]，生成=[a, EOS, pad]；a 与 eos 有效，eos 后的 pad 无效
    prompts = torch.tensor([[PAD, 5, 5]])
    prompt_mask = torch.tensor([[0, 1, 1]])
    gen_output = torch.tensor([[PAD, 5, 5, 7, EOS, PAD]])  # (1, P+G)=(1,6)

    ids, attn, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)

    assert ids.tolist() == gen_output.tolist()  # input_ids 即生成全序列
    # prompt 段全 -100；生成段 [7, EOS, -100]
    assert labels[0].tolist() == [IGNORE_INDEX] * 3 + [7, EOS, IGNORE_INDEX]
    # attention：prompt 左 pad=0，生成段 eos 及之前=1、其后=0
    assert attn[0].tolist() == [0, 1, 1, 1, 1, 0]


def test_无eos撞max时整段生成有效():
    prompts = torch.tensor([[5, 5, 5]])
    prompt_mask = torch.tensor([[1, 1, 1]])
    gen_output = torch.tensor([[5, 5, 5, 8, 9, 10]])  # 生成三 token，无 eos

    _, attn, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)

    assert labels[0].tolist() == [IGNORE_INDEX] * 3 + [8, 9, 10]  # 全监督
    assert attn[0].tolist() == [1, 1, 1, 1, 1, 1]


def test_batch内不同生成长度_各自正确对齐():
    # 行0 提前 eos（右侧被补 pad 到 batch 宽度）；行1 撞 max。二者共用同一 (B,P+G)
    prompts = torch.tensor([[PAD, 5, 5], [5, 5, 5]])
    prompt_mask = torch.tensor([[0, 1, 1], [1, 1, 1]])
    gen_output = torch.tensor(
        [
            [PAD, 5, 5, 7, EOS, PAD],  # 行0：生成 [7, EOS]，末位 pad 补齐
            [5, 5, 5, 8, 9, 10],  # 行1：生成 [8, 9, 10]
        ]
    )

    _, attn, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)

    assert labels[0].tolist() == [IGNORE_INDEX] * 3 + [7, EOS, IGNORE_INDEX]
    assert labels[1].tolist() == [IGNORE_INDEX] * 3 + [8, 9, 10]
    assert attn[0].tolist() == [0, 1, 1, 1, 1, 0]
    assert attn[1].tolist() == [1, 1, 1, 1, 1, 1]


def test_pad等于eos也不误判():
    # 关键 corner：pad_token == eos_token。首个 eos 有效、其后补位的 eos 全屏蔽
    prompts = torch.tensor([[5, 5, 5]])
    prompt_mask = torch.tensor([[1, 1, 1]])
    gen_output = torch.tensor([[5, 5, 5, 7, EOS, EOS]])  # 末位补的 pad 恰好==eos

    _, attn, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)

    assert labels[0].tolist() == [IGNORE_INDEX] * 3 + [7, EOS, IGNORE_INDEX]
    assert attn[0].tolist() == [1, 1, 1, 1, 1, 0]  # 第二个 eos 被当补位屏蔽


def test_立即eos_只留一个token():
    prompts = torch.tensor([[5, 5, 5]])
    prompt_mask = torch.tensor([[1, 1, 1]])
    gen_output = torch.tensor([[5, 5, 5, EOS, PAD, PAD]])  # 第一步就 eos

    _, attn, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)

    assert labels[0].tolist() == [IGNORE_INDEX] * 3 + [EOS, IGNORE_INDEX, IGNORE_INDEX]
    assert attn[0].tolist() == [1, 1, 1, 1, 0, 0]


def test_prompt段恒为负100():
    prompts = torch.tensor([[PAD, PAD, 5, 5]])
    prompt_mask = torch.tensor([[0, 0, 1, 1]])
    gen_output = torch.tensor([[PAD, PAD, 5, 5, 8, 9]])

    _, _, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)
    assert labels[0, :4].tolist() == [IGNORE_INDEX] * 4  # prompt 段（含左 pad）全 -100