"""Pair raw records and generate inspectable image comparisons (no invented quality score)."""
import argparse
import json
import math
import statistics
from pathlib import Path
import numpy as np
from PIL import Image,ImageDraw
from scripts.common import sha256,write_json

PAIR_KEYS=('case_id','prompt','seed','width','height','steps','cfg','vae_tiling','kv_cache','phase_offload','input_sha256','warmup')
REQUIRED_CASES={'portrait','english_text','chinese_text','composition','texture','rgba','edit'}

def records(root):
    root=Path(root)
    if not (root/'COMPLETE.json').exists(): raise ValueError(f'Incomplete evaluation: {root}')
    data=[json.loads(s) for s in (root/'results.jsonl').read_text().splitlines()]
    out={}
    for row in data:
        key=(row['case_id'],row['seed'])
        if key in out: raise ValueError('Duplicate case/seed')
        if sha256(root/row['output'])!=row['output_sha256']: raise ValueError('Image hash mismatch')
        out[key]=row
    return out

def pair_records(baseline,candidate):
    if set(baseline)!=set(candidate): raise ValueError('Different evaluation cases')
    for key,a in baseline.items():
        b=candidate[key]
        for k in PAIR_KEYS:
            if a.get(k)!=b.get(k): raise ValueError(f'Incomparable pair: {key}/{k}')
        yield a,b

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument('--baseline',type=Path,default=Path('artifacts/eval/bf16'))
    ap.add_argument('--candidate',type=Path,default=Path('artifacts/eval/8bit'))
    ap.add_argument('--output',type=Path,default=Path('artifacts/eval'))
    args=ap.parse_args(); args.output.mkdir(parents=True,exist_ok=True)
    pairs=list(pair_records(records(args.baseline),records(args.candidate)))
    rows=[]; sheet=Image.new('RGB',(1024,len(pairs)*550),'#eeeeee'); draw=ImageDraw.Draw(sheet)
    for i,(a,b) in enumerate(pairs):
        images=[Image.open(root/row['output']).convert('RGBA') for root,row in [(args.baseline,a),(args.candidate,b)]]
        av,bv=[np.asarray(img).astype(np.float32) for img in images]
        mse=float(np.mean((av[:,:,:3]-bv[:,:,:3])**2))
        row=dict(case_id=a['case_id'],seed=a['seed'],bf16_seconds=a['seconds'],q8_seconds=b['seconds'],
                 bf16_peak_gib=a['mlx_peak_gib'],q8_peak_gib=b['mlx_peak_gib'],
                 bf16_swap_delta_gib=a.get('swap_delta_gib'),q8_swap_delta_gib=b.get('swap_delta_gib'),
                 rgb_psnr_db=10*math.log10(255**2/max(mse,1e-12)),
                 alpha_mae=float(np.mean(np.abs(av[:,:,3]-bv[:,:,3]))))
        rows.append(row)
        for j,(img,label) in enumerate(zip(images,('BF16','Q8'))):
            checker=Image.new('RGBA',img.size,'white'); cd=ImageDraw.Draw(checker)
            for y in range(0,img.height,32):
                for x in range(0,img.width,32):
                    if (x//32+y//32)%2: cd.rectangle((x,y,x+31,y+31),fill='#d7d7d7')
            thumb=Image.alpha_composite(checker,img).convert('RGB').resize((512,512),Image.Resampling.LANCZOS)
            sheet.paste(thumb,(j*512,i*550+30)); draw.text((j*512+10,i*550+8),f'{a["case_id"]} / {label} / seed {a["seed"]}',fill='black')
    sheet.save(args.output/'comparison.png')
    summary=dict(pairs=len(pairs),cases=len({x['case_id'] for x in rows}),size=pairs[0][0]['width'],steps=pairs[0][0]['steps'],
                 warmup=pairs[0][0]['warmup'],phase_offload=pairs[0][0].get('phase_offload',False),
                 bf16_mean_seconds=statistics.mean(x['bf16_seconds'] for x in rows),
                 q8_mean_seconds=statistics.mean(x['q8_seconds'] for x in rows),
                 bf16_t2i_mean_seconds=statistics.mean(x['bf16_seconds'] for x in rows if x['case_id']!='edit'),
                 q8_t2i_mean_seconds=statistics.mean(x['q8_seconds'] for x in rows if x['case_id']!='edit'),
                 bf16_max_peak_gib=max(x['bf16_peak_gib'] for x in rows),
                 q8_max_peak_gib=max(x['q8_peak_gib'] for x in rows),
                 rows=rows,limitations='One seed per case, one machine, same pinned MLX runtime. Sequential BF16-then-Q8 desktop run without repeated or interleaved trials; order, thermal state and other applications may affect timing. Pixel similarity is not a perceptual quality score; no cross-runtime parity claim.')
    write_json(args.output/'summary.json',summary)
    text=['# Image21-MLX informal evaluation','',summary['limitations'],'',
          '| Case | BF16 s | Q8 s | BF16 peak GiB | Q8 peak GiB | RGB PSNR dB |',
          '|---|---:|---:|---:|---:|---:|']
    for r in rows: text.append(f'| {r["case_id"]} | {r["bf16_seconds"]:.2f} | {r["q8_seconds"]:.2f} | {r["bf16_peak_gib"]:.2f} | {r["q8_peak_gib"]:.2f} | {r["rgb_psnr_db"]:.2f} |')
    text+=['','MLX peak allocated memory excludes the OS, other applications and some process allocations. It is not minimum machine RAM.',
           '', 'Timing includes prompt encoding, phase-by-phase component loading, denoising and VAE decoding; PNG writing is excluded. The after-image RSS field is a snapshot after component release, not a peak.',
           '', 'Existing system swap is recorded separately from per-image swap change. This desktop-session run leaves other applications open. See system-context.json and the per-model environment.json files.',
           '', 'All editing pairs use the same BF16 portrait input. Visual findings and release limitations are in visual-review.json.',
           '', '![All evaluated pairs](comparison.png)','']
    (args.output/'report.md').write_text('\n'.join(text))
    print(json.dumps(summary,indent=2))

if __name__=='__main__': main()