"""ZeroPoint 架构自检。运行: python3 tests/test_zeropoint.py"""
import itertools
import json
import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, ".."))
sys.path.insert(0, HERE)
import zerocg as zc  # noqa: E402
from models import ZeroX  # noqa: E402
from models.zeropoint import CHUNK, MemoryMixer  # noqa: E402
from test_zerocg import module_grad_check  # noqa: E402

zc.set_default_dtype(np.float64)
rng = np.random.default_rng(0)

# 所有开关的组合（消融实验会用到的每一种都要验证两种算法一致）
OPTIONS = [dict(zip(("gate", "write", "init", "learn_lambda", "pos", "ema"), combo))
           for combo in itertools.product((True, False), ("softmax", "sigmoid"), ("hier", "uniform"),
                                          (True, False), ("none", "sin"), (False, True))]


def test_parallel_equals_recurrent():
    """训练用的分块并行形式和推理用的递推形式，在所有开关组合下都必须一致。
    长度 300：跨好几个块，也远超旧版并行形式 120 字的上限。"""
    idx = rng.integers(0, 11, size=(2, 300))
    worst = 0.0
    for opts in OPTIONS:
        np.random.seed(1)
        model = ZeroX(vocab_size=11, n_embd=16, n_head=2, n_slot=5, n_layer=2, **opts).eval()
        worst = max(worst, np.abs(model(idx).data - model.forward_recurrent(idx)).max())
    assert worst < 1e-8, worst


def test_matches_golden_reference():
    """和 tests/golden/zeropoint_ref.json 对答案：防止以后改代码时并行或递推形式悄悄变了。"""
    with open(os.path.join(HERE, "golden", "zeropoint_ref.json"), encoding="utf-8") as f:
        ref = json.load(f)
    cfg = {k[7:]: int(v) for k, v in ref.items() if k.startswith("config:")}
    model = ZeroX(**cfg).eval()
    model.load_state_dict({k[6:]: np.array(v, dtype=np.float64) for k, v in ref.items()
                           if k.startswith("param:")})
    idx, logits = np.array(ref["idx"]), np.array(ref["logits"])
    assert np.abs(model(idx).data - logits).max() < 1e-10
    assert np.abs(model.forward_recurrent(idx) - logits).max() < 1e-8


def test_causal_across_chunks():
    np.random.seed(2)
    model = ZeroX(vocab_size=7, n_embd=16, n_head=2, n_slot=4).eval()
    idx = rng.integers(0, 7, size=(1, 3 * CHUNK))
    y1 = model(idx).data
    p = CHUNK + 5            # 改第二块里的一个字
    idx[0, p] = (idx[0, p] + 1) % 7
    y2 = model(idx).data
    assert np.allclose(y1[0, :p], y2[0, :p]) and not np.allclose(y1[0, p:], y2[0, p:])


def test_ema_state_is_bounded():
    """ema=True：写入恒为常数 v 时，槽会收敛到 v 本身，不会随长度无限变大（ema=False 时会涨到约 v/(1-λ)）。"""
    np.random.seed(6)
    for ema, bounded in ((True, True), (False, False)):
        mixer = MemoryMixer(8, 1, 4, ema=ema)
        S = mixer.init_state(1) * 0
        x = np.ones((1, 8))
        for _ in range(3000):
            _, S = mixer.step(x, S)
        v = np.abs(x @ mixer.value.weight.data + mixer.value.bias.data).max()
        assert (np.abs(S).max() <= v * 1.0001) == bounded, (ema, np.abs(S).max(), v)


def test_mixer_gradients():
    np.random.seed(3)
    module_grad_check(MemoryMixer(8, 2, 3), rng.standard_normal((2, 6, 8)))


def test_mixer_gradients_across_chunks():
    """梯度要能穿过块与块之间传递的状态，也要传到可学习的初始状态 s0。"""
    np.random.seed(3)
    mixer = MemoryMixer(4, 1, 2)
    module_grad_check(mixer, rng.standard_normal((1, CHUNK + 3, 4)))
    assert np.abs(mixer.s0.grad).max() > 0


def test_initial_state_is_learned():
    model = ZeroX(vocab_size=5, n_embd=8, n_head=2, n_slot=3, n_layer=1)
    assert np.allclose(model.init_state(3)[0][2], model.blocks[0].mix.s0.data)
    assert "blocks.0.mix.s0" in dict(model.named_parameters())


def test_retention_init_is_hierarchical():
    lam = MemoryMixer(16, 2, 6, max_half_life=64).retention().data
    half_life = np.log(0.5) / np.log(lam[0])
    assert np.all(np.diff(lam[0]) > 0) and abs(half_life[0] - 1.5) < 1e-3 and abs(half_life[-1] - 64) < 1e-2


def test_usage_recorded():
    model = ZeroX(vocab_size=5, n_embd=8, n_head=2, n_slot=3, n_layer=1)
    model(rng.integers(0, 5, size=(2, 10)))
    read, write = model.blocks[0].mix.usage
    assert read.shape == write.shape == (2, 3) and np.allclose(read.sum(-1), 1)


def test_state_is_constant_size():
    model = ZeroX(vocab_size=5, n_embd=16, n_head=2, n_slot=4, n_layer=3)
    states = model.init_state(1)
    for t in range(300):
        _, states = model.step(np.array([t % 5]), states, t)
    assert sum(s.size for s in states) == model.state_floats() == 3 * 2 * 4 * 8
    assert all(np.isfinite(s).all() for s in states)


def test_float32_long_sequence_is_finite():
    """float32 下 1000 个字的分块并行不溢出，并且和递推形式一致。"""
    zc.set_default_dtype(np.float32)
    try:
        np.random.seed(5)
        model = ZeroX(vocab_size=9, n_embd=16, n_head=2, n_slot=6).eval()
        idx = rng.integers(0, 9, size=(1, 1000))
        par = model(idx).data
        assert par.dtype == np.float32 and np.isfinite(par).all()
        assert np.abs(par - model.forward_recurrent(idx)).max() < 1e-3
    finally:
        zc.set_default_dtype(np.float64)


def test_learns_to_copy_from_far_back():
    """记忆任务：开头一个字，中间 40 个干扰字（跨块），最后复述开头那个字。"""
    np.random.seed(4)
    V, gap = 6, 40
    model = ZeroX(vocab_size=V + 1, n_embd=32, n_head=2, n_slot=8, n_layer=1)
    opt = zc.optim.AdamW(model.parameters(), lr=0.01)
    for _ in range(300):
        first = np.random.randint(0, V, 32)
        idx = np.concatenate([first[:, None], np.full((32, gap), V)], axis=1)
        loss = zc.functional.cross_entropy(model(idx)[np.arange(32), gap], first)
        opt.zero_grad()
        loss.backward()
        opt.step()
    assert loss.item() < 0.05, loss.item()


if __name__ == "__main__":
    tests = [(k, v) for k, v in list(globals().items()) if k.startswith("test_")]
    failed = 0
    for name, fn in tests:
        try:
            fn()
            print(f"  ok   {name}")
        except Exception as e:  # noqa: BLE001
            failed += 1
            print(f"  FAIL {name}: {type(e).__name__}: {e}")
    print(f"\n{len(tests) - failed}/{len(tests)} 通过")
    sys.exit(1 if failed else 0)