"""Regression: BF16 counters above 256 can stop incrementing after resume."""

import json

import torch

from ngram_train_resume_safe import restore_adamw_steps

parameter = torch.nn.Parameter(torch.ones(3))
optimizer = torch.optim.AdamW([parameter], lr=0.01)
parameter.sum().backward()
optimizer.step()
counter = torch.tensor(1001.0, dtype=torch.bfloat16)
before = counter.item()
counter.add_(1)
assert counter.item() == before == 1000.0
optimizer.state[parameter]["step"] = counter
assert restore_adamw_steps([optimizer], 1001) == 1
assert optimizer.state[parameter]["step"].dtype == torch.float32
optimizer.zero_grad()
parameter.sum().backward()
optimizer.step()
assert optimizer.state[parameter]["step"].item() == 1002
optimizer.state[parameter]["step"] = torch.tensor(900.0)
try:
    restore_adamw_steps([optimizer], 1001)
except ValueError:
    pass
else:
    raise AssertionError("Mismatched counters must not be silently overwritten")

# The real checkpointer saves BF16 then promotes all floating states to the
# model dtype on load. Large, non-BF16-exact training positions must still pass.
checkpoint_step = 25680
saved = optimizer.state_dict()
saved["state"][0]["step"] = torch.tensor(checkpoint_step, dtype=torch.bfloat16).float()
optimizer.load_state_dict(saved)
assert optimizer.state[parameter]["step"].dtype == torch.float32
assert optimizer.state[parameter]["step"].item() == 25728
moments = {
    name: optimizer.state[parameter][name].clone()
    for name in ("exp_avg", "exp_avg_sq")
}
assert restore_adamw_steps([optimizer], checkpoint_step) == 1
assert optimizer.state[parameter]["step"].item() == checkpoint_step
for name, value in moments.items():
    torch.testing.assert_close(optimizer.state[parameter][name], value, rtol=0, atol=0)
optimizer.zero_grad()
parameter.sum().backward()
optimizer.step()
assert optimizer.state[parameter]["step"].item() == checkpoint_step + 1
print(
    json.dumps(
        {
            "passed": True,
            "legacy_bf16_counter_stalled_at": before,
            "repaired_fp32_counter_after_update": 1002,
            "mismatched_counter_rejected": True,
            "bf16_promoted_to_fp32_roundtrip_passed": True,
            "moments_unchanged_by_repair": True,
            "large_counter_after_update": checkpoint_step + 1,
        }
    )
)