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

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\_v3_verify")

BORDER      = 3
PAD         = 36          # horizontal padding -> recover side (y-axis) labels
VPAD        = 16          # vertical re-pad after row-strip (keep labels adjacent to figure)
D0ACAC      = (208, 172, 172)
MIN_AREA    = 600
HEADER_SKIP = 880
HEADER_H_GATE = 1500
DILATE_K    = 25          # REVERTED to v1
TEXT_H      = 30          # rows whose components are all shorter than this = text
INNER_FRAC  = 0.20
MIN_DENSITY = 0.008

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)

    skip = HEADER_SKIP if h > HEADER_H_GATE else 0      # height-aware (kept)
    if skip:
        binary[:skip, :] = 0

    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

    # --- v1 component keep + grouping (identical kernel/merge) ---
    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

    kernel = np.ones((DILATE_K, DILATE_K), 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
        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])

    merged = []                                          # v1 merge (identical)
    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[:])

    # --- NEW: per-region row-based text strip ---
    # Within each region, mark rows covered by a "tall" component (>= TEXT_H).
    # Trim leading/trailing rows that contain only short components (text lines).
    final = []
    for b in merged:
        bx1, by1, bx2, by2 = b
        sub = binary[by1:by2, bx1:bx2]
        if sub.size == 0:
            continue
        sn, slab, sst, _ = cv2.connectedComponentsWithStats(sub)
        tall_rows = np.zeros(sub.shape[0], dtype=bool)
        for i in range(1, sn):
            if sst[i, cv2.CC_STAT_HEIGHT] >= TEXT_H:
                t = sst[i, cv2.CC_STAT_TOP]; hh = sst[i, cv2.CC_STAT_HEIGHT]
                tall_rows[t:t+hh] = True
        idx = np.where(tall_rows)[0]
        if len(idx) == 0:
            continue                                     # all-text region -> drop
        ty1 = by1 + int(idx[0]);  ty2 = by1 + int(idx[-1]) + 1
        # re-tighten X to illus pixels in the trimmed band; small vertical re-pad
        band = illus[ty1:ty2, bx1:bx2]
        ys, xs = np.where(band > 0)
        if len(xs) == 0:
            continue
        x1 = max(0, bx1 + int(xs.min()) - PAD)
        x2 = min(w,  bx1 + int(xs.max()) + PAD)
        y1 = max(0, ty1 - VPAD)
        y2 = min(h,  ty2 + VPAD)
        if (y2-y1) < 40 or (x2-x1) < 60:
            continue
        mx = int((x2-x1)*INNER_FRAC); my = int((y2-y1)*INNER_FRAC)
        inner = binary[y1+my:y2-my, x1+mx:x2-mx]
        if inner.size == 0 or (inner > 0).mean() < MIN_DENSITY:
            continue
        final.append((x1, y1, x2, y2))

    return pil, final

SAMPLES = ["P2.sh-2023-t2", "P3.acsp-2025-eoy", "P3.ph-2022-wa2", "P3.rg-2022-sa2"]

if __name__ == "__main__":
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    for old in OUTPUT_DIR.glob("*.png"):
        old.unlink()
    for stem in SAMPLES:
        src = INPUT_DIR / f"{stem}.png"
        if not src.exists():
            print(f"{stem}: SOURCE MISSING"); continue
        pil, regions = extract(src)
        for i, (x1, y1, x2, y2) in enumerate(regions, 1):
            crop = pil.crop((x1, y1, x2, y2))
            bw, bh = crop.size
            bd = Image.new("RGB", (bw+BORDER*2, bh+BORDER*2), (255,255,255))
            ImageDraw.Draw(bd).rectangle([0,0,bw+BORDER*2-1,bh+BORDER*2-1], outline=D0ACAC, width=BORDER)
            bd.paste(crop, (BORDER, BORDER))
            bd.save(OUTPUT_DIR / f"{stem}_illus_{i:02d}.png")
        print(f"{stem}: {len(regions)} crop(s)")