import json,statistics,sys
from pathlib import Path
import numpy as np
import torch
import lpips
from PIL import Image,ImageDraw
from safetensors.torch import load_file
from skimage.metrics import structural_similarity,peak_signal_noise_ratio
ROOT=Path(__file__).resolve().parents[1]
CAND=Path(sys.argv[1]); REF=ROOT/'evaluation'/sys.argv[2] if len(sys.argv)>2 else ROOT/'evaluation/fp8'
torch.set_num_threads(6)
out=CAND;out.mkdir(parents=True,exist_ok=True)
metric=lpips.LPIPS(net='alex').cpu().eval()
def composite(im):
    im=im.convert('RGBA');bg=Image.new('RGBA',im.size,(255,255,255,255));return Image.alpha_composite(bg,im).convert('RGB')
def tensor(im):return torch.from_numpy(np.asarray(im).copy()).permute(2,0,1).unsqueeze(0).float()/127.5-1
rows=[]
with torch.inference_mode():
    for key in [p.stem for p in sorted(CAND.glob('*.png'))]:
        ai=Image.open(REF/f'{key}.png');bi=Image.open(CAND/f'{key}.png')
        a=composite(ai).resize((512,512),Image.Resampling.LANCZOS);b=composite(bi).resize((512,512),Image.Resampling.LANCZOS)
        aa=np.asarray(a)/255.;bb=np.asarray(b)/255.
        row={'id':key,'lpips_alex_512':metric(tensor(a),tensor(b)).item(),'ssim_512':structural_similarity(aa,bb,channel_axis=2,data_range=1.),'psnr_512':peak_signal_noise_ratio(aa,bb,data_range=1.)}
        if not key.startswith('edit'):
            la=load_file(str(REF/f'{key}-latents.safetensors'))['latents'].float()
            lb=load_file(str(CAND/f'{key}-latents.safetensors'))['latents'].float()
            row['latent_nrmse']=(((la-lb).square().mean()/la.square().mean()).sqrt()).item()
            row['latent_cosine']=torch.nn.functional.cosine_similarity(la.flatten(),lb.flatten(),dim=0).item()
        if ai.mode=='RGBA' and bi.mode=='RGBA':row['alpha_mae']=float(np.abs(np.asarray(ai.getchannel('A')).astype(float)-np.asarray(bi.getchannel('A')).astype(float)).mean()/255)
        rows.append(row)
        canvas=Image.new('RGB',(1024,544),'white');canvas.paste(a,(0,32));canvas.paste(b,(512,32))
        ImageDraw.Draw(canvas).text((12,10),f'{key} {REF.name.upper()} reference',fill='black');ImageDraw.Draw(canvas).text((524,10),'Optimized FP8',fill='black')
        dst=CAND/('comparisons-'+REF.name);dst.mkdir(exist_ok=True);canvas.save(dst/f'{key}.jpg',quality=94)
    summary={'protocol':'16 disjoint text-to-image prompts plus2 disjoint edits. Matched seeds,40steps,same pipeline and BF16 encoder/VAE. RGB composited over white and resized512 for LPIPS/SSIM/PSNR. Latents compared at full resolution. This is a finite fidelity check, not a broad human preference benchmark.',
             'pairs':rows,'mean_lpips':statistics.mean(r['lpips_alex_512'] for r in rows),'max_lpips':max(r['lpips_alex_512'] for r in rows),
             'mean_ssim':statistics.mean(r['ssim_512'] for r in rows),'min_ssim':min(r['ssim_512'] for r in rows),
             'mean_latent_cosine':statistics.mean(r['latent_cosine'] for r in rows if 'latent_cosine' in r)}
    summary.update(reference=str(REF),candidate=str(CAND),pair_count=len(rows))
    (out/('quality_metrics-'+REF.name+'.json')).write_text(json.dumps(summary,indent=2))
    summary['reference']=str(REF);summary['candidate']=str(CAND);print(json.dumps({k:v for k,v in summary.items() if k!='pairs'},indent=2),flush=True)