"""Round 4: matched multiscale structure and DDColor distillation experiments."""
import os,sys,json,time,random,hashlib,traceback,shutil
from pathlib import Path
import numpy as np
import cv2
from PIL import Image,ImageDraw
from skimage import data as sample_data
from skimage.color import rgb2lab,lab2rgb
import pyarrow as pa
import pyarrow.parquet as pq
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset,DataLoader
from huggingface_hub import HfApi,hf_hub_download,PyTorchModelHubMixin
from model import load_model,save_model
from data import decode_image,ColorBins
from metrics import soft_ce,per_image_metrics
from spatial import guided_chroma

OUT=Path('round6');OUT.mkdir(exist_ok=True)
DURABLE=None
START=time.monotonic();DEADLINE=START+3.5*3600
DEVICE='cuda';SEED=409;STEPS=3000
BASE_REV='704fa80d792c3d759db91daa00b2dcfe6f0f6412'
BASE_HASH='0e4c417375684a044860f8af3ac3a2fb44e1a5729254ca33abda758ea71aea6e'
TEACHER_ID='piddnad/ddcolor_artistic'
torch.set_num_threads(6)

def deadline():
    if time.monotonic()>DEADLINE:raise TimeoutError('Graceful export before remote hard timeout')

def gray_arrays(im,style='gray'):
    im=im.resize((256,256),Image.Resampling.BILINEAR)
    rgb=np.asarray(im,np.float32)/255
    lab=rgb2lab(rgb).astype(np.float32)
    if style=='film':
        g=(rgb*np.array([.42,.45,.13],np.float32)).sum(-1)
        g=np.clip((g**1.15-.5)*.8+.5,0,1)
    else:g=np.asarray(im.convert('L'),np.float32)/255
    # Feed the same neutral sRGB input to teacher and student.
    gray=np.repeat(g[...,None],3,-1).astype(np.float32)
    L=rgb2lab(gray)[...,0].astype(np.float32)
    return (L[None]/50-1).copy(),lab[...,1:].transpose(2,0,1).copy(),gray

class Photos(Dataset):
    def __init__(self,table,ids,bins=None,cache=None,augment=False,style='gray'):
        self.table,self.ids,self.bins,self.cache,self.augment,self.style=table,ids,bins,cache,augment,style
    def __len__(self):return len(self.ids)
    def __getitem__(self,k):
        L,ab,gray=gray_arrays(decode_image(self.table,self.ids[k]),self.style)
        teach=np.array(self.cache[k],np.float32) if self.cache is not None else np.zeros((2,64,64),np.float32)
        if self.augment:
            if random.random()<.5:L=L[...,::-1].copy();ab=ab[...,::-1].copy();teach=teach[...,::-1].copy()
            if random.random()<.3:L=np.clip(L*random.uniform(.85,1.15)+random.uniform(-.08,.08),-1,1)
        if self.bins is not None:
            idx,w=self.bins.encode(ab.transpose(1,2,0));return torch.from_numpy(L),torch.from_numpy(ab),torch.from_numpy(idx),torch.from_numpy(w),torch.from_numpy(teach)
        return torch.from_numpy(L),torch.from_numpy(ab),torch.from_numpy(gray.transpose(2,0,1)),self.ids[k]

def patch_excess(p,t):
    """Excess color differences across 4–64px regions, in target-flat areas."""
    values=[]
    for scale in [4,16]:
        a=F.avg_pool2d(p,scale);b=F.avg_pool2d(t,scale)
        for offset in [1,4]:
            for dim in [-1,-2]:
                x=a.narrow(dim,offset,a.shape[dim]-offset)-a.narrow(dim,0,a.shape[dim]-offset)
                y=b.narrow(dim,offset,b.shape[dim]-offset)-b.narrow(dim,0,b.shape[dim]-offset)
                dx=x.norm(dim=1);dy=y.norm(dim=1);mask=(dy<3).float()
                values.append(((dx-dy-1).relu()*mask).sum((1,2))/mask.sum((1,2)).clamp_min(1))
    return sum(values)/len(values)

def structural_loss(p,t):
    terms=[]
    for scale in [4,16]:
        a=F.avg_pool2d(p,scale);b=F.avg_pool2d(t,scale)
        for offset in [1,4]:
            for dim in [-1,-2]:
                x=a.narrow(dim,offset,a.shape[dim]-offset)-a.narrow(dim,0,a.shape[dim]-offset)
                y=b.narrow(dim,offset,b.shape[dim]-offset)-b.narrow(dim,0,b.shape[dim]-offset)
                terms.append(F.smooth_l1_loss(x/10,y/10,beta=.3))
    return sum(terms)/len(terms)

def additional_losses(p,t,teacher):
    small=F.avg_pool2d(p,4);target=F.avg_pool2d(t,4)
    # Teacher is fallible; reduce guidance when it conflicts strongly with ground truth.
    confidence=torch.exp(-(teacher-target).norm(dim=1,keepdim=True)/30).detach()
    kd=(F.smooth_l1_loss(small/20,teacher/20,beta=.5,reduction='none')*confidence).mean()
    neutral=(t.norm(dim=1)<3).float()
    neutral_loss=(p.norm(dim=1)*neutral).sum()/neutral.sum().clamp_min(1)/20
    return structural_loss(p,t),kd,neutral_loss

def measure(L,p,t,raw=None):
    r=per_image_metrics(L,p,t);C=p.norm(dim=1);T=t.norm(dim=1)
    for name,mask,event in [('missed_color',T>12,C<5),('neutral_spill',T<3,C>10)]:
        r[name]=(mask&event).sum((1,2))/mask.sum((1,2)).clamp_min(1)
    r['color_coverage']=(C>10).float().mean((1,2));r['patch_excess']=patch_excess(p,t)
    r['raw_patch_excess']=patch_excess(p if raw is None else raw,t)
    return r

@torch.inference_mode()
def teacher_predict(teacher,gray):
    gray=F.interpolate(gray.to(DEVICE),size=(512,512),mode='bilinear',align_corners=False)
    # Full precision is deliberate: spectral normalization / attention are not assumed BF16-safe.
    out=teacher(gray).float()
    assert out.shape[1]==2 and torch.isfinite(out).all()
    return F.interpolate(out,size=(256,256),mode='bilinear',align_corners=False)

@torch.inference_mode()
def evaluate(model,table,ids,is_teacher=False):
    result={};model.eval()
    for style in ['gray','film']:
        rows=[]
        for L,t,gray,indices in DataLoader(Photos(table,ids,style=style),batch_size=2 if is_teacher else 8,num_workers=4):
            deadline();L,t=L.to(DEVICE),t.to(DEVICE)
            raw=teacher_predict(model,gray) if is_teacher else model.decode(model(L),.38)
            pred=raw if is_teacher else guided_chroma(L,raw,8)
            metrics=measure(L,pred,t,raw)
            for j,index in enumerate(indices):rows.append({'index':int(index)}|{k:float(v[j]) for k,v in metrics.items()})
        keep=[x for x in rows if x['target_chroma']>=5]
        result[style]={'n_total':len(rows),'n_color':len(keep),'summary':{k:float(np.mean([x[k] for x in keep])) for k in keep[0] if k!='index'},'per_image':rows}
    return result

def selection(result,baseline):
    scores=[];eligible=True
    for style in ['gray','film']:
        s=result[style]['summary'];b=baseline[style]['summary']
        eligible &= (s['ab_error']<=b['ab_error']*1.05 and s['neutral_spill']<=b['neutral_spill']+.01
                     and s['missed_color']<=b['missed_color']+.01 and s['color_coverage']>=.95*b['color_coverage']
                     and s['patch_excess']<.85*b['patch_excess'] and s['raw_patch_excess']<.9*b['raw_patch_excess'])
        scores.append(s['patch_excess']/max(b['patch_excess'],1e-6)+.3*s['neutral_spill']+.2*s['ab_error']/b['ab_error'])
    return float(np.mean(scores)),bool(eligible)

@torch.inference_mode()
def probes(model,tag,teacher=False):
    names=['astronaut','coffee','chelsea','rocket','camera','coins','moon'];images=[]
    for name in names:
        original=Image.fromarray(getattr(sample_data,name)()).convert('RGB')
        L,_,gray=gray_arrays(original);x=torch.from_numpy(L)[None].to(DEVICE)
        ab=teacher_predict(model,torch.from_numpy(gray.transpose(2,0,1))[None]) if teacher else guided_chroma(x,model.decode(model(x),.38),8)
        lab=np.concatenate([(L[0]*50+50)[...,None],ab[0].cpu().numpy().transpose(1,2,0)],-1)
        rgb=np.uint8(np.clip(lab2rgb(lab),0,1)*255)
        images.append((name,Image.fromarray(np.uint8(gray*255)),Image.fromarray(rgb)))
    canvas=Image.new('RGB',(512,len(names)*280),'white');d=ImageDraw.Draw(canvas)
    for i,(name,gray,col) in enumerate(images):
        canvas.paste(gray,(0,i*280+24));canvas.paste(col,(256,i*280+24));d.text((4,i*280+4),name+' input | '+tag,fill='black')
    canvas.save(OUT/(tag+'.jpg'),quality=90)

def write(name,obj):
    (OUT/name).write_text(json.dumps(obj,indent=2));return obj