"""Compare cached ONNX inference against the official PyTorch trajectories."""
import json
from pathlib import Path
import time
import numpy as np
import onnxruntime as ort

out = Path('validation')
fixtures = json.loads((out / 'fixtures.json').read_text())
bf16 = np.load(out / 'reference-bf16.npz')
fp16 = np.load(out / 'reference-fp16.npz')


def metrics(reference, actual):
    r = reference.astype(np.float64); a = actual.astype(np.float64)
    p = np.exp(r - r.max(-1, keepdims=True)); p /= p.sum(-1, keepdims=True)
    q = np.exp(a - a.max(-1, keepdims=True)); q /= q.sum(-1, keepdims=True)
    return {'rmse': float(np.sqrt(np.mean((a-r)**2))),
            'max_abs': float(np.max(np.abs(a-r))),
            'cosine': float(np.mean(np.sum(a*r,-1)/(np.linalg.norm(a,axis=-1)*np.linalg.norm(r,axis=-1)))),
            'argmax_agreement': float(np.mean(np.argmax(a,-1)==np.argmax(r,-1))),
            'mean_kl_reference_to_actual': float(np.mean(np.sum(p*np.log((p+1e-30)/(q+1e-30)),axis=-1)))}


reports = {}
for name, path in [('fp16', 'models/own-fp16/model.onnx'),
                   ('int4', 'models/own-int4/model.onnx'),
                   ('packaged', 'models/minicpm5-webgpu/onnx/model_q4f16.onnx')]:
    options = ort.SessionOptions(); options.intra_op_num_threads = 8; options.inter_op_num_threads = 1
    start = time.time()
    session = ort.InferenceSession(path, sess_options=options, providers=['CPUExecutionProvider'])
    print('Loaded', name, time.time()-start, flush=True)
    output_names = [o.name for o in session.get_outputs()]
    samples = {}
    for fixture in fixtures:
        cache = {i.name: np.zeros((1,2,0,128), dtype=np.float16)
                 for i in session.get_inputs() if i.name.startswith('past_key_values')}
        ids = np.array([fixture['input_ids']], dtype=np.int64)
        total = ids.shape[1]
        logits = []
        for next_id in fixture['continuation']:
            feeds = {'input_ids': ids, 'attention_mask': np.ones((1,total),dtype=np.int64), **cache}
            results = dict(zip(output_names, session.run(None, feeds)))
            assert results['logits'].shape == (1,1,130560)
            assert np.isfinite(results['logits']).all()
            logits.append(results['logits'][0,-1].astype(np.float32))
            cache = {key.replace('present', 'past_key_values'): value
                     for key,value in results.items() if key.startswith('present')}
            ids = np.array([[next_id]],dtype=np.int64); total += 1
        samples[fixture['id']] = np.stack(logits)
        print(name, fixture['id'], 'checked', len(logits), 'steps', flush=True)
    np.savez(out / f'onnx-{name}.npz', **samples)
    actual = np.concatenate(list(samples.values()))
    reports[name] = {'vs_bf16': metrics(np.concatenate([bf16[f['id']] for f in fixtures]),actual),
                     'vs_fp16': metrics(np.concatenate([fp16[f['id']] for f in fixtures]),actual),
                     'positions': len(actual), 'elapsed_s': time.time()-start}
    if name == 'packaged':
        raw = np.load(out/'onnx-int4.npz')
        reports[name]['vs_unmodified_export'] = metrics(np.concatenate([raw[f['id']] for f in fixtures]), actual)
        assert all(np.array_equal(samples[f['id']],raw[f['id']]) for f in fixtures), 'Packaging changed logits'
    print(json.dumps(reports[name]), flush=True)
    (out/'onnx-validation.json').write_text(json.dumps(reports,indent=2))
    del session, cache, results
assert reports['fp16']['vs_fp16']['cosine'] > 0.999, 'Unquantized export deviates from PyTorch'
assert reports['fp16']['vs_fp16']['mean_kl_reference_to_actual'] < 0.01, 'Unquantized export KL too high'
print('Export parity and lossless packaging checks passed', flush=True)