#!/usr/bin/env python3
"""Apple Silicon AuroraASR qualification on a pinned acoustic tensor.

Run each backend in an independent fresh process, interleaving benchmark arms.
Produces a JSON record and fails numeric/transcription gates. Prints requested
Core ML compute unit only; ANE placement requires MLComputePlan or Instruments.
"""
import argparse
import json
import os
import platform
import resource
import statistics
import time
from collections import Counter
from pathlib import Path
import numpy as np

p=argparse.ArgumentParser()
p.add_argument('--backend',choices=['mlx','onnx_cpu','coreml'],required=True)
p.add_argument('--model-dir',type=Path,required=True)
p.add_argument('--runs',type=int,default=40)
p.add_argument('--precision',choices=['fp32','fp16'],default='fp32')
p.add_argument('--compute-unit',choices=['all','cpu','gpu'],default='all')
p.add_argument('--dynamic',action='store_true')
a=p.parse_args()
if a.dynamic and a.backend=='coreml':
    raise SystemExit('--dynamic applies to ONNX, not the static Core ML package')
if a.backend!='coreml' and a.precision!='fp32':
    raise SystemExit('The ONNX exports are FP32; --precision fp16 applies only to Core ML')
if platform.system() != 'Darwin' or platform.machine() != 'arm64':
    raise SystemExit('Requires macOS Apple Silicon (arm64).')
fix=np.load(a.model_dir/'reference_fixture.npz')
x=fix['acoustic_input'];expected=fix['reference_logits']
assert x.dtype==np.float32 and expected.ndim==3
mac_profile=None
provider_info={}
start=time.perf_counter()
if a.backend in ('mlx','onnx_cpu'):
    import onnxruntime as ort
    providers=['CPUExecutionProvider']
    if a.backend=='mlx':
        import onnxruntime_ep_mlx
        onnxruntime_ep_mlx.register_execution_provider_library()
        providers=['MLXExecutionProvider','CPUExecutionProvider']
    opts=ort.SessionOptions()
    opts.enable_profiling=a.backend=='mlx'
    fname='aurora_dynamic.onnx' if a.dynamic else 'aurora_asr.onnx'
    sess=ort.InferenceSession(str(a.model_dir/fname),opts,providers=providers)
    if a.backend=='mlx' and 'MLXExecutionProvider' not in sess.get_providers():
        raise RuntimeError('MLX provider not active. Do not call this an MLX benchmark.')
    name=sess.get_inputs()[0].name
    def run():
        return sess.run(None,{name:x})[0]
    provider_info['registered_providers']=sess.get_providers()
else:
    import coremltools as ct
    package=a.model_dir/f'aurora_{a.precision}.mlpackage'
    if not package.is_dir() or not (package/'Manifest.json').exists():
        raise SystemExit(
            f'Core ML package missing: {package.resolve()}\n'
            'Use the sibling AuroraASR-1M-CoreML directory; AuroraASR-1M-MLX '
            'contains ONNX weights only. Do not mix --backend coreml with an MLX --model-dir.')
    unit={'all':ct.ComputeUnit.ALL,'cpu':ct.ComputeUnit.CPU_ONLY,
          'gpu':ct.ComputeUnit.CPU_AND_GPU}[a.compute_unit]
    model=ct.models.MLModel(str(package),compute_units=unit)
    inname=model.get_spec().description.input[0].name
    outname=model.get_spec().description.output[0].name
    def run():
        return np.asarray(model.predict({inname:x})[outname])
    provider_info.update({'coreml_package':package.name,'coreml_requested_compute_unit':a.compute_unit,
                          'actual_ane_placement':'NOT_VERIFIED'})
load_ms=(time.perf_counter()-start)*1000
for _ in range(8):
    out=np.asarray(run())
if out.shape!=expected.shape:
    raise RuntimeError(f'Output shape mismatch: got {out.shape}, expected {expected.shape}')
times=[]
for _ in range(max(1,a.runs)):
    start=time.perf_counter();out=np.asarray(run());times.append((time.perf_counter()-start)*1000)
absdiff=np.abs(out-expected)
max_abs=float(absdiff.max());mean_abs=float(absdiff.mean())
agreement=float(np.mean(out.argmax(axis=-1)==expected.argmax(axis=-1)))
# Meaningful parity gate, avoiding a blanket 0.05 threshold for FP32.
atol=0.05 if (a.backend=='coreml' and a.precision=='fp16') else 0.005
rtol=0.05 if (a.backend=='coreml' and a.precision=='fp16') else 0.005
numerical=bool(np.allclose(out,expected,atol=atol,rtol=rtol) and agreement>=0.995)
try:
    from aurora_runtime import ctc_text
    config=json.loads((a.model_dir/'upstream_reference'/'config.json').read_text())
    spm=a.model_dir/'upstream_reference'/'bpe256.model'
    transcript_ref=ctc_text(expected,spm,config.get('blank_token_id',0))
    transcript_got=ctc_text(out,spm,config.get('blank_token_id',0))
    transcript_exact=bool(transcript_ref==transcript_got)
except Exception as exc:
    transcript_ref=None;transcript_got=None;transcript_exact=False
    provider_info['transcript_error']=f'{type(exc).__name__}: {exc}'
node_providers={};node_durations_us={};mlx_claimed=False
if a.backend=='mlx':
    profile=Path(sess.end_profiling())
    try:
        events=json.loads(profile.read_text())
        counts=Counter(); durations=Counter()
        for e in events:
            if e.get('cat')!='Node':continue
            provider=e.get('args',{}).get('provider')
            if not provider: continue
            counts[str(provider)]+=1
            durations[str(provider)]+=float(e.get('dur',0))
        node_providers=dict(counts)
        node_durations_us=dict(durations)
        mlx_claimed=counts.get('MLXExecutionProvider',0)>0
    finally:
        profile.unlink(missing_ok=True)
# ru_maxrss is bytes on macOS. Process peak is not model-specific allocated memory.
peak_rss_mib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1048576
import hashlib
model_input_path=(a.model_dir/('aurora_dynamic.onnx' if a.dynamic else 'aurora_asr.onnx')) if a.backend!='coreml' else package
if model_input_path.is_file():
    model_sha256=hashlib.sha256(model_input_path.read_bytes()).hexdigest()
else:
    model_sha256=hashlib.sha256(b''.join(
        str(p.relative_to(model_input_path)).encode('utf-8')+b'\x00'+hashlib.sha256(p.read_bytes()).digest()
        for p in sorted(model_input_path.rglob('*')) if p.is_file())).hexdigest()
summary={'backend':a.backend,'precision':a.precision,'dynamic':bool(a.dynamic),
    'model_file':model_input_path.name,'model_sha256':model_sha256,'shape':list(x.shape),
    'cold_load_ms':load_ms,'p50_warm_ms':float(np.median(times)),
    'p95_warm_ms':float(np.percentile(times,95)),
    'max_abs_logit_error':max_abs,'mean_abs_logit_error':mean_abs,
    'frame_argmax_agreement':agreement,'allclose_atol_rtol':[atol,rtol],
    'reference_transcript':transcript_ref,'actual_transcript':transcript_got,
    'transcript_exact':transcript_exact,'peak_process_rss_mib':peak_rss_mib,
    'ep_node_events':node_providers,'ep_node_duration_us':node_durations_us,
    'mlx_claimed_any_subgraph':mlx_claimed if a.backend=='mlx' else None,
    'provider_details':provider_info, 'numeric_pass':numerical}
filename=f"mac_{a.backend}_{a.precision}_{a.compute_unit}_{'dynamic' if a.dynamic else 'static'}_v4_qualification.json"
(a.model_dir/filename).write_text(json.dumps(summary,indent=2))
print(json.dumps(summary,indent=2))
if not numerical:raise SystemExit('FAIL numeric gate')
if not transcript_exact:raise SystemExit('FAIL CTC transcript gate / missing tokenizer')
if a.backend=='mlx' and not mlx_claimed:raise SystemExit('FAIL no MLX subgraph proved')
print('PASS: numeric & decoded CTC parity. ANE device placement remains unverified.')