"""Source integrity and native image-codebook-decode tests; no language model required."""
import gc
import hashlib
import json
import time
from pathlib import Path
import torch
from safetensors.torch import load_file
from component import load_model

root = Path(__file__).parent
torch.set_num_threads(4)
checks = []
started = time.monotonic()


def check(name, ok, **details):
    assert ok, (name, details)
    checks.append({'name': name, 'passed': True, **details})
    print(name, details, flush=True)


with (root / 'model.safetensors').open('rb') as f:
    file_sha = hashlib.file_digest(f, 'sha256').hexdigest()
check('official_LFS_SHA256', file_sha == 'a808dc617e489c37be7f51a63a0e7cb77a97a54ceb59f674255d2e7bc7b2c080')
m = load_model(root)
source = load_file(root / 'model.safetensors')
state = m.state_dict()
check('native_strict_543_entries', set(state) == set(source) and len(state) == 543)
check('all_original_tensor_values', all(torch.equal(source[k], state[k]) for k in source))
del source, state
with torch.inference_mode():
    for height, width in [(256, 256), (128, 192)]:
        label = f'{height}x{width}'
        x = torch.linspace(-1, 1, 3 * height * width).reshape(1, 3, height, width)
        q, _, info = m.encode(x)
        ids = info[-1]
        check(label + '_encode_geometry', tuple(q.shape) == (1, 256, height // 16, width // 16), shape=list(q.shape))
        check(label + '_ID_range', ids.numel() == height * width // 256 and ids.min().item() >= 0 and ids.max().item() < 131072)
        lookup = m.quantize.get_codebook_entry(ids, shape=(1, height // 16, width // 16))
        check(label + '_lookup_exact', torch.equal(q, lookup))
        h = m.quant_conv(m.encoder(x))
        expected = (h.permute(0, 2, 3, 1).reshape(-1, 256) @ m.quantize.embedding.weight.T).argmax(-1)
        check(label + '_max_dot_assignment', torch.equal(ids, expected))
        decoded = m.decode(q)
        by_code = m.decode_code(ids, shape=(1, height // 16, width // 16))
        check(label + '_decode_code_exact', torch.equal(decoded, by_code))
        check(label + '_reconstruction', decoded.shape == x.shape and torch.isfinite(decoded).all().item(), mse=(decoded - x).square().mean().item())
        del q, lookup, h, decoded, by_code
        gc.collect()
report = {'strict_load': True, 'source_tensor_bytes_equal': True, 'checks': checks, 'smoke': {'passed': True, 'device': 'cpu', 'dtype': 'float32'}, 'elapsed_seconds': time.monotonic() - started, 'limitations': ['Synthetic codec fixtures only; no visual quality benchmark or full Emu3.5 language/generation backbone evaluation.', 'This is a 2D image VQ codec. No temporal video codec behavior is asserted.']}
(root / 'validation.json').write_text(json.dumps(report, indent=2) + '\n')
print('EMU3.5 VALIDATED', len(checks), flush=True)