"""Held-out output preservation check; does not select or alter quantisation."""
import argparse,hashlib,json,math
from pathlib import Path
import numpy as np
import torch
from PIL import Image,ImageDraw,ImageOps
from semantic_model import load_semantic
from inference import prepare,chroma_coefficients,render
from base93_codec import load_model,decode_state
ROOT=Path(__file__).resolve().parent
p=argparse.ArgumentParser();p.add_argument('--reference',required=True);p.add_argument('--images',default=str(ROOT/'images'));p.add_argument('--archives',default=str(ROOT/'archive_inputs'));args=p.parse_args()
OUT=ROOT/'evaluation';OUT.mkdir(exist_ok=True)
torch.set_num_threads(4)
reference=load_semantic(Path(args.reference)).eval()
quantized=load_model(ROOT/'weights_base93.txt')
source=json.loads((ROOT/'image_manifest.json').read_text())
records=[dict(r,path=Path(args.images)/r['name'],subset='COCO_holdout') for r in source['rows'] if r['split']=='validation']
archives=json.loads((ROOT/'nara_sources.json').read_text())
records += [dict(r,path=Path(args.archives)/r['name'],subset='historical') for r in archives]
results=[];thumbs=[]
@torch.inference_mode()
def run(record,size=256):
    im=Image.open(record['path']).convert('RGB')
    light,alpha,x=prepare(im,size)
    baseline_ab=reference(x);quantized_ab=quantized(x)
    delta=baseline_ab-quantized_ab
    baseline=render(light,alpha,chroma_coefficients(reference,x,8))
    candidate=render(light,alpha,chroma_coefficients(quantized,x,8))
    aa=np.asarray(baseline,dtype=np.float32);bb=np.asarray(candidate,dtype=np.float32)
    d=aa-bb;mse=float(np.mean(d*d));mae=float(np.mean(np.abs(d)))
    row={'name':record['name'],'subset':record['subset'],'size':size,'source_sha256':hashlib.sha256(record['path'].read_bytes()).hexdigest(),'ab_mae':delta.abs().mean().item(),'ab_rmse':delta.square().mean().sqrt().item(),'rgb_mae_255':mae,'rgb_psnr_db':10*math.log10(255**2/max(mse,1e-15)),'rgb_max_difference':float(np.max(np.abs(d)))}
    if size==256:
        baseline.save(OUT/(record['name']+'_fp32.png'));candidate.save(OUT/(record['name']+'_base93.png'))
        neutral=render(light,alpha,(torch.zeros_like(baseline_ab),torch.zeros_like(baseline_ab)))
        thumbs.append((row,[neutral,baseline,candidate]))
    return row
for i,record in enumerate(records):
    row=run(record);results.append(row)
    if (i+1)%10==0:print('EVALUATED',i+1,flush=True)
# Resolution/alpha checks supplement the default-size held-out check.
large=[run(r,512) for r in records[:4]]
rgba=Image.open(records[0]['path']).convert('RGBA');arr=np.asarray(rgba).copy();arr[...,3]=np.linspace(0,255,arr.shape[1],dtype=np.uint8)[None,:];rgba=Image.fromarray(arr)
light,alpha,x=prepare(rgba,256);output=render(light,alpha,chroma_coefficients(quantized,x,8))
assert np.array_equal(np.asarray(output)[...,3],arr[...,3]) and output.size==rgba.size
for page in range(math.ceil(len(thumbs)/8)):
    sheet=Image.new('RGB',(1002,8*264+32),'#202020');draw=ImageDraw.Draw(sheet)
    for j,label in enumerate(['Grayscale input','Production FP32','Base93 decoded']):draw.text((j*334+8,10),label,fill='white')
    for r,(row,ims) in enumerate(thumbs[page*8:page*8+8]):
        y=32+r*264
        draw.text((8,y),f"{row['name']}  RGB MAE {row['rgb_mae_255']:.2f}/255; PSNR {row['rgb_psnr_db']:.1f} dB",fill='white')
        for j,im in enumerate(ims):
            im=ImageOps.contain(im,(330,236));sheet.paste(im,(j*334+(334-im.width)//2,y+23))
    sheet.save(OUT/f'comparison_{page+1:02}.jpg',quality=92)
def aggregate(rows):
    return {'images':len(rows),**{key:float(np.mean([r[key] for r in rows])) for key in ['ab_mae','ab_rmse','rgb_mae_255','rgb_psnr_db']},'worst_rgb_mae':max(rows,key=lambda r:r['rgb_mae_255']),'lowest_psnr':min(rows,key=lambda r:r['rgb_psnr_db'])}
report={'purpose':'Preservation of production FP32 outputs, not colourisation accuracy or ground-truth recovery. No validation input used to select quantisation.','reference_revision':'1a9eb8af2754ad2329a24cfe50d388cb559441d0','base93_sha256':hashlib.sha256((ROOT/'weights_base93.txt').read_bytes()).hexdigest(),'default_settings':{'size':256,'smoothing_radius':8,'saturation':1},'COCO_holdout':aggregate([r for r in results if r['subset']=='COCO_holdout']),'historical':aggregate([r for r in results if r['subset']=='historical']),'size512':aggregate(large),'alpha_preserved':True,'rows':results,'size512_rows':large}
(ROOT/'evaluation.json').write_text(json.dumps(report,indent=2))
print(json.dumps({k:v for k,v in report.items() if k not in ('rows','size512_rows')},indent=2),flush=True)
# Hash every decoded array in little-endian form for the independent JS decoder.
arrays,_,_=decode_state((ROOT/'weights_base93.txt').read_text())
hashes={n:hashlib.sha256(a.astype('<f4' if a.dtype.kind=='f' else '<i8').tobytes()).hexdigest() for n,a in arrays.items()}
(ROOT/'decoded_tensor_hashes.json').write_text(json.dumps(hashes,indent=2))