"""Tests for the ADR-013 altered_minds integration glue + the B4 SDPO-fires proof.

Covers the ADR-013 acceptance gate:
  - MMLUFormatReward: correct→+1, wrong→0, unparseable→−0.2, multiple→−0.1,
    length-penalty, and an "always C" option-prior exploit is DETECTABLE via the
    logged option distribution. Rationale style is NOT scored.
  - dual_kl_logger: KL(p‖p)==0 and KL grows as the policy moves.
  - channel_ladder_configs: A1 both off, A2 SDPO-only, A3 replay-only.
  - B4: the SDPO channel actually FIRES (NONZERO loss) with REAL collator-built
    alignment indices. See the module docstring on test_b4_* for the honest
    stub-vs-real note.

All CPU-only and fast (stub tokenizer + tiny model — no model download).
"""
from __future__ import annotations

import pytest
import torch

from composer_replication.integrations.altered_minds import (
    MMLUFormatReward,
    channel_ladder_configs,
    dual_kl_logger,
    randomize_options,
)
from composer_replication.integrations.altered_minds.reward import parse_final_answer


def test_hedged_answer_is_penalized_not_full_credit():
    """Final-verify 2026-05-29: 'Answer: A or B' / 'Answer: A/B' must NOT score
    full credit for A — a hedge naming a second distinct option is a multiple-
    answers format hack (n_distinct >= 2)."""
    for hedge in ("Answer: A or B", "Answer: A/B", "Answer: A, B", "Answer: A and C"):
        letter, n_distinct = parse_final_answer(hedge)
        assert n_distinct >= 2, f"hedge {hedge!r} not detected (n_distinct={n_distinct})"
    # A clean single answer is still n_distinct == 1.
    letter, n_distinct = parse_final_answer("After reasoning, Answer: C")
    assert letter == "C" and n_distinct == 1
    # Two clean markers of the SAME letter are not a hedge.
    _, n_same = parse_final_answer("Answer: C ... wait, Answer: C")
    assert n_same == 1


def test_hedged_answer_scores_multiple_penalty_via_reward():
    r = MMLUFormatReward()
    out = r(prompts=["p"], completions=["Answer: A or B"], answers=["A"])
    # Even though the gold is A and A is the lead letter, the hedge is penalized.
    assert out[0] == r.multiple_answers_reward


# ===========================================================================
# MMLUFormatReward
# ===========================================================================

def test_reward_correct_wrong_unparseable_multiple():
    r = MMLUFormatReward()
    completions = [
        "Reasoning blah. Answer: B",                 # correct
        "I think it's Answer: A",                    # wrong (gold C)
        "no marker here at all",                     # unparseable
        "Answer: A then actually Answer: D",         # multiple distinct
        '{"answer": "C"}',                           # JSON correct
    ]
    answers = ["B", "C", "B", "A", "C"]
    out = r(prompts=None, completions=completions, answers=answers)
    assert out[0] == pytest.approx(1.0)    # correct
    assert out[1] == pytest.approx(0.0)    # wrong
    assert out[2] == pytest.approx(-0.2)   # unparseable
    assert out[3] == pytest.approx(-0.1)   # multiple distinct markers
    assert out[4] == pytest.approx(1.0)    # JSON correct


def test_reward_last_match_wins_same_letter_not_penalized():
    """Two markers of the SAME letter is not 'multiple distinct' — last wins."""
    r = MMLUFormatReward()
    out = r(completions=["Answer: C ... so my final Answer: C"], answers=["C"])
    assert out[0] == pytest.approx(1.0)


def test_reward_case_insensitive_and_json_variants():
    r = MMLUFormatReward()
    out = r(
        completions=["answer: d", '{"answer":"a"}'],
        answers=["D", "A"],
    )
    assert out[0] == pytest.approx(1.0)
    assert out[1] == pytest.approx(1.0)


def test_reward_length_penalty_only_past_cap():
    """A correct-but-long completion is penalized by ~0.001/char past the cap;
    a short one is not. Rationale CONTENT is never scored — only length."""
    r = MMLUFormatReward(rationale_char_cap=20, length_penalty_per_char=0.001)
    short = "Answer: B"                       # under cap
    long = "x" * 120 + " Answer: B"           # ~130 chars, 110 over cap
    out = r(completions=[short, long], answers=["B", "B"])
    assert out[0] == pytest.approx(1.0)
    # 130 - 20 = 110 over => penalty 0.110; reward 1.0 - 0.110
    assert out[1] < 1.0
    assert out[1] == pytest.approx(1.0 - 0.001 * (len(long) - 20))


def test_reward_always_C_exploit_is_detectable():
    """An 'always C' policy that happens to be right when gold==C scores well on
    those items, but the logged option distribution reveals the exploit."""
    r = MMLUFormatReward()
    completions = [f"Answer: C" for _ in range(10)]
    golds = ["C", "A", "B", "C", "D", "A", "C", "B", "C", "D"]
    r(completions=completions, answers=golds)
    report = r.exploit_report()
    assert report["most_common"] == "C"
    # Every parsed answer was C => fraction 1.0 — the exploit signature.
    assert report["max_fraction"] == pytest.approx(1.0)
    assert report["counts"] == {"C": 10}


def test_reward_requires_answers():
    r = MMLUFormatReward()
    with pytest.raises(ValueError, match="requires `answers`"):
        r(completions=["Answer: A"])


def test_randomize_options_tracks_label_remap_and_updates_gold():
    item = {"question": "q", "options": ["w", "x", "y", "z"], "answer": "A"}
    shuffled, remap = randomize_options(item, seed=7)
    # All four letters map to four distinct new letters (a permutation).
    assert sorted(remap.keys()) == ["A", "B", "C", "D"]
    assert sorted(remap.values()) == ["A", "B", "C", "D"]
    # The gold option's text ("w", originally A) now lives at its remapped letter.
    new_gold_letter = shuffled["answer"]
    new_gold_idx = ord(new_gold_letter) - ord("A")
    assert shuffled["options"][new_gold_idx] == "w"
    assert remap["A"] == new_gold_letter
    # Determinism.
    shuffled2, remap2 = randomize_options(item, seed=7)
    assert remap == remap2 and shuffled["options"] == shuffled2["options"]


# ===========================================================================
# dual_kl_logger
# ===========================================================================

def test_dual_kl_self_is_zero():
    """KL(p‖p) == 0 for both diagnostics."""
    logits = torch.randn(2, 5, 16)
    out = dual_kl_logger(logits, logits, logits)
    assert out["kl_to_altered_init"] == pytest.approx(0.0, abs=1e-6)
    assert out["kl_to_base"] == pytest.approx(0.0, abs=1e-6)


def test_dual_kl_grows_as_policy_moves():
    """As the policy distribution moves further from a fixed reference, the KL
    grows monotonically. Both diagnostics are non-negative."""
    torch.manual_seed(0)
    ref = torch.randn(1, 4, 16)
    base = torch.randn(1, 4, 16)

    near = ref + 0.1 * torch.randn_like(ref)
    far = ref + 2.0 * torch.randn_like(ref)

    kl_near = dual_kl_logger(near, ref, base)["kl_to_altered_init"]
    kl_far = dual_kl_logger(far, ref, base)["kl_to_altered_init"]

    assert kl_near >= -1e-9
    assert kl_far > kl_near, f"KL should grow as policy moves: {kl_near} -> {kl_far}"


def test_dual_kl_mask_restricts_tokens():
    """A token mask restricts the mean to the masked answer+reasoning tokens."""
    torch.manual_seed(1)
    policy = torch.randn(1, 4, 8)
    ref = torch.randn(1, 4, 8)
    base = torch.randn(1, 4, 8)
    mask = torch.tensor([[1, 1, 0, 0]])
    out = dual_kl_logger(policy, ref, base, mask=mask)
    # Masked-all-zero => 0.0 (guarded), nonzero mask => finite non-negative.
    assert out["kl_to_altered_init"] >= -1e-9
    zero = dual_kl_logger(policy, ref, base, mask=torch.zeros(1, 4))
    assert zero["kl_to_altered_init"] == 0.0
    assert zero["kl_to_base"] == 0.0


# ===========================================================================
# channel_ladder_configs
# ===========================================================================

def test_ladder_arms_and_order():
    arms = channel_ladder_configs()
    assert [a["arm"] for a in arms] == ["A0", "A1", "A2", "A3", "A4"]


def test_ladder_a0_is_no_rl_sentinel():
    a0 = channel_ladder_configs()[0]
    assert a0["arm"] == "A0"
    assert a0["alpha_sdpo"] is None
    assert a0["beta_replay"] is None
    assert a0["kl_beta"] is None


def test_ladder_a1_both_off():
    a1 = channel_ladder_configs()[1]
    assert a1["alpha_sdpo"] == 0.0
    assert a1["beta_replay"] == 0.0
    assert a1["kl_beta"] == 0.02


def test_ladder_a2_sdpo_only():
    a2 = channel_ladder_configs()[2]
    assert a2["alpha_sdpo"] == 0.02
    assert a2["beta_replay"] == 0.0
    assert a2["kl_beta"] == 0.02


def test_ladder_a3_replay_only():
    a3 = channel_ladder_configs()[3]
    assert a3["alpha_sdpo"] == 0.0
    assert a3["beta_replay"] == 0.05
    assert a3["kl_beta"] == 0.02


def test_ladder_a4_combined():
    a4 = channel_ladder_configs()[4]
    assert a4["alpha_sdpo"] == 0.02
    assert a4["beta_replay"] == 0.05


# ===========================================================================
# B4 — the SDPO channel actually FIRES (NONZERO) with REAL collator indices
# ===========================================================================
#
# HONEST NOTE ON STUB-VS-REAL (ADR-013 B4 acceptance):
#
# This proof uses the same TinyLM stub pattern as
# trainer/tests/test_sdpo_alignment_indices.py, NOT a real Qwen checkpoint
# (kept offline/CPU and deterministic). The alignment indices are REAL: they are
# built by the production ComposerDataCollator from a trace that HAS an error
# turn (so ctx_teacher_input_ids + student/teacher_response_idx are genuinely
# emitted by the shipped collator, exactly as in a real run).
#
# Why we must perturb the student tokens to get a NONZERO loss: the collator's
# placeholder-alignment trick makes student and teacher carry the SAME token ids
# at the SAME absolute positions at valid aligned indices, so a deterministic
# stub yields JSD≈0 there (the CORRECT answer for a perfectly-aligned identical
# model — see that test's gate-3 note). To prove the channel genuinely GATHERS
# the aligned positions and computes nonzero divergence, we make the student's
# input_ids DIFFER from the teacher's at exactly the aligned response positions
# — this mimics the hint actually changing the recovery tokens (the real-world
# case where SDPO has a signal to distill). With a position-dependent stub,
# different aligned token ids => different logits => provably NONZERO JSD on a
# grad path, through the real collator-built indices.

from composer_replication.trainer.data_collator import (  # noqa: E402
    CollatorConfig,
    ComposerDataCollator,
)


class _StubTok:
    """Word-level deterministic tokenizer; apply_chat_template space-joins."""

    pad_token_id = 0

    def __init__(self) -> None:
        self._v: dict[str, int] = {"<pad>": 0, "<bos>": 1, "<eos>": 2}

    def _id(self, w: str) -> int:
        if w not in self._v:
            self._v[w] = len(self._v)
        return self._v[w]

    def __call__(self, text, **_k):
        return {"input_ids": [self._id(w) for w in text.split()] if text else []}

    def apply_chat_template(self, messages, tokenize=True, **_k):  # noqa: ARG002
        return [self._id(w) for w in " ".join(m.get("content", "") for m in messages).split()]


class _TinyLM(torch.nn.Module):
    """Position-dependent minimal model: model(input_ids=...).logits."""

    def __init__(self, vocab: int = 64, hidden: int = 8, max_pos: int = 512):
        super().__init__()
        torch.manual_seed(0)
        self.embed = torch.nn.Embedding(vocab, hidden)
        self.pos = torch.nn.Embedding(max_pos, hidden)
        self.head = torch.nn.Linear(hidden, vocab)

    def forward(self, input_ids: torch.Tensor):
        T = input_ids.size(1)
        positions = torch.arange(T, device=input_ids.device).unsqueeze(0)
        h = self.embed(input_ids) + self.pos(positions)

        class _Out:
            pass

        out = _Out()
        out.logits = self.head(h)
        return out


def _hint_gen(_kind, _meta):
    return "HINT search before reading"


def _error_trace(trace_id: str, recovery: str = "let me use a real tool instead now"):
    return {
        "trace_id": trace_id,
        "turns": [
            {"role": "user", "content": "do the task now"},
            {"role": "user", "content": "tool not found error occurred"},
            {
                "role": "assistant",
                "content": recovery,
                "tool_error": "tool_not_found",
                "error_meta": {},
            },
        ],
        "final_reward": 0.0,
    }


def _make_sdpo_trainer(alpha_sdpo: float):
    from composer_replication.trainer.composer_trainer import ComposerReplicationTrainer

    obj = ComposerReplicationTrainer.__new__(ComposerReplicationTrainer)
    obj.alpha_sdpo = alpha_sdpo
    obj.sdpo_jsd_beta = 0.5
    obj.sdpo_temperature = 1.0
    obj.sdpo_token_clip = None
    obj.strict_sdpo_alignment = True  # production default
    return obj


def test_b4_sdpo_fires_nonzero_with_real_collator_indices():
    """B4: with REAL collator-built alignment indices and the student tokens
    differing from the teacher at the aligned response positions (hint changed
    the recovery tokens), the SDPO channel gathers those positions and produces
    a NONZERO JSD on a grad path — proving the channel actually FIRES."""
    tok = _StubTok()
    cfg = CollatorConfig(hint_generator=_hint_gen, enable_replay_dpo=False)
    collator = ComposerDataCollator(tokenizer=tok, config=cfg)
    batch = collator([_error_trace("b4-fires")])

    # Sanity: the collator genuinely emitted error-site teacher context + indices.
    assert batch["ctx_teacher_input_ids"].numel() > 0
    s_idx = batch["student_response_idx"]
    t_idx = batch["teacher_response_idx"]
    s_valid = batch["student_response_valid"]
    assert int(s_valid.sum()) > 0, "no valid aligned positions — collator emitted nothing"

    # Perturb the STUDENT tokens at the aligned response positions so they differ
    # from the teacher's tokens there (the hint changed the recovery tokens). We
    # keep the REAL collator-built indices; only the student input_ids change.
    student_ids = batch["input_ids"].clone()
    vocab_ceiling = int(
        max(batch["input_ids"].max(), batch["ctx_teacher_input_ids"].max())
    ) + 8
    for b in range(s_idx.shape[0]):
        for k in range(s_idx.shape[1]):
            if bool(s_valid[b, k]):
                pos = int(s_idx[b, k])
                # bump to a different, in-vocab token id (deterministic).
                student_ids[b, pos] = (int(student_ids[b, pos]) + 3) % vocab_ceiling
    batch["input_ids"] = student_ids

    model = _TinyLM(vocab=max(vocab_ceiling, 8))
    obj = _make_sdpo_trainer(alpha_sdpo=0.02)  # A2 config (SDPO-only small)

    loss = obj._compute_sdpo_loss(model, batch)
    val = float(loss.detach())

    assert val == val and val not in (float("inf"), float("-inf")), "loss not finite"
    assert loss.requires_grad, "SDPO loss must be on a grad path"
    assert val > 1e-6, (
        f"SDPO channel did not fire: JSD={val} (expected NONZERO once the "
        "aligned student/teacher tokens differ). The channel must gather the "
        "real collator indices and compute a positive divergence."
    )

    # Prove it is differentiable end-to-end: backward populates a real gradient.
    (obj.alpha_sdpo * loss).backward()
    grad_norm = sum(
        float(p.grad.norm()) for p in model.parameters() if p.grad is not None
    )
    assert grad_norm > 0.0, "no gradient flowed from the SDPO loss into the model"