#!/usr/bin/env python3
from pathlib import Path
from PIL import Image
import numpy as np

KEYS=['seated-box-jump','pull-up','lateral-line-hops','lateral-power-shuffle','kneeling-soccer-toss','buddy-hamstring-curls']
root=Path.cwd()
src=root/'incoming/greenscreen-source'
dst=root/'incoming/approved-figures'
qa=root/'qa_screenshots'
dst.mkdir(parents=True,exist_ok=True); qa.mkdir(parents=True,exist_ok=True)

def contain(img, size=(1600,900)):
    w,h=size
    ratio=min(w/img.width,h/img.height)
    nw,nh=max(1,int(img.width*ratio)),max(1,int(img.height*ratio))
    im=img.resize((nw,nh),Image.LANCZOS)
    out=Image.new('RGBA',size,(0,0,0,0))
    out.alpha_composite(im,((w-nw)//2,(h-nh)//2))
    return out

def key_green(im):
    arr=np.array(im.convert('RGBA'),dtype=np.float32)
    r,g,b=arr[...,0],arr[...,1],arr[...,2]
    # Transparent when green dominates strongly and is bright.
    green_dom=g-np.maximum(r,b)
    bg=(g>120)&(green_dom>35)
    # Soft alpha transition around the edge
    alpha=np.where(bg,0.0,255.0)
    edge=(g>95)&(green_dom>15)&(~bg)
    alpha[edge]=np.clip((35-green_dom[edge])/20*255,0,255)
    # Keep blue arrows and navy clothes
    keep_blue=(b>80)&(g<230)
    alpha[keep_blue]=255
    # Keep skin and black equipment/shadows except pure green background
    keep_skin=(r>90)&(g<210)&(b<190)
    alpha[keep_skin]=255
    # Despill: reduce green in semi/nontransparent pixels where green over-dominates
    visible=alpha>0
    arr[...,1]=np.where(visible, np.minimum(g, np.maximum(r,b)+25), g)
    arr[...,3]=alpha
    return Image.fromarray(np.clip(arr,0,255).astype('uint8'),'RGBA')

def comp(img,bg):
    base=Image.new('RGBA',img.size,bg+(255,))
    base.alpha_composite(img)
    return base.convert('RGB')

for k in KEYS:
    p=src/f'{k}.png'
    im=Image.open(p)
    keyed=contain(key_green(im))
    out=dst/f'{k}.webp'
    keyed.save(out,'WEBP',quality=96,method=0)
    comp(keyed,(0,0,0)).save(qa/f'{k}_qa_black.png')
    comp(keyed,(30,102,255)).save(qa/f'{k}_qa_blue.png')
    comp(keyed,(245,245,245)).save(qa/f'{k}_qa_light.png')
    print('FAST KEYED',k,out)