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_v2'
OUT.mkdir(parents=True, exist_ok=True)

BLUE_PAGES={4458,4459,4460,4461}
PENCIL_PAGES={4447,4448,4449,4450}
# normalized answer regions for pencil pages only, kept narrower to avoid printed prompts
PENCIL_REGIONS={
 4447:[(0.30,0.18,0.88,0.39),(0.28,0.45,0.94,0.57),(0.10,0.62,0.94,0.96),(0.70,0.04,0.94,0.12),(0.88,0.05,0.98,0.55)],
 4448:[(0.28,0.08,0.92,0.18),(0.12,0.22,0.94,0.56),(0.30,0.60,0.94,0.67),(0.10,0.70,0.94,0.96),(0.88,0.02,0.98,0.22)],
 4449:[(0.28,0.08,0.92,0.19),(0.10,0.24,0.94,0.50),(0.28,0.56,0.94,0.64),(0.10,0.67,0.94,0.96),(0.88,0.02,0.98,0.55)],
 4450:[(0.18,0.13,0.94,0.50),(0.22,0.52,0.94,0.62),(0.10,0.70,0.94,0.94),(0.04,0.86,0.98,0.99),(0.88,0.02,0.98,0.18)],
}
# mask answer/name/signature regions for blue pages too, but only colored strokes there/global
BLUE_REGIONS={
 4458:[(0.70,0.04,0.94,0.12),(0.12,0.14,0.92,0.38),(0.88,0.05,0.98,0.22)],
 4459:[(0.12,0.10,0.96,0.80),(0.88,0.02,0.98,0.18)],
 4460:[(0.08,0.20,0.95,0.42),(0.88,0.02,0.98,0.18)],
 4461:[(0.08,0.10,0.94,0.50),(0.04,0.48,0.98,0.62),(0.12,0.62,0.82,0.88),(0.80,0.03,0.98,0.16)],
}

def color_mask(img):
    hsv=cv2.cvtColor(img, cv2.COLOR_BGR2HSV); H,S,V=cv2.split(hsv)
    red=((H<12)|(H>168)) & (S>28) & (V<250)
    blue=((H>88)&(H<145)&(S>35)&(V<235))
    purple=((H>130)&(H<170)&(S>28)&(V<235))
    # remove colored marks but not yellow/green side bars
    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))
    return mask

def region_mask(shape, regs):
    h,w=shape[:2]
    m=np.zeros((h,w),np.uint8)
    for x1,y1,x2,y2 in regs:
        m[int(y1*h):int(y2*h), int(x1*w):int(x2*w)] = 255
    return m

def clean_one(num,outdir):
    img=cv2.imread(str(PRE/f'IMG_{num}.jpg'))
    h,w=img.shape[:2]
    mask=color_mask(img)
    hsv=cv2.cvtColor(img, cv2.COLOR_BGR2HSV); H,S,V=cv2.split(hsv)
    gray=cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    if num in PENCIL_PAGES:
        rm=region_mask(img.shape, PENCIL_REGIONS[num])
        # pencil/graphite: medium/dark low-saturation strokes. threshold conservative to spare printed/grid text.
        pencil=((gray<150)&(V<205)).astype('uint8')*255
        # avoid the yellow/green sidebars
        sidebar=((H>18)&(H<90)&(S>45)&(V>70)).astype('uint8')*255
        pencil=cv2.bitwise_and(pencil, cv2.bitwise_not(sidebar))
        pencil=cv2.bitwise_and(pencil, rm)
        mask=cv2.bitwise_or(mask,pencil)
    elif num in BLUE_PAGES:
        # keep only colored marks; include in defined regions and globally red/blue corrections.
        pass
    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)
    clean=cv2.inpaint(img, mask, 2, cv2.INPAINT_TELEA)
    out_img=outdir/f'IMG_{num}_bereinigt.jpg'
    cv2.imwrite(str(out_img), clean, [int(cv2.IMWRITE_JPEG_QUALITY), 94])
    cv2.imwrite(str(outdir/f'IMG_{num}_mask.jpg'), mask)
    return out_img

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_one(n,sub) for n in nums]
    pil=[]
    for p in imgs: pil.append(Image.open(p).convert('RGB'))
    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)