from pathlib import Path
import cv2, numpy as np
from PIL import Image

BASE=Path('/home/agent/family/Lernen/Pruefungen_Emilia_4_Klasse/Mathe')
PRE=Path('/home/agent/tmp/emilia_pruefungen_import/preprocessed_jpg')
OUT=BASE/'_Uebungsblaetter_bereinigt_final'
OUT.mkdir(parents=True, exist_ok=True)

# Manual whiteout rectangles, normalized. grid=True redraws light graph grid.
ACTIONS={
4458:[('color',),('white',(0.68,0.04,0.96,0.12),False),('white',(0.12,0.17,0.90,0.39),True),('white',(0.88,0.04,0.98,0.18),False)],
4459:[('color',),('white',(0.22,0.11,0.96,0.28),False),('white',(0.12,0.29,0.94,0.77),True),('white',(0.88,0.02,0.98,0.18),False)],
4460:[('color',),('white',(0.10,0.24,0.94,0.62),True),('white',(0.88,0.02,0.98,0.18),False)],
4461:[('color',),('white',(0.10,0.13,0.94,0.56),True),('white',(0.04,0.48,0.98,0.62),False),('white',(0.10,0.62,0.86,0.90),False),('white',(0.78,0.02,0.98,0.16),False)],
4447:[('color',),('white',(0.68,0.04,0.96,0.12),False),('white',(0.15,0.15,0.93,0.41),True),('white',(0.12,0.42,0.96,0.61),False),('white',(0.08,0.60,0.97,0.99),True),('white',(0.88,0.04,0.98,0.56),False)],
4448:[('color',),('white',(0.12,0.06,0.98,0.22),False),('white',(0.08,0.18,0.96,0.58),True),('white',(0.12,0.55,0.96,0.72),False),('white',(0.08,0.66,0.96,0.99),True),('white',(0.88,0.02,0.98,0.25),False)],
4449:[('color',),('white',(0.12,0.06,0.98,0.25),False),('white',(0.08,0.20,0.96,0.53),True),('white',(0.12,0.53,0.96,0.68),False),('white',(0.08,0.64,0.96,0.99),True),('white',(0.88,0.02,0.98,0.58),False)],
4450:[('color',),('white',(0.08,0.10,0.96,0.54),True),('white',(0.14,0.50,0.96,0.65),False),('white',(0.08,0.67,0.96,0.95),True),('white',(0.04,0.84,0.98,0.99),False),('white',(0.88,0.02,0.98,0.18),False)],
}

def remove_colored(img):
    hsv=cv2.cvtColor(img, cv2.COLOR_BGR2HSV); H,S,V=cv2.split(hsv)
    red=((H<14)|(H>166)) & (S>24) & (V<250)
    blue=((H>82)&(H<150)&(S>24)&(V<245))
    purple=((H>125)&(H<170)&(S>20)&(V<245))
    sidebar=((H>18)&(H<90)&(S>45)&(V>70)).astype('uint8')*255
    mask=((red|blue|purple).astype('uint8'))*255
    mask=cv2.bitwise_and(mask, cv2.bitwise_not(sidebar))
    mask=cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((2,2),np.uint8))
    mask=cv2.dilate(mask,np.ones((3,3),np.uint8),iterations=1)
    return cv2.inpaint(img,mask,2,cv2.INPAINT_TELEA)

def paper_color(img):
    # sample lower right-ish clean paper
    h,w=img.shape[:2]
    crop=img[int(.72*h):int(.9*h), int(.60*w):int(.82*w)]
    med=np.median(crop.reshape(-1,3),axis=0)
    med=np.clip(med+8,210,248).astype(np.uint8)
    return tuple(int(x) for x in med.tolist())

def apply_white(img, rect, grid=False):
    h,w=img.shape[:2]
    x1,y1,x2,y2=rect
    x1=int(x1*w); x2=int(x2*w); y1=int(y1*h); y2=int(y2*h)
    col=paper_color(img)
    cv2.rectangle(img,(x1,y1),(x2,y2),col,-1)
    if grid:
        line=(175,182,178)  # BGR light grid
        step=max(18, int(w*0.018))
        for x in range(x1, x2, step): cv2.line(img,(x,y1),(x,y2),line,1)
        for y in range(y1, y2, step): cv2.line(img,(x1,y),(x2,y),line,1)
    return img

def clean(num,outdir):
    img=cv2.imread(str(PRE/f'IMG_{num}.jpg'))
    for act in ACTIONS[num]:
        if act[0]=='color': img=remove_colored(img)
        elif act[0]=='white': img=apply_white(img,act[1],act[2])
    out=outdir/f'IMG_{num}_uebung.jpg'
    cv2.imwrite(str(out),img,[int(cv2.IMWRITE_JPEG_QUALITY),95])
    return out

sets={
 'Mathe_Emilia_4Kl_LZK_4_12_Textaufgaben_Schriftliche_Multiplikation_Uebungsblatt.pdf':[4458,4459,4460,4461],
 'Mathe_Emilia_4Kl_LZK_4_13_Schriftliche_Division_Monsterzahl_Uebungsblatt.pdf':[4447,4448,4449,4450],
}
for pdfname,nums in sets.items():
    sub=OUT/pdfname.replace('.pdf',''); sub.mkdir(exist_ok=True)
    imgs=[clean(n,sub) for n in nums]
    pil=[Image.open(p).convert('RGB') for p in imgs]
    pdf=OUT/pdfname
    pil[0].save(pdf,save_all=True,append_images=pil[1:],resolution=150.0)
    for im in pil: im.close()
    print(pdf)