"""Compare each exported graph against its pinned source on macOS Core AI.

Exercises encoder masks, recurrent decoder state, and joint logits. This is
conversion parity, not an accuracy benchmark or a confidence calibration.
"""
import argparse
import asyncio
import json
from pathlib import Path
import numpy as np
import torch
from coreai.runtime import AIModel, NDArray, SpecializationOptions, ComputeUnitKind
from export_models import load_source
import apple_parakeet_export as apple

async def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--bundle', type=Path, required=True)
    parser.add_argument('--sources', type=Path, required=True)
    parser.add_argument('--report', type=Path, required=True)
    args = parser.parse_args()
    provenance = json.loads((args.bundle / 'provenance.json').read_text())
    model, processor, _ = load_source(provenance['source'], provenance['revision'], args.sources)
    model = model.to(torch.float16)
    torch.manual_seed(42)
    metadata = json.loads((args.bundle / 'metadata.json').read_text())
    config = model.config
    geometry = metadata['streaming']
    encoder = {'input_features': torch.randn(1, geometry['window_mel_frames'], 128, dtype=torch.float16),
               'attention_mask': torch.ones(1, geometry['window_mel_frames'], dtype=torch.bool)}
    encoder['attention_mask'][:, -201:] = False
    decoder = apple._decoder_step_inputs(config, torch.float16)
    decoder['input_ids'].fill_(42)
    decoder['hidden_state'] = torch.randn_like(decoder['hidden_state']) * 0.1
    decoder['cell_state'] = torch.randn_like(decoder['cell_state']) * 0.1
    joint = apple._joint_inputs(config, torch.float16)
    cases = [('encoder', apple.ParakeetEncoderModule(model), encoder, ['encoder_hidden_states']),
             ('decoder_step', apple.ParakeetDecoderStepModule(model), decoder, ['decoder_output', 'new_hidden_state', 'new_cell_state']),
             ('joint', apple.ParakeetJointModule(model), joint, ['logits'])]
    rows = []
    for name, module, inputs, names in cases:
        with torch.no_grad(), torch.autocast(device_type='cpu', dtype=torch.float16):
            ref = module(**inputs)
        ref = ref if isinstance(ref, tuple) else (ref,)
        asset = await AIModel.load(args.bundle / metadata['assets'][name], SpecializationOptions.from_preferred_compute_unit_kind(ComputeUnitKind.gpu()))
        function = asset.load_function(asset.function_names[0])
        out = await function({key: NDArray(value.detach().numpy()) for key, value in inputs.items()})
        for expected, key in zip(ref, names):
            expected = expected.float().numpy().flatten()
            actual = out[key].numpy().astype(np.float32).flatten()
            assert actual.shape == expected.shape
            assert np.isfinite(actual).all()
            cosine = float(np.dot(actual, expected) / max(1e-12, np.linalg.norm(actual) * np.linalg.norm(expected)))
            relative_rmse = float(np.sqrt(np.mean((actual - expected)**2)) / max(1e-12, np.sqrt(np.mean(expected**2))))
            row = dict(graph=name, output=key, cosine=cosine, relative_rmse=relative_rmse)
            rows.append(row)
            print(json.dumps(row), flush=True)
            assert cosine > 0.995 and relative_rmse < 0.05, row
    args.report.write_text(json.dumps(dict(provenance=provenance, passed=True, graph_parity=rows), indent=2) + '\n')

if __name__ == '__main__':
    torch.set_num_threads(4)
    asyncio.run(main())