"""CPU checks for teacher identity, detachment, and native-reference null signal."""
import copy
import json
from pathlib import Path
import torch

from run_baseline import FastModel, QWEN_CONFIG_06_B
from latent_deliberation import LatentDeliberation
from contextual_rl_common import accumulate_question
from reference_contextual_operator import accumulate_reference_question


class Tokens:
    eos_token_id = 2

    def decode(self, ids):
        return ' '.join('t' + str(i) for i in ids)

    def encode(self, text):
        return [3 + sum(map(ord, word)) % 61 for word in text.split()]


def gradients(model):
    return torch.cat([(torch.zeros_like(p) if p.grad is None else p.grad).flatten()
                      for p in model.parameters()])


def main():
    torch.set_num_threads(2)
    torch.manual_seed(21641202)
    cfg = dict(QWEN_CONFIG_06_B, vocab_size=64, emb_dim=32, n_heads=4,
               n_kv_groups=2, head_dim=8, n_layers=2, hidden_dim=64,
               context_length=512, dtype=torch.float32)
    initial = LatentDeliberation(FastModel(cfg), 0, True)
    tok = Tokens()
    args = (tok, [3, 4, 5], [[8, 9, 2], [10, 11, 2], [12, 13, 14],
                           [15, 2], [16, 17, 18, 2], [19, 2], [20, 21, 2], [22, 2]],
            [1., 0., 0., 1., 0., 0., 0., 0.],
            [True, True, False, True, True, True, True, True])
    kwargs = dict(batch_questions=1, topk=7, chunk=2)
    old = copy.deepcopy(initial)
    old_record = accumulate_question(old, *args, arm='C', **kwargs)
    current = copy.deepcopy(initial)
    current_record = accumulate_reference_question(current, current, *args,
                                                   context_mode='contextual', **kwargs)
    torch.testing.assert_close(gradients(current), gradients(old), atol=1e-6, rtol=1e-5)
    assert current_record['work'] == old_record['work']
    assert current_record['receivers'] == [1, 4, 5, 6, 7]
    assert current_record['backward_calls'] == 7

    student = copy.deepcopy(initial)
    teacher = copy.deepcopy(initial).requires_grad_(False)
    before = {k: v.clone() for k, v in teacher.state_dict().items()}
    accumulate_reference_question(student, teacher, *args,
                                  context_mode='contextual', **kwargs)
    torch.testing.assert_close(gradients(student), gradients(old), atol=1e-6, rtol=1e-5)
    assert all(p.grad is None for p in teacher.parameters())
    torch.optim.SGD(student.parameters(), lr=.1).step()
    assert all(torch.equal(before[k], v) for k, v in teacher.state_dict().items())
    assert any(not torch.equal(before[k], v) for k, v in student.state_dict().items())

    ordinary = copy.deepcopy(initial)
    accumulate_question(ordinary, *args, arm='G', **kwargs)
    native = copy.deepcopy(initial)
    native_record = accumulate_reference_question(native, teacher, *args,
                                                  context_mode='native', **kwargs)
    assert abs(native_record['dense_jsd_sum']) < 1e-5
    torch.testing.assert_close(gradients(native), gradients(ordinary), atol=1e-6, rtol=1e-5)

    changed_teacher = copy.deepcopy(teacher)
    with torch.no_grad():
        changed_teacher.base.out_head.weight.mul_(-1)
    changed_student = copy.deepcopy(initial)
    accumulate_reference_question(changed_student, changed_teacher, *args,
                                  context_mode='contextual', **kwargs)
    effect = float((gradients(changed_student) - gradients(old)).abs().max())
    assert effect > 1e-6
    assert all(p.grad is None for p in changed_teacher.parameters())

    empty_student = copy.deepcopy(initial)
    zero_args = (*args[:3], [0.] * 8, args[4])
    zero_record = accumulate_reference_question(empty_student, teacher, *zero_args,
                                                context_mode='contextual', **kwargs)
    assert zero_record['backward_calls'] == 0
    assert all(p.grad is None for p in empty_student.parameters())
    result = {
        'status': 'PASS', 'device': 'CPU',
        'checks': ['current-teacher gradient/work agreement with completed operator',
                   'identical frozen teacher gives matching initial student gradient',
                   'teacher receives no gradient and remains exact after student update',
                   'native identical-weight teacher has zero auxiliary signal',
                   'changed teacher changes the student gradient',
                   'all-wrong/no-donor group remains zero signal'],
        'changed_teacher_max_student_gradient_difference': effect,
        'scope': 'Tiny-model implementation check; no full-size GPU or native learning claim.',
    }
    destination = Path(__file__).resolve().parents[1] / 'notes/reference_contextual_operator_cpu_check.json'
    destination.write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()
