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

INPUT_DIR  = Path(r"C:\allgifted\mathapi11v2\docs\past year questions\Old Images")
OUTPUT_DIR = Path(r"C:\allgifted\mathapi11v2\docs\past year questions\Only Images")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

BORDER    = 3
PAD       = 22          # enough to keep oval/circle borders intact
D0ACAC    = (208, 172, 172)
MIN_AREA  = 600         # filters body text; keeps illustration shapes
HEADER_SKIP = 880       # skips school logo/name/date/section header

def extract(image_path):
    pil = Image.open(image_path).convert("RGB")
    img = np.array(pil)
    h, w = img.shape[:2]

    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    _, binary = cv2.threshold(gray, 210, 255, cv2.THRESH_BINARY_INV)

    # Blank school header
    binary[:HEADER_SKIP, :] = 0

    # Blank yellow section-header bands
    R, G, B = img[:,:,0], img[:,:,1], img[:,:,2]
    yellow = (R.astype(int)>200) & (G.astype(int)>200) & (B.astype(int)<180)
    for y in np.where(yellow.mean(axis=1) > 0.2)[0]:
        binary[max(0,y-2):min(h,y+28), :] = 0

    # Keep only large components (removes text characters)
    n_lbl, labels, stats, _ = cv2.connectedComponentsWithStats(binary)
    illus = np.zeros_like(binary)
    for i in range(1, n_lbl):
        if stats[i, cv2.CC_STAT_AREA] >= MIN_AREA:
            illus[labels == i] = 255

    # Dilate to group nearby illustration elements
    kernel = np.ones((25, 25), np.uint8)
    dilated = cv2.dilate(illus, kernel)

    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    regions = []
    for c in contours:
        x, y, bw, bh = cv2.boundingRect(c)
        if bh < 40 or bw < 80 or bw * bh < 8000:
            continue
        # Skip thin horizontal rules/separators
        if bh < 45 and bw > 400:
            continue
        roi = illus[y:y+bh, x:x+bw]
        ys, xs = np.where(roi > 0)
        if len(xs) == 0:
            continue
        x1 = max(0, x + int(xs.min()) - PAD)
        y1 = max(0, y + int(ys.min()) - PAD)
        x2 = min(w,  x + int(xs.max()) + PAD)
        y2 = min(h,  y + int(ys.max()) + PAD)
        regions.append([x1, y1, x2, y2])

    regions.sort(key=lambda b: b[1])

    # Merge vertically close regions
    merged = []
    for b in regions:
        if merged and b[1] <= merged[-1][3] + 15:
            merged[-1][0] = min(merged[-1][0], b[0])
            merged[-1][1] = min(merged[-1][1], b[1])
            merged[-1][2] = max(merged[-1][2], b[2])
            merged[-1][3] = max(merged[-1][3], b[3])
        else:
            merged.append(b[:])

    # Re-tighten each box to actual illustration pixels + PAD
    final = []
    for b in merged:
        roi = illus[b[1]:b[3], b[0]:b[2]]
        ys, xs = np.where(roi > 0)
        if len(xs) == 0:
            continue
        x1 = max(0, b[0] + int(xs.min()) - PAD)
        y1 = max(0, b[1] + int(ys.min()) - PAD)
        x2 = min(w,  b[0] + int(xs.max()) + PAD)
        y2 = min(h,  b[1] + int(ys.max()) + PAD)
        if (y2-y1) >= 40 and (x2-x1) >= 60:
            final.append((x1, y1, x2, y2))

    return pil, final

# -- Process every image in Old Images --
exts = {".png", ".jpg", ".jpeg", ".bmp"}
files = [f for f in INPUT_DIR.iterdir() if f.suffix.lower() in exts]
print(f"Found {len(files)} paper images.")

log = []
for img_path in sorted(files):
    print(f"\nProcessing: {img_path.name}")
    try:
        pil, regions = extract(img_path)
    except Exception as e:
        print(f"  ERROR: {e}")
        log.append(f"{img_path.name} | ERROR | {e}")
        continue

    stem = img_path.stem
    saved = 0
    for i, (x1, y1, x2, y2) in enumerate(regions, 1):
        crop = pil.crop((x1, y1, x2, y2))
        # Add D0ACAC border
        bw, bh = crop.size
        bordered = Image.new("RGB", (bw + BORDER*2, bh + BORDER*2), (255,255,255))
        draw = ImageDraw.Draw(bordered)
        draw.rectangle([0, 0, bw+BORDER*2-1, bh+BORDER*2-1],
                       outline=D0ACAC, width=BORDER)
        bordered.paste(crop, (BORDER, BORDER))

        out_name = f"{stem}_illus_{i:02d}.png"
        bordered.save(OUTPUT_DIR / out_name)
        saved += 1
        print(f"  Saved: {out_name} ({bw}x{bh}px)")

    log.append(f"{img_path.name} | {len(regions)} illustrations | {saved} saved")

# -- Summary log --
log_path = OUTPUT_DIR / "_extraction_log.txt"
log_path.write_text("\n".join(log))
print(f"\nDone. Log: {log_path}")