legend_parser.py

"""
legend_parser.py
Parse a screenshot of a window/door schedule legend table into a pandas DataFrame.

Usage:
    from legend_parser import parse_legend_image
    df = parse_legend_image("path/to/schedule_screenshot.png")

The parser uses EasyOCR to locate all text bounding boxes, clusters them
into rows by vertical proximity, then sorts each row left-to-right to
reconstruct the column structure.  No table lines need to be detected —
the spatial layout of text alone is sufficient.
"""

import re
from pathlib import Path

import cv2
import easyocr
import numpy as np
import pandas as pd

# Re-use the shared reader if already initialised, otherwise create one here.
try:
    from classical_cv_detector import _READER as _reader
except ImportError:
    _reader = easyocr.Reader(["en"], gpu=False, verbose=False)


# ── Core table-extraction logic ───────────────────────────────────────────────

def _extract_text_boxes(image: np.ndarray) -> list[dict]:
    """
    Run EasyOCR on the image and return a list of
    {text, x, y, w, h, cx, cy} dicts sorted top-to-bottom.
    """
    results = _reader.readtext(image, detail=1, paragraph=False)
    boxes = []
    for (pts, text, conf) in results:
        text = text.strip()
        if not text:
            continue
        xs = [p[0] for p in pts]
        ys = [p[1] for p in pts]
        x, y = int(min(xs)), int(min(ys))
        w, h = int(max(xs) - x), int(max(ys) - y)
        boxes.append({
            "text": text,
            "x": x, "y": y, "w": w, "h": h,
            "cx": x + w // 2,
            "cy": y + h // 2,
            "conf": conf,
        })
    return sorted(boxes, key=lambda b: b["cy"])


def _cluster_rows(boxes: list[dict], row_gap: int | None = None) -> list[list[dict]]:
    """
    Group boxes into rows.  Two boxes belong to the same row if their
    vertical centres are within `row_gap` pixels of each other.
    If row_gap is None it is estimated as the median text height.
    """
    if not boxes:
        return []

    if row_gap is None:
        median_h = float(np.median([b["h"] for b in boxes]))
        row_gap  = max(4, int(median_h * 0.8))

    rows: list[list[dict]] = []
    current_row: list[dict] = [boxes[0]]
    ref_cy = boxes[0]["cy"]

    for box in boxes[1:]:
        if abs(box["cy"] - ref_cy) <= row_gap:
            current_row.append(box)
        else:
            rows.append(sorted(current_row, key=lambda b: b["cx"]))
            current_row = [box]
            ref_cy = box["cy"]

    if current_row:
        rows.append(sorted(current_row, key=lambda b: b["cx"]))

    return rows


# Patterns used during column assignment
_MEAS_RE     = re.compile(r"\d+'\-\d+\"")          # dimension: 3'-6"
_TYPEMARK_RE = re.compile(r"^[A-Z]\d{1,2}$")       # type mark: A1, J2
_TEXT_ONLY_RE = re.compile(r"^[A-Za-z@/\s]+$")     # pure text (no digits/quotes)

# Known OCR substitutions in the TYPE MARK column
_TYPEMARK_FIXES = {
    "41": "H1",   # 4 mis-scanned as H
    "HI": "H1",
    "H|": "H1",
    "Il": "J1",
    "J|": "J1",
}

def _correct_typemark(text: str) -> str:
    """Apply known single-character OCR corrections to a type-mark candidate."""
    upper = text.upper().strip()
    if upper in _TYPEMARK_FIXES:
        return _TYPEMARK_FIXES[upper]
    # Single stray letter that looks like it could be the letter part of a code
    # e.g. "R" in a row where F-series codes surround it → keep as-is; the
    # caller will drop it if it doesn't match ^[A-Z]\d{1,2}$ after full parsing.
    return text


def _assign_columns(
    rows: list[list[dict]],
    col_names: list[str],
    col_x_centers: list[int],
) -> list[dict]:
    """
    Content-aware column assignment:
    - Boxes matching a dimension pattern (3\'−6\") → W or H only
    - Boxes matching a type-mark pattern (A1, J2) → TYPE MARK only
    - Pure-text boxes (no digits or quote marks) → never assigned to W or H;
      routed to the nearest text column instead.
    - Everything else → nearest column by X distance.

    This prevents OCR fragments like "INTERIOR" or "OPERABLE" from landing
    in the W/H columns when they appear at an ambiguous X position.
    """
    dim_cols  = [c for c in col_names if c in ("W", "H")]
    text_cols = [c for c in col_names if c not in ("W", "H")]

    records = []
    for row in rows:
        record: dict[str, list[str]] = {c: [] for c in col_names}
        for box in row:
            text = box["text"]
            if "TYPE MARK" in col_names and _TYPEMARK_RE.match(text):
                record["TYPE MARK"].append(_correct_typemark(text))
            elif _MEAS_RE.search(text) and dim_cols:
                # Assign to the W or H column whose centre is nearest
                d_xs  = [col_x_centers[col_names.index(c)] for c in dim_cols]
                col   = dim_cols[int(np.argmin([abs(box["cx"] - x) for x in d_xs]))]
                record[col].append(text)
            elif _TEXT_ONLY_RE.match(text) and text_cols:
                # Never put pure text into a dimension column
                t_xs = [col_x_centers[col_names.index(c)] for c in text_cols]
                col  = text_cols[int(np.argmin([abs(box["cx"] - x) for x in t_xs]))]
                record[col].append(text)
            else:
                dists = [abs(box["cx"] - cx) for cx in col_x_centers]
                col   = col_names[int(np.argmin(dists))]
                record[col].append(text)
        records.append({k: " ".join(v) for k, v in record.items()})
    return records


# ── Column detection ──────────────────────────────────────────────────────────

def _detect_columns(
    header_row: list[dict],
    all_rows: list[list[dict]],
) -> tuple[list[str], list[int]]:
    """
    Infer column names and their X-centre positions from the header row,
    falling back to equal-spacing if the header can't be parsed.

    For window/door schedules the expected columns are:
      TYPE MARK | DESCRIPTION | W | H | COMMENTS
    """
    # Try to read standard schedule headers from the detected header row text
    header_text = " ".join(b["text"].upper() for b in header_row)

    standard = {
        "TYPE MARK": ["TYPE", "MARK"],
        "DESCRIPTION": ["DESCRIPTION", "DESC"],
        "W": ["W"],
        "H": ["H"],
        "COMMENTS": ["COMMENT", "NOTE", "REMARKS"],
    }

    # Match each box in the header row to a standard column name
    col_names: list[str] = []
    col_x_centers: list[int] = []

    for box in header_row:
        t = box["text"].upper()
        matched = None
        for name, keywords in standard.items():
            if any(k in t for k in keywords):
                matched = name
                break
        if matched and matched not in col_names:
            col_names.append(matched)
            col_x_centers.append(box["cx"])

    # Fallback: generic column numbering based on equally-spaced boxes
    if len(col_names) < 3:
        col_names     = [f"col_{i}" for i in range(len(header_row))]
        col_x_centers = [b["cx"] for b in header_row]

    return col_names, col_x_centers


# ── Shared header-row detection (used by both the diamond and hexagon parsers) ─

def _find_header_span(
    rows: list[list[dict]],
    is_header: "callable[[list[dict]], bool]",
    max_span: int = 2,
) -> tuple[int, int]:
    """
    Scan `rows` for the column-header row(s), returning (start, end) such
    that rows[start:end] (concatenated) is the header and rows[end:] is the
    data.

    Some schedule exports wrap column headers onto two lines (e.g. "TYPE"
    alone on one line, with "MARK" and the rest of the header row below it —
    see Original's "Test Legend.png"). Assuming the header is always exactly
    one row previously made column detection silently fail whenever a
    schedule wrapped its header this way — `_detect_columns` would see only
    one header box, decide it didn't recognise enough columns, and fall back
    to generic col_0/col_1/... names, permanently losing the type-mark
    column.

    Tries each row index in turn as a candidate header start, first alone
    then merged with the following row (up to `max_span` rows), and returns
    the first span for which `is_header(merged_boxes)` returns True. Falls
    back to treating row 0 alone as the header if nothing matches (the old
    fixed-position behaviour), so a legend image with an unrecognisable
    header still parses as before rather than raising.
    """
    for i in range(len(rows)):
        for span in range(1, max_span + 1):
            if i + span > len(rows):
                break
            merged = [b for r in rows[i:i + span] for b in r]
            if is_header(merged):
                return i, i + span
    return 0, 1


# ── Public API ────────────────────────────────────────────────────────────────

def parse_legend_image(
    image_path: str | Path,
    type_col: str = "TYPE MARK",
    skip_title_rows: int = 1,
    manual_codes: list[str] | None = None,
) -> pd.DataFrame:
    """
    Parse a schedule legend screenshot into a DataFrame.

    Parameters
    ----------
    image_path      Path to the image file (PNG, JPG, etc.).
    type_col        Name of the column that contains the type-mark code
                    (e.g. 'A1', 'J2').  Used to filter out non-data rows.
    skip_title_rows Number of title rows above the column-header row to skip
                    (e.g. the 'WINDOW SCHEDULE' banner).
    manual_codes    Optional list of type-mark codes to add if OCR missed them
                    (e.g. ['E4', 'F2', 'J1']).  Added as rows with empty
                    description/dimension fields for validation purposes.

    Returns
    -------
    pd.DataFrame with one row per legend entry and a column for each
    detected table column.  Rows where the type-mark cell is empty or
    does not look like a code (letter + digit) are dropped.
    Prints a warning for any row whose detected type-mark looks like an
    OCR misread (single letter, all-digit string, etc.).
    """
    img_path = Path(image_path)
    if not img_path.exists():
        raise FileNotFoundError(f"Image not found: {img_path}")

    image = cv2.imread(str(img_path))
    if image is None:
        raise ValueError(f"Could not read image: {img_path}")

    # Upscale small images for better OCR accuracy
    h, w = image.shape[:2]
    if max(h, w) < 1200:
        scale  = 1200 / max(h, w)
        image  = cv2.resize(image, None, fx=scale, fy=scale,
                            interpolation=cv2.INTER_CUBIC)

    boxes = _extract_text_boxes(image)
    rows  = _cluster_rows(boxes)

    if len(rows) < skip_title_rows + 2:
        raise ValueError("Too few text rows detected — check the image quality.")

    # Auto-detect the header row(s) (searching from skip_title_rows onward),
    # merging up to 2 consecutive rows so a header wrapped onto two lines
    # (e.g. "TYPE" alone above "MARK | DESCRIPTION | W | H | COMMENTS") is
    # still recognised as one header instead of two under-populated ones.
    search_rows = rows[skip_title_rows:]
    rel_start, rel_end = _find_header_span(
        search_rows,
        is_header=lambda boxes: len(_detect_columns(boxes, [])[0]) >= 3,
    )
    header_row = [b for r in search_rows[rel_start:rel_end] for b in r]
    data_rows  = search_rows[rel_end:]

    col_names, col_x_centers = _detect_columns(header_row, data_rows)

    records = _assign_columns(data_rows, col_names, col_x_centers)
    df      = pd.DataFrame(records, columns=col_names)

    # Drop rows where the type-mark cell is empty or is just a number
    if type_col in df.columns:
        # Apply type-mark corrections before filtering
        df[type_col] = df[type_col].map(
            lambda v: _correct_typemark(v) if isinstance(v, str) else v
        )

        # If TYPE MARK has extra words (e.g. "D1 1x2 FIXED"), the first token
        # is the real code and the rest belong in DESCRIPTION.
        def _split_typemark(row):
            tm = str(row.get(type_col, "")).strip()
            if " " in tm:
                parts = tm.split(None, 1)
                if re.match(r"^[A-Z]\d", parts[0]):
                    row = row.copy()
                    row[type_col] = parts[0]
                    desc_col = "DESCRIPTION" if "DESCRIPTION" in row else type_col
                    if desc_col != type_col:
                        existing = str(row.get(desc_col, "")).strip()
                        row[desc_col] = (parts[1] + " " + existing).strip()
            return row

        df = df.apply(_split_typemark, axis=1)

        # Surface rows dropped here instead of silently discarding them -- a
        # row skipped because its type-mark cell was empty or unreadable
        # (e.g. a missing OCR glyph, or a digit/letter misread like
        # "D2" -> "02") is exactly the kind of failure that previously only
        # surfaced after a multi-hour detection run had already finished.
        # Printing the row's other cell contents lets a human immediately
        # recognise which legend entry silently vanished.
        mask    = df[type_col].str.match(r"^[A-Z]\d", na=False)
        dropped = df[~mask]
        for _, r in dropped.iterrows():
            other = {c: r[c] for c in df.columns if c != type_col and str(r[c]).strip()}
            if other or str(r[type_col]).strip():
                print(f"  ⚠  Legend row skipped (unrecognisable {type_col}="
                      f"{r[type_col]!r}): other cells = {other}")

        df = df[mask].copy()
        df = df.reset_index(drop=True)

    # Warn about single-character or all-digit type marks that survived filtering
    # (they are likely OCR misreads that the correction pass did not cover)
    if type_col in df.columns:
        suspicious = df[~df[type_col].str.match(r"^[A-Z]\d{1,2}$", na=False)][type_col]
        if not suspicious.empty:
            print(f"  ⚠  Suspicious type marks (may be OCR misreads): "
                  f"{sorted(suspicious.tolist())} — check manually.")

    # Merge any manually-supplied codes that OCR missed
    if manual_codes:
        existing = set(df[type_col].tolist()) if type_col in df.columns else set()
        new_rows = [
            {c: (code if c == type_col else "") for c in df.columns}
            for code in manual_codes
            if re.match(r"^[A-Z]\d{1,2}$", code) and code not in existing
        ]
        if new_rows:
            df = pd.concat([df, pd.DataFrame(new_rows)], ignore_index=True)
            df = df.sort_values(type_col).reset_index(drop=True)
            print(f"  ℹ  Manually added codes: {[r[type_col] for r in new_rows]}")

    return df


def valid_codes(df: pd.DataFrame, type_col: str = "TYPE MARK") -> set[str]:
    """
    Return the set of valid type-mark codes extracted from a legend DataFrame:
    a letter followed by 1-2 digits (matches shapes.diamond.CODE_RE and
    _TYPEMARK_RE above) -- e.g. "A1" or "J12".
    """
    if type_col not in df.columns:
        raise KeyError(f"Column '{type_col}' not found in DataFrame.")
    codes = df[type_col].dropna().str.strip()
    return {c for c in codes if re.match(r"^[A-Z]\d{1,2}$", c)}


# ── Hexagon / single-letter legend parsing ────────────────────────────────────

# Hexagon legend headers that identify the code column ("Window Letter", "Letter", etc.)
_HEX_CODE_HEADERS = {"WINDOW LETTER", "WINDOW NUMBER", "LETTER", "WIN LETTER", "TYPE", "MARK"}


def _detect_hexagon_columns(
    header_row: list[dict],
) -> tuple[str | None, int | None]:
    """
    Find the code column in a hexagon window-schedule header row.
    Returns (column_label, x_center) or (None, None) if not found.
    """
    for box in header_row:
        text_upper = box["text"].upper().strip()
        # Accept "LETTER", "WINDOW LETTER", "WIN LETTER" or "TYPE" as the code column
        if any(kw in text_upper for kw in _HEX_CODE_HEADERS):
            return text_upper, box["cx"]
    # Fallback: the first column (leftmost box in the header row) is often the code
    if header_row:
        return "LETTER", header_row[0]["cx"]
    return None, None


def parse_hexagon_legend_image(
    image_path,
    skip_title_rows: int = 1,
    manual_codes=None,
) -> pd.DataFrame:
    """
    Parse a hexagon/window-schedule legend image into a DataFrame.

    Hexagon schedules use single capital letters (A, B, C …) as codes,
    unlike diamond schedules which use letter+digit codes (A1, J2 …).

    Returns a DataFrame with a 'WINDOW LETTER' column containing the
    single-letter codes found in the schedule, plus any additional columns
    the OCR can identify (Width, Height, Count, etc.).
    """
    from pathlib import Path as _Path

    img_path = _Path(image_path)
    if not img_path.exists():
        raise FileNotFoundError(f"Image not found: {img_path}")

    image = cv2.imread(str(img_path))
    if image is None:
        raise ValueError(f"Could not read image: {img_path}")

    h, w = image.shape[:2]
    if max(h, w) < 1200:
        scale = 1200 / max(h, w)
        image = cv2.resize(image, None, fx=scale, fy=scale,
                           interpolation=cv2.INTER_CUBIC)

    boxes = _extract_text_boxes(image)
    rows  = _cluster_rows(boxes)

    if len(rows) < 2:
        raise ValueError("Too few text rows detected — check the image quality.")

    # Auto-detect the header row(s) (searching from skip_title_rows onward),
    # merging up to 2 consecutive rows so a header wrapped onto two lines is
    # still recognised as one header rather than missed entirely -- shares
    # _find_header_span with parse_legend_image (see there for why the
    # multi-row merge matters).
    search_rows = rows[skip_title_rows:]
    rel_start, rel_end = _find_header_span(
        search_rows,
        is_header=lambda boxes: any(
            kw in " ".join(b["text"].upper() for b in boxes) for kw in _HEX_CODE_HEADERS
        ),
    )
    if rel_end >= len(search_rows):
        raise ValueError("Could not locate column-header row in the legend image.")

    header_row = [b for r in search_rows[rel_start:rel_end] for b in r]
    data_rows  = search_rows[rel_end:]

    _, code_x = _detect_hexagon_columns(header_row)

    # Extract the leftmost token from each data row; keep it when it's a letter
    # optionally followed by up to 2 digits (possibly misread by OCR as
    # lowercase) — matches shapes.hexagon.CODE_RE. Most hexagon schedules
    # (Acama, Lexington) use a bare letter; Colorado Grand Oaks uses
    # letter+digit codes (W1..W11), so a single-letter-only pattern here
    # would truncate "W1" down to "W" and make every letter+digit code in
    # that legend fail to validate.
    _HEX_CODE_RE = re.compile(r"^[A-Z][0-9]{0,2}$")
    records: list[dict] = []
    for row in data_rows:
        if not row:
            continue
        # The code is the leftmost box (nearest to code_x if we found a header col)
        if code_x is not None:
            row_sorted = sorted(row, key=lambda b: abs(b["cx"] - code_x))
        else:
            row_sorted = sorted(row, key=lambda b: b["cx"])

        code_text = row_sorted[0]["text"].strip().upper()
        # Accept multi-char OCR fragments where the code may have been merged
        # with the next token (e.g. "A2'-6\"" → "A2", "W1'-6\"" → "W1"), but
        # only when the character immediately after the matched prefix is
        # NOT itself a letter -- otherwise this would truncate an unrelated
        # word into a false code (e.g. "PICTURE" -> "P", or an OCR digit/
        # letter confusion like "W4" misread as "WL" -> "W", "W10" misread
        # as "W1O" -> "W1"). In those letter-follows cases the row's leading
        # cell isn't a real code at all (or is unrecoverably garbled), so it
        # must be dropped rather than guessed at.
        if not _HEX_CODE_RE.match(code_text):
            m = re.match(r"^[A-Z][0-9]{0,2}", code_text)
            if m and not (m.end() < len(code_text) and code_text[m.end()].isalpha()):
                code_text = m.group(0)
            else:
                # Surface the skip instead of silently discarding the row --
                # this is exactly the kind of failure (missing OCR glyph,
                # digit/letter misread like "W1"->"WI") that previously only
                # showed up after a multi-hour detection run had finished.
                other = [b["text"] for b in row_sorted[1:]]
                print(f"  ⚠  Legend row skipped (unrecognisable code {code_text!r}): "
                      f"other cells = {other}")
                continue   # skip rows with no recognisable (or unrecoverably garbled) code

        records.append({"WINDOW LETTER": code_text})

    df = pd.DataFrame(records).drop_duplicates("WINDOW LETTER").reset_index(drop=True)

    # Merge manually-supplied codes
    if manual_codes:
        existing = set(df["WINDOW LETTER"].tolist())
        new_rows = [
            {"WINDOW LETTER": c.upper()}
            for c in manual_codes
            if re.match(r"^[A-Z][0-9]{0,2}$", c.upper()) and c.upper() not in existing
        ]
        if new_rows:
            df = pd.concat([df, pd.DataFrame(new_rows)], ignore_index=True)
            df = df.sort_values("WINDOW LETTER").reset_index(drop=True)

    return df


def valid_codes_hexagon(df: pd.DataFrame, type_col: str = "WINDOW LETTER") -> set[str]:
    """
    Return the set of valid hexagon codes from a legend DataFrame: a letter
    optionally followed by up to 2 digits (matches shapes.hexagon.CODE_RE) —
    covers both bare-letter schedules (Acama, Lexington) and letter+digit
    schedules (Colorado Grand Oaks' W1..W11).
    """
    if type_col not in df.columns:
        raise KeyError(f"Column '{type_col}' not found in DataFrame.")
    codes = df[type_col].dropna().str.strip().str.upper()
    return {c for c in codes if re.match(r"^[A-Z][0-9]{0,2}$", c)}


# ── CLI entry point ───────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse, sys

    parser = argparse.ArgumentParser(
        description="Parse a schedule legend screenshot into a CSV table."
    )
    parser.add_argument("image",  help="Path to the legend screenshot")
    parser.add_argument("--out",  default=None,
                        help="Save parsed table to this CSV path")
    parser.add_argument("--skip", type=int, default=1,
                        help="Title rows to skip above the header (default 1)")
    args = parser.parse_args()

    try:
        df = parse_legend_image(args.image, skip_title_rows=args.skip)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

    print(df.to_string(index=False))
    print(f"\nValid codes: {sorted(valid_codes(df))}")

    if args.out:
        df.to_csv(args.out, index=False)
        print(f"\nSaved to '{args.out}'")