"""Compare converted checkpoints against pinned Mistral CPU FP32 inference."""
import argparse
import gc
import importlib.metadata
import json
import platform
import sys
import subprocess
import time
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "upstream" / "src"))
import numpy as np
import torch
import mlx.core as mx
from timestral import TimestralForecaster as SourceForecaster
from timestral_mlx import TimestralForecaster
from timestral_mlx.loading import sha256

VARIANTS = ("5M", "28M", "123M", "209M")


def fixtures():
    rng = np.random.default_rng(20261020)
    def signal(b, v, t):
        x = (3*np.sin(np.arange(t)*2*np.pi/24) + rng.normal(0,.2,(b,v,t)).cumsum(-1) + 10).astype(np.float32)
        return x
    x = signal(1, 1, 257)
    x[..., 10:16] = np.nan
    yield "univariate_missing_96", x, 96, {}
    x = signal(2, 2, 129); x[0,0,:40] = np.nan; x[1,1,60:70] = np.nan
    yield "multivariate_batched_65", x, 65, {}
    x = signal(1, 1, 193)
    past,known = signal(1,2,193),signal(1,2,193+97)
    past[0,0,20:28] = np.nan
    known[:,1] = np.arange(193+97)%2
    known[0,0,210:215] = np.nan
    yield "numeric_covariates_97", x, 97, dict(past_covariates=past,known_covariates=known)
    codes = (np.arange(193+97)%4 + 2**26).astype(np.float64)[None,None]
    codes[...,5:9] = np.nan
    yield "categorical_covariates_97", x, 97, dict(categorical_past_covariates=codes[...,:193],categorical_known_covariates=codes)
    x = signal(1,1,257)
    yield "autoregressive_1057", x, 1057, dict(symmetric_averaging=False)
    yield "constant_65", np.full(65, 5, np.float32), 65, {}
    yield "all_missing_65", np.full(65, np.nan, np.float32), 65, {}


def validate(size):
    print(f"Validating {size} against upstream...", flush=True)
    src = SourceForecaster.from_pretrained(ROOT / "models" / "source" / size)
    dst = TimestralForecaster.from_pretrained(ROOT / "models" / size)
    records=[]
    for name,x,h,kw in fixtures():
        start=time.perf_counter(); expected=src.forecast(x,h,**kw).quantiles.numpy(); source_s=time.perf_counter()-start
        start=time.perf_counter(); actual=dst.forecast(x,h,**kw).quantiles; mlx_s=time.perf_counter()-start
        delta=actual.astype(np.float64)-expected
        rms=float(np.sqrt(np.mean(expected.astype(np.float64)**2)))
        err=float(np.sqrt(np.mean(delta**2)))
        normalized_rmse=err/max(rms,1e-6)
        normalized_max=float(np.max(np.abs(delta)))/max(rms,1e-6)
        # Fixed, declared acceptance criteria; not adjusted per model.
        passed=bool(np.isfinite(actual).all() and (np.diff(actual,axis=-1)>=0).all()
                    and normalized_rmse <= 1e-4 and normalized_max <= 2e-3)
        record=dict(fixture=name,shape=list(actual.shape),source_seconds=source_s,mlx_seconds=mlx_s,
                    max_absolute_error=float(np.max(np.abs(delta))),rmse=err,reference_rms=rms,
                    normalized_rmse=normalized_rmse,normalized_max_error=normalized_max,
                    finite=bool(np.isfinite(actual).all()),monotone=bool((np.diff(actual,axis=-1)>=0).all()),passed=passed)
        records.append(record)
        print(f"  {name}: normalized RMSE={normalized_rmse:.3g}, max error={record['max_absolute_error']:.3g}, pass={passed}",flush=True)
    report=dict(variant=size,passed=all(r['passed'] for r in records),
                acceptance=dict(normalized_rmse_max=1e-4,normalized_max_error_max=2e-3),
                bundle_manifest_sha256=sha256(ROOT/'models'/size/'manifest.json'),
                runtime_sha256={p.name:sha256(p) for p in sorted((ROOT/'timestral_mlx').glob('*.py')) if not p.name.startswith('._')},
                environment=dict(python=platform.python_version(),platform=platform.platform(),
                                 torch=torch.__version__,mlx=importlib.metadata.version('mlx'),numpy=np.__version__,
                                 source_device='cpu',mlx_device=str(mx.default_device()),torch_threads=torch.get_num_threads()),
                notes='Deterministic synthetic parity fixtures; timings include first-call overhead and are not a controlled benchmark. No GIFT-Eval/TIME/fev-bench evaluation was run.',fixtures=records)
    (ROOT/'reports'/f'{size}.json').write_text(json.dumps(report,indent=2)+'\n')
    del src,dst; gc.collect(); mx.clear_cache()
    if not report['passed']:
        raise RuntimeError(f'{size} did not pass numerical validation')


if __name__=='__main__':
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('variants',nargs='*',choices=VARIANTS,default=list(VARIANTS))
    revision=subprocess.check_output(['git','-C',str(ROOT/'upstream'),'rev-parse','HEAD'],text=True).strip()
    if revision!='a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff':
        raise ValueError('upstream source revision mismatch')
    subprocess.run(['git','-C',str(ROOT/'upstream'),'diff','--exit-code','HEAD','--','src'],check=True)
    torch.set_num_threads(1)
    for size in p.parse_args().variants:
        validate(size)