#!/usr/bin/env python3
"""Reference-only check of Vision-Exp weights and the official PyTorch image graph.

Production code never invokes this script. Inputs are the pinned DS4 source and
DeepSeek source revisions recorded in docs/issue100-vision, and native dumps
from the ignored issue100_encoder_real_weights test.
"""
import argparse
import importlib.util
import json
from pathlib import Path
import sys
from types import SimpleNamespace

import numpy as np
import torch

p = argparse.ArgumentParser()
p.add_argument('--ds4-source', type=Path, required=True)
p.add_argument('--official-source', type=Path, required=True)
p.add_argument('--artifacts', type=Path, required=True)
p.add_argument('--native', type=Path, required=True)
p.add_argument('--image', type=Path, required=True)
p.add_argument('--device', default='mps')
a = p.parse_args()
parser_path = a.ds4_source / 'gguf-tools/mixed/splice_mixed_expert_layers_gguf.py'
spec = importlib.util.spec_from_file_location('ds4_gguf', parser_path)
gguf = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = gguf
spec.loader.exec_module(gguf)
sys.path.insert(0, str(a.official_source / 'inference'))
from vision import ViT, Aligner
from image_processor import load_image, build_image_block

config = json.loads((a.official_source / 'config.json').read_text())
config['dim'] = config['hidden_size']
args = SimpleNamespace(**config)
file = a.artifacts / 'DeepSeek-V4-Flash-Vision-Encoder.gguf'
gguf.GGML_QUANT_SIZES[30] = (1, 2, "BF16")  # GGML BF16, raw source weights.
info = gguf.parse_gguf(file)
torch.set_default_dtype(torch.bfloat16)
torch.set_default_device(a.device)
vit, aligner = ViT(args), Aligner(args)
with file.open('rb') as f:
    def state(prefix):
        result = {}
        for name, t in info.tensor_by_name.items():
            if not name.startswith(prefix):
                continue
            assert t.ggml_type == 30, (name, t.ggml_type)
            f.seek(t.data_offset)
            values = np.frombuffer(f.read(t.n_bytes), dtype='<u2').copy()
            weight = torch.from_numpy(values).view(torch.bfloat16).reshape(tuple(reversed(t.dims)))
            result[name.removeprefix(prefix)] = weight.to(a.device)
        return result
    vit.load_state_dict(state('vision.'), strict=True)
    aligner.load_state_dict(state('aligner.'), strict=True)
vit.eval(); aligner.eval()
patches, h, w, lh, lw = load_image({'url': str(a.image)}, args)
meta = json.loads(a.native.with_suffix('.json').read_text())
assert meta['grid'] == [h, w]
assert meta['aligned_grid'] == [lh, lw]
with torch.inference_mode():
    output = aligner(vit(patches.to(a.device), h, w), h, w).float().cpu().numpy()
native = np.fromfile(a.native.with_suffix('.embedding.f32'), dtype='<f4').reshape(output.shape)
error = native - output
result = {
    'source_revision': 'e46e16bf6035c6f317eb2ac7458eb0362926d402',
    'ds4_revision': '0aaea5a238fb41a35106a551e73c8409dfb751ac',
    'torch': torch.__version__, 'device': a.device,
    'shape': list(output.shape),
    'max_abs': float(np.abs(error).max()),
    'rmse': float(np.sqrt(np.mean(error**2))),
    'reference_rms': float(np.sqrt(np.mean(output**2))),
    'cosine': float(np.dot(native.ravel(), output.ravel()) / (np.linalg.norm(native) * np.linalg.norm(output))),
}
# DS4 uses FP32 patch inputs and its explicit BF16 rounding points, whereas
# the published Python graph uses BF16 module operations. Record that numerical
# comparison as such; it is not a claim of bitwise DS4 oracle equivalence.
assert np.isfinite(native).all() and np.isfinite(output).all()
assert result['cosine'] > .99, result
print(json.dumps(result), flush=True)