import hashlib
import json
import time
from pathlib import Path
import torch
import torch.nn.functional as F
from safetensors import safe_open
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)


expected_hash = json.loads((root / 'tensor_sha256.json').read_text())
with safe_open(root / 'model.safetensors', framework='pt') as f:
    for k in f.keys():
        t = f.get_tensor(k)
        assert hashlib.sha256(t.contiguous().view(torch.uint8).numpy().tobytes()).hexdigest() == expected_hash[k]
check('649_source_tensor_payloads', len(expected_hash) == 649)
m = load_model(root)
check('native_strict_load', len(m.state_dict()) == 649)
check('official_scale_schedule', m.scale_rq_layers == [1, 2, 4, 6, 8, 10, 12, 14, 16])
with torch.inference_mode():
    x = torch.linspace(-1, 1, 3*256*256).reshape(1, 3, 256, 256)
    q, losses, ids = m.encode(x)
    check('multiscale_token_geometry', tuple(ids.shape) == (1, sum(s*s for s in m.scale_rq_layers)), shape=list(ids.shape))
    check('ID_bounds', ids.min().item() >= 0 and ids.max().item() < 32768)
    canonical = 0.
    cursor = 0
    for scale in m.scale_rq_layers:
        n = scale*scale
        vectors = m.quantize.get_codebook_entry(ids[:, cursor:cursor+n])
        vectors = vectors.reshape(1, scale, scale, m.code_dim).permute(0, 3, 1, 2)
        canonical = canonical + F.interpolate(vectors, size=(16,16), mode='bicubic')
        cursor += n
    check('residual_lookup_bound', torch.allclose(q, canonical, rtol=2e-5, atol=2e-5), max_abs=(q-canonical).abs().max().item())
    semantic, rgb = m.decode(q)
    check('semantic_output', tuple(semantic.shape) == (1, 256, 512) and torch.isfinite(semantic).all().item(), shape=list(semantic.shape))
    check('RGB_output', rgb.shape == x.shape and torch.isfinite(rgb).all().item(), shape=list(rgb.shape), mse=(rgb-x).square().mean().item())
    by_ids = m.decode_code(ids)
    by_lookup = m.decode(canonical)
    check('semantic_decode_code_lookup_exact', torch.equal(by_ids[0], by_lookup[0]))
    check('RGB_decode_code_lookup_exact', torch.equal(by_ids[1], by_lookup[1]))
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 256px native codec fixtures, not image quality benchmarks or complete VQA/T2I pipelines.', 'Includes the trained CLIP visual branch, semantic and pixel codebooks, semantic decoder and enhanced RGB decoder. No CLIP text model or external pretrained download required.', 'Native straight-through arithmetic can differ slightly from direct lookup; canonical lookup and decode-code paths are compared exactly.']}
(root / 'validation.json').write_text(json.dumps(report, indent=2)+'\n')
print('TOKENFLOW VALIDATED',len(checks),flush=True)