"""How high should the SCSI profile go?

Two questions the current profiles never answered:
  1. With the payload deflated, where does the hybrid's lam->0 end actually land?
     (Un-deflated it is 439 KB/s, which is NOT comparable to the 247 KB/s
      lossless changed-spans+deflate path.)
  2. Is there any point shipping lossy VQ on SCSI at all, or does the lossless
     path dominate once both are entropy-coded?
"""
import sys, zlib; sys.path.insert(0,'tools/encoder')
import vq_hybrid as H, vq as VQ, numpy as np, struct
S=sys.argv[1]

def pack_modes(mode):
    n=len(mode); out=bytearray((n*2+7)//8)
    for i,m in enumerate(mode): out[i//4] |= (int(m)&3)<<(6-2*(i%4))
    return bytes(out)

def payload(mode,l1,l4g,src,nbx):
    b=bytearray()
    for i,mo in enumerate(mode):
        if mo==1: b.append(int(l1[i]))
        elif mo==2:
            for j in range(4): b.append(int(l4g[i][j]))
        elif mo==3:
            by,bx=divmod(i,nbx); b+=src[by*4:by*4+4,bx*4:bx*4+4].tobytes()
    return bytes(b)

def spans(a,b,gap=8):
    out=bytearray()
    for y in range(a.shape[0]):
        d=np.nonzero(a[y]!=b[y])[0]
        if not len(d): continue
        sp=[];st=d[0];p=d[0]
        for x in d[1:]:
            if x-p>gap: sp.append((st,p)); st=x
            p=x
        sp.append((st,p))
        for s0,e0 in sp: out+=bytes([y,s0,e0-s0])+b[y][s0:e0+1].tobytes()
    return bytes(out)

for s in ['00020','00146']:
    m=H.build(f'{S}/fr_{s}',k1=256,k4=256,iters=16)
    pal,idx,W=m['pal'],m['idx'],m['W']; nbx=W//4
    ceil_db=np.mean([VQ.psnr(o,pal[i]) for o,i in zip(m['rgb'],idx)])
    # lossless reference, same frames
    lz=np.mean([len(zlib.compress(spans(idx[i-1],idx[i]),9)) for i in range(1,len(idx))])
    print(f'=== {s}  palette ceiling {ceil_db:.2f} dB ===',flush=True)
    print(f'  LOSSLESS changed-spans+deflate : {lz*12/1024:6.1f} KB/s  (pixel-exact)',flush=True)
    print(f'  {"lam":>6}{"PSNR":>8}{"raw KB/s":>10}{"defl KB/s":>11}{"gain":>7}{"RAW%":>7}',flush=True)
    for lam in [0.,10.,25.,60.,150.,300.,800.]:
        e=H.encode(m,lam=lam); r=H.evaluate(m,e)
        tot_r=tot_d=0
        for f,im in enumerate(idx):
            B1=H.blocks_of(im,pal,4,4); l1=VQ.assign(B1,m['C1s'])
            B4=H.blocks_of(im,pal,2,2); l4=VQ.assign(B4,m['C4s'])
            q=H._group_2x2_into_4x4(np.arange(len(l4)),W); l4g=l4[q].reshape(-1,4)
            blob=pack_modes(e['modes'][f])+payload(e['modes'][f],l1,l4g,im,nbx)
            tot_r+=len(blob); tot_d+=len(zlib.compress(blob,9))
        n=len(idx)
        rk=tot_r/n*12/1024; dk=tot_d/n*12/1024
        print(f'  {lam:6.0f}{r["psnr"]:8.2f}{rk:10.1f}{dk:11.1f}{rk/dk:7.2f}{r["raw"]:7.1f}',flush=True)