#!/usr/bin/env python3
"""
Shared, shape-agnostic primitives used by both classical_cv_detector.py and
the per-shape modules under shapes/: PDF loading, binarization, tiling,
OCR helpers, and deduplication. No shape-specific (diamond/circle/hexagon)
logic lives here.
"""

import os
import re
import subprocess

import cv2
import easyocr
import fitz  # PyMuPDF
import numpy as np

# Resolve bundled model directory relative to this file so the path works both
# locally and on Streamlit Cloud (where the working directory may differ).
_MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "Model", "easyocr")

# GPU acceleration for EasyOCR's neural OCR calls, which dominate per-tile
# detection time (every candidate that clears the geometry gates pays for at
# least one OCR call, often two for text-first re-verification). EasyOCR
# auto-selects cuda -> mps -> cpu when gpu=True (checked directly against the
# installed easyocr package: see easyocr.py's Reader.__init__), so this also
# picks up Apple-Silicon MPS acceleration with no code change needed on Macs,
# not just CUDA machines. Was hard-coded gpu=False, leaving free hardware
# acceleration unused. Escape hatch (CV_FORCE_CPU=1) kept for machines where
# GPU/MPS OCR misbehaves, and for direct CPU-vs-GPU benchmarking.
_GPU_ENABLED = os.environ.get("CV_FORCE_CPU") != "1"

# Initialise EasyOCR once at import time. model_storage_directory points at the
# bundled weights so no network download is needed on any deployment.
_READER    = easyocr.Reader(["en"], gpu=_GPU_ENABLED, verbose=False,
                             model_storage_directory=_MODEL_DIR)
_ALLOWLIST = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"


# ── Tunable parameters (shared across shapes) ──────────────────────────────────
DPI          = 200    # PDF render DPI.  Higher → bigger symbols, better OCR, slower.
TILE_SIZE    = 640    # Sliding-window tile size in pixels.
TILE_OVERLAP = 0.15   # Fraction of tile_size used as overlap between adjacent tiles.
MIN_AREA     = 100    # px²  — drops sub-pixel noise
MAX_AREA     = 8000   # px²  — drops room outlines and wall regions
OCR_UPSCALE  = 6      # resize factor before OCR
MASK_ERODE   = 2      # erosion iterations to strip a shape's outline from its interior mask
NMS_DIST     = 25     # px — two detections closer than this are considered duplicates
HV_LINE_MIN  = 80     # px — H/V lines longer than this are erased in the line-removal pass
# ──────────────────────────────────────────────────────────────────────────────


# ── PDF / image loading ────────────────────────────────────────────────────────

def load_pages(pdf_path: str, dpi: int = DPI) -> list[np.ndarray]:
    """Return each PDF page as a BGR numpy array."""
    doc = fitz.open(pdf_path)
    mat = fitz.Matrix(dpi / 72, dpi / 72)
    pages = []
    for page in doc:
        pix = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB)
        arr = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, 3)
        pages.append(arr[:, :, ::-1].copy())  # RGB → BGR
    return pages


# ── Preprocessing ──────────────────────────────────────────────────────────────

def binarize(bgr: np.ndarray) -> np.ndarray:
    """
    Adaptive threshold → inverted binary (white lines on black background).
    Working on a tile rather than the full page makes blockSize=15 meaningful
    relative to the symbol size.
    """
    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (3, 3), 0)
    return cv2.adaptiveThreshold(
        blur, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        blockSize=15,
        C=6,
    )


# ── Line removal ──────────────────────────────────────────────────────────────

def remove_hv_lines(binary: np.ndarray, min_length: int = HV_LINE_MIN) -> np.ndarray:
    """
    Erase long horizontal and vertical lines from the binary image.

    Morphological OPEN with a long thin kernel extracts only the portions of
    white-pixel runs that span at least min_length in one direction — i.e. the
    wall lines, dimension lines, and grid lines that typically cross symbol
    outlines in blueprints. Subtracting them breaks the connected-component
    link between a symbol outline and whatever crosses it.
    """
    h_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (min_length, 1))
    v_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, min_length))
    h_lines  = cv2.morphologyEx(binary, cv2.MORPH_OPEN, h_kernel)
    v_lines  = cv2.morphologyEx(binary, cv2.MORPH_OPEN, v_kernel)
    return cv2.subtract(binary, cv2.add(h_lines, v_lines))


# ── OCR ───────────────────────────────────────────────────────────────────────

def _tess(img: np.ndarray, psm: int) -> str:
    """
    Call Tesseract via stdin/stdout subprocess — bypasses pytesseract's
    error handler which crashes on Tesseract 5.5's binary stderr output.
    Kept as a utility but no longer used in the main pipeline.
    """
    _, buf = cv2.imencode(".png", img)
    r = subprocess.run(
        ["tesseract", "stdin", "stdout",
         "--psm", str(psm),
         "-c", "tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"],
        input=buf.tobytes(),
        capture_output=True,
    )
    return r.stdout.decode("utf-8", errors="replace").strip().replace(" ", "")


def _eocr(img: np.ndarray) -> str:
    """
    Run EasyOCR on a uint8 image (grayscale or BGR) and return the
    concatenated recognised text, stripped of spaces.

    EasyOCR's neural recognition substantially outperforms Tesseract on
    small architectural text — in particular it correctly separates visually
    similar glyphs such as E/C and B/8 that Tesseract confuses at small sizes.
    The allowlist restricts output to the valid character set for diamond codes.
    """
    try:
        results = _READER.readtext(
            img,
            allowlist=_ALLOWLIST,
            detail=0,
            paragraph=False,
        )
        return "".join(results).strip().replace(" ", "")
    except Exception:
        return ""


def _ocr_with_mask(
    tile: np.ndarray,
    mask: np.ndarray,
    bbox: tuple[int, int, int, int],
    upscale: int = OCR_UPSCALE,
) -> str:
    """
    Mask out a shape's outline using a precomputed full-tile binary mask
    (255 = interior, 0 = background/outline), leaving only the interior
    (letter + digit), then OCR the clean crop.

    `mask` must already be eroded inward enough to clear outline pixels.
    Shared by every shape's OCR helper (mask built from a detected contour)
    and the exemplar-driven matcher (mask built once from the user's example
    image and repositioned per match).
    """
    x, y, w, h = bbox
    gray = cv2.cvtColor(tile, cv2.COLOR_BGR2GRAY)
    interior = gray.copy()
    interior[mask == 0] = 255

    # Crop to the bounding box and upscale for recognition.
    # EasyOCR operates on grayscale — no binarisation needed (the neural model
    # uses full tonal information and is less sensitive to thresholding artefacts
    # that caused Tesseract to confuse glyphs like E/C).
    crop = interior[max(0, y):y + h, max(0, x):x + w]
    if crop.size == 0:
        return ""
    crop = cv2.resize(crop, None, fx=upscale, fy=upscale, interpolation=cv2.INTER_CUBIC)

    text = _eocr(crop)
    return text if re.match(r"^[A-Z]\d$", text) else ""


# ── Cross-process progress reporting ───────────────────────────────────────────
# Shared by every shape's detect_page() pipeline (diamond, hexagon, ...) and
# the generic exemplar correlation pass in classical_cv_detector.py, so
# tile-level progress surfaces in the UI regardless of which shape is being
# detected. Two prior channels both failed to surface this for a real stall
# (print()-based logging -- a multi-hour capture across two full pool
# lifetimes captured zero of ~90+ expected lines -- and a
# multiprocessing.Manager().dict() proxy, which also showed nothing).
# Writing to a plain file both the worker and the main Streamlit process
# share unconditionally is about as hard to silently break as cross-process
# communication gets. Never let a write failure crash real detection work
# over a diagnostic.
_REPORT_FAILED_KEYS: set = set()  # progress_keys already logged a write failure for -- avoid spamming


def report_progress(progress_dir: str | None, progress_key, msg: str) -> None:
    if progress_dir is None:
        return
    try:
        with open(os.path.join(progress_dir, f"page_{progress_key}.txt"), "w") as f:
            f.write(msg)
    except Exception as e:
        if progress_key not in _REPORT_FAILED_KEYS:
            _REPORT_FAILED_KEYS.add(progress_key)
            print(f"[cv_core pid={os.getpid()}] progress write FAILED for "
                  f"page {progress_key}: {e!r} (dir={progress_dir!r})", flush=True)


# ── Spatial NMS index ───────────────────────────────────────────────────────
# Every exemplar-driven correlation pass (classical_cv_detector's generic
# exemplar_template_match_pass, shapes/diamond.py's exemplar_match_pass) does
# an "is any point within NMS_DIST of (x,y) already known" check for every
# raw correlation peak, against a plain growing list -- O(n) per query.
# Measured directly on a real hexagon page: exemplar_template_match_pass's
# NMS+OCR loop was still only 27% through 100,581 raw peaks after 124
# minutes. A permissive correlation threshold swept across up to 11 scales
# x 2 binarised variants x overlapping tiles naturally rediscovers the same
# handful of true local maxima dozens of times over, so raw_peaks routinely
# reaches 10^5-10^6 entries even though the true number of distinct
# candidate positions is far smaller.
class SpatialIndex:
    """
    Grid-based spatial index for the `abs(dx) < dist and abs(dy) < dist`
    square-window membership test used throughout the NMS passes. Bucketing
    by a grid cell sized to `dist` makes each query O(points in the 3x3
    neighbouring cells) instead of O(n) -- close to O(1) for realistically
    sparse detections, regardless of how many points have been indexed.
    """

    def __init__(self, dist: int):
        self.dist = dist
        self._cells: dict[tuple[int, int], list[tuple[int, int]]] = {}

    def _key(self, x: int, y: int) -> tuple[int, int]:
        return (x // self.dist, y // self.dist)

    def near_any(self, x: int, y: int) -> bool:
        gx, gy = self._key(x, y)
        for dgx in (-1, 0, 1):
            for dgy in (-1, 0, 1):
                for ex, ey in self._cells.get((gx + dgx, gy + dgy), ()):
                    if abs(x - ex) < self.dist and abs(y - ey) < self.dist:
                        return True
        return False

    def add(self, x: int, y: int) -> None:
        self._cells.setdefault(self._key(x, y), []).append((x, y))


# ── Tiling ────────────────────────────────────────────────────────────────────

def iter_tiles(
    bgr: np.ndarray, tile_size: int, overlap: float
):
    """Yield (tile_bgr, x_offset, y_offset) for every overlapping tile."""
    H, W = bgr.shape[:2]
    stride = max(1, int(tile_size * (1 - overlap)))
    for y0 in range(0, H, stride):
        for x0 in range(0, W, stride):
            x1 = min(W, x0 + tile_size)
            y1 = min(H, y0 + tile_size)
            yield bgr[y0:y1, x0:x1], x0, y0


def deduplicate(dets: list[dict], min_dist: int = NMS_DIST) -> list[dict]:
    """
    Remove duplicate detections caused by a symbol appearing in multiple tiles.
    Two detections are duplicates if their bbox centres are within min_dist px.
    """
    kept: list[dict] = []
    for d in dets:
        x, y, w, h = d["bbox"]
        cx, cy = x + w / 2.0, y + h / 2.0
        is_dup = any(
            abs(cx - (k["bbox"][0] + k["bbox"][2] / 2.0)) < min_dist
            and abs(cy - (k["bbox"][1] + k["bbox"][3] / 2.0)) < min_dist
            for k in kept
        )
        if not is_dup:
            kept.append(d)
    return kept


# ── Ground-truth recall check (pre-existing "Marked" PDF highlight boxes) ──────
# Some source PDFs (the "... Marked.pdf" variants) already have every target
# window symbol highlighted with a red rectangle/hexagon outline, drawn by
# whoever prepared the marked-up drawing set -- independent of and pre-dating
# anything this pipeline draws. That makes them usable as free ground truth:
# comparing our own detections against these pre-existing boxes tells us how
# many real symbols we're missing (false negatives), which the confirmed/
# flagged/shape-only split can't answer on its own (that split only concerns
# OCR/legend outcome for symbols we *did* find).
#
# IMPORTANT: find_marked_boxes() must be called on the RAW page (straight out
# of load_pages(), before annotate() runs) -- annotate() draws its own
# "flagged" boxes in a similar red, and the two would be indistinguishable on
# an already-annotated image.
MARKED_BOX_MIN_AREA = 400     # px^2 -- filters sub-symbol-sized red specks
MARKED_BOX_MAX_AREA = 20000   # px^2 -- filters large unrelated red regions


def find_marked_boxes(
    page_bgr: np.ndarray,
    min_area: int = MARKED_BOX_MIN_AREA,
    max_area: int = MARKED_BOX_MAX_AREA,
) -> list[tuple[int, int, int, int]]:
    """
    Find pre-existing red highlight boxes on a RAW (un-annotated) page.

    Thresholds directly on colour (high red channel, low green/blue) rather
    than via binarize()/contour-shape logic, since the markup is a solid,
    saturated red distinct from the black/grey line work everywhere else on
    a blueprint page. Returns one (x, y, w, h) bounding box per contiguous
    red outline whose bounding-box area falls in [min_area, max_area].
    """
    b, g, r = cv2.split(page_bgr.astype(np.int16))
    red_mask = ((r > 150) & (g < 100) & (b < 100)).astype(np.uint8) * 255

    contours, _ = cv2.findContours(red_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    boxes = []
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if min_area <= w * h <= max_area:
            boxes.append((x, y, w, h))
    return boxes


def match_marked_boxes(
    marked_boxes: list[tuple[int, int, int, int]],
    dets: list[dict],
    min_dist: int = NMS_DIST,
) -> dict:
    """
    Cross-reference pre-existing marked boxes (see find_marked_boxes) against
    this pipeline's own detections to measure recall.

    A marked box counts as "matched" if any detection's centre falls within
    tolerance of the marked box's centre. Tolerance is the larger of
    min_dist (NMS_DIST) and half the marked box's own footprint plus a
    margin, since these are hand-placed highlight boxes and won't align
    pixel-for-pixel with a detection's own bbox the way two independent
    passes of this pipeline would.

    Returns {"total_marked", "matched", "missed": [(x,y,w,h), ...]}.
    """
    det_centres = [
        (d["bbox"][0] + d["bbox"][2] / 2.0, d["bbox"][1] + d["bbox"][3] / 2.0)
        for d in dets
    ]

    missed: list[tuple[int, int, int, int]] = []
    for (mx, my, mw, mh) in marked_boxes:
        mcx, mcy = mx + mw / 2.0, my + mh / 2.0
        tol_x = max(min_dist, mw / 2.0 + 10)
        tol_y = max(min_dist, mh / 2.0 + 10)
        found = any(
            abs(mcx - dcx) < tol_x and abs(mcy - dcy) < tol_y
            for dcx, dcy in det_centres
        )
        if not found:
            missed.append((mx, my, mw, mh))

    return {
        "total_marked": len(marked_boxes),
        "matched": len(marked_boxes) - len(missed),
        "missed": missed,
    }