"""Run real image-conditioned decisions with the released JEVision checkpoint."""
import argparse
import json
from pathlib import Path
import statistics
import time
import torch
from PIL import Image
from training.judge.model import ImageJudge
from training.judge.fast_engine import FastEngine, convert_projections

p = argparse.ArgumentParser()
p.add_argument('--root', default='./data', help='Base model cache root')
p.add_argument('--checkpoint', default='.')
p.add_argument('--image', help='Optional PNG, JPEG or WebP image')
p.add_argument('--question', default='What is the dominant color in the image?')
p.add_argument('--choices', default='red,blue,green')
p.add_argument('--output', default='verification.json')
a = p.parse_args()
torch.set_num_threads(4)
started = time.perf_counter()
model = ImageJudge(Path(a.root), checkpoint=Path(a.checkpoint)).eval()
engine = FastEngine(model, resident=False, row_buckets=(1, 2, 3, 4), max_call_tokens=2560)
convert_projections(engine, 'w8a8', stage_on_host=True)
engine.release_original(cpu_embeddings=True)
engine.load_costs(Path(__file__).parent / 'training/judge/engine_costs_rtx5060ti_w8a8.json')
engine.prepare(graphs=False)
engine.enable_static_vision()
setup = time.perf_counter() - started
q = {'instructions': a.question, 'criteria': {x: x for x in a.choices.split(',')}}
images = [(Path(a.image).name, Image.open(a.image).convert('RGB'))] if a.image else [(color, Image.new('RGB', (224, 224), color)) for color in ('red','blue','green')]
rows = []
for label, im in images:
    engine.predict_many('Use the supplied image.', [q], image=im)
    for repeat in range(3):
        torch.cuda.synchronize(); begin = time.perf_counter()
        r = engine.predict_many('Use the supplied image.', [q], image=im)
        torch.cuda.synchronize(); elapsed = (time.perf_counter()-begin)*1000
        probs = r['results'][0]['probabilities']
        rows.append({'image': label, 'prediction': max(probs, key=probs.get), 'probabilities': probs, 'wall_ms': elapsed})
report = {'gpu': torch.cuda.get_device_name(), 'setup_seconds': setup, 'cases': rows, 'median_ms': statistics.median(x['wall_ms'] for x in rows), 'scope': 'Synthetic color image smoke test; warm inference excludes startup, network and queue time.'}
Path(a.output).write_text(json.dumps(report, indent=2))
print(json.dumps(report, indent=2))
if not a.image:
    assert all(x['image'] == x['prediction'] for x in rows), 'Image-conditioned smoke test failed'