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

# region tuples are normalized (x1,y1,x2,y2); handwriting/answer regions only.
REGIONS={
  4458:[(0.70,0.04,0.94,0.12),(0.17,0.17,0.90,0.36),(0.88,0.05,0.98,0.22)],
  4459:[(0.18,0.12,0.95,0.30),(0.08,0.28,0.94,0.80),(0.88,0.02,0.98,0.18)],
  4460:[(0.07,0.22,0.94,0.42),(0.88,0.02,0.98,0.18)],
  4461:[(0.09,0.12,0.91,0.47),(0.05,0.48,0.98,0.61),(0.12,0.62,0.82,0.88),(0.80,0.03,0.98,0.16)],
  4447:[(0.70,0.04,0.94,0.12),(0.16,0.16,0.90,0.38),(0.12,0.42,0.94,0.60),(0.08,0.62,0.95,0.96),(0.88,0.05,0.98,0.52)],
  4448:[(0.16,0.06,0.92,0.18),(0.08,0.18,0.94,0.57),(0.16,0.54,0.94,0.69),(0.08,0.68,0.94,0.96),(0.88,0.02,0.98,0.22)],
  4449:[(0.16,0.06,0.92,0.20),(0.08,0.18,0.94,0.50),(0.16,0.54,0.94,0.66),(0.08,0.62,0.94,0.96),(0.88,0.02,0.98,0.55)],
  4450:[(0.08,0.10,0.94,0.50),(0.20,0.52,0.94,0.62),(0.08,0.70,0.94,0.94),(0.04,0.86,0.98,0.99),(0.88,0.02,0.98,0.18)],
}

def build_mask(img, num):
    h,w=img.shape[:2]
    hsv=cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    H,S,V=cv2.split(hsv)
    # Global red teacher corrections and blue/purple pen/signature (avoid yellow/green side bars)
    red=((H<12)|(H>168)) & (S>35) & (V<245)
    blue=((H>85)&(H<150)&(S>30)&(V<245))
    purple=((H>125)&(H<170)&(S>20)&(V<245))
    mask=((red|blue|purple).astype('uint8'))*255

    gray=cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    regmask=np.zeros((h,w),np.uint8)
    for x1,y1,x2,y2 in REGIONS[num]:
        regmask[int(y1*h):int(y2*h), int(x1*w):int(x2*w)] = 255
    # In answer regions remove: colored pen, pencil/graphite grey, and handwriting strokes.
    # Keep very black printed text better by focusing mostly on medium-dark pixels plus saturated strokes.
    medium_dark=((gray>55)&(gray<205)).astype('uint8')*255
    sat_pen=((S>20)&(V<230)).astype('uint8')*255
    region_strokes=cv2.bitwise_and(cv2.bitwise_or(medium_dark, sat_pen), regmask)
    mask=cv2.bitwise_or(mask, region_strokes)

    # Do not mask yellow/green sidebars
    sidebar=((H>18)&(H<90)&(S>35)).astype('uint8')*255
    mask=cv2.bitwise_and(mask, cv2.bitwise_not(sidebar))

    # Remove tiny noise, connect strokes a bit
    k=np.ones((3,3),np.uint8)
    mask=cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((2,2),np.uint8))
    mask=cv2.dilate(mask,k,iterations=1)
    return mask

def clean_one(num, outdir):
    src=PRE/f'IMG_{num}.jpg'
    img=cv2.imread(str(src))
    if img is None: raise FileNotFoundError(src)
    mask=build_mask(img,num)
    # Inpaint masked handwriting with paper/grid neighbourhood.
    clean=cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)
    # Light denoise masked areas by blending with original where no mask
    out_img=outdir/f'IMG_{num}_bereinigt.jpg'
    out_mask=outdir/f'IMG_{num}_mask.jpg'
    cv2.imwrite(str(out_img), clean, [int(cv2.IMWRITE_JPEG_QUALITY), 94])
    cv2.imwrite(str(out_mask), 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]
    pdf=OUT/pdfname
    pil_imgs=[]
    for p in imgs:
        im=Image.open(p).convert('RGB')
        pil_imgs.append(im)
    pil_imgs[0].save(pdf, save_all=True, append_images=pil_imgs[1:], resolution=150.0)
    for im in pil_imgs:
        im.close()
    print(pdf)