#!/usr/bin/env python3
"""
ShotTrack OCR Engine - Phase 4/6/10
Extracts counter/date/time from injection-moulding machine HMI photos.
Usage:
  Full-image:  python ocr_engine.py <photo_path>
  With regions: python ocr_engine.py <photo_path> --regions '[{"fieldName":"counter",...}]'
Output: JSON to stdout
"""
import sys
import json
import time
import re
import argparse
from collections import Counter as Ctr
from pathlib import Path

try:
    import easyocr
    import cv2
    import numpy as np
except ImportError as e:
    print(json.dumps({"error": f"Missing dependency: {e}. Run: pip install easyocr opencv-python-headless numpy"}), flush=True)
    sys.exit(1)

# Phase 0 fix: remove \b boundary to catch fused labels like "04947Cntl"
CNT_LABEL = re.compile(r"(?i)(cnt|count|pcs|shot|parts|total)")
NON_COUNTER = re.compile(r"(?i)(screw|mold|rpm|kwh|run|hour|temp|pressure|barrel|nozzle|inj|fill|hold|cool|clamp|eject|gate|cushion|transfer|position|speed|alarm|status)")

DATE_PAT = re.compile(r"(\d{2})[/\-\.](\d{2})[/\-\.](\d{2,4})")
TIME_PAT = re.compile(r"\b(\d{1,2})[:\.](\d{2})[:\.](\d{2})\b")


def preprocess(img, variant: str):
    if variant == "original":
        return img
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    if variant == "inverted":
        return cv2.cvtColor(cv2.bitwise_not(gray), cv2.COLOR_GRAY2BGR)
    if variant == "clahe":
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        return cv2.cvtColor(clahe.apply(gray), cv2.COLOR_GRAY2BGR)
    return img


def extract_counter(boxes):
    """
    3-strategy counter extraction (Phase 0 findings):
    1. Fused: label + digits in same text box
    2. Adjacent: digits box next to label-only box
    3. Standalone: 4-6 digit box with no nearby non-counter words
    """
    results = []

    for i, (bbox, text, conf) in enumerate(boxes):
        t = text.strip()

        # Strategy 1 — fused label
        if CNT_LABEL.search(t):
            cleaned = CNT_LABEL.sub("", t).strip()
            m = re.match(r"(\d{3,8})", re.sub(r"[^0-9]", " ", cleaned).strip())
            if m:
                results.append((int(m.group(1)), float(conf), "fused"))
                continue

        # Strategy 2 — adjacent label
        if re.fullmatch(r"\d{3,8}", t):
            for j in range(max(0, i - 3), min(len(boxes), i + 4)):
                if j == i:
                    continue
                adj_text = boxes[j][1].strip()
                adj_conf = float(boxes[j][2])
                if CNT_LABEL.search(adj_text) and not re.search(r"\d{3,}", adj_text):
                    results.append((int(t), (float(conf) + adj_conf) / 2, "adjacent"))
                    break

    # Strategy 3 — standalone (only if strategies 1+2 found nothing)
    if not results:
        for i, (bbox, text, conf) in enumerate(boxes):
            t = text.strip()
            if not re.fullmatch(r"\d{4,6}", t) or float(conf) < 0.45:
                continue
            bad_nearby = any(
                NON_COUNTER.search(boxes[j][1])
                for j in range(max(0, i - 4), min(len(boxes), i + 5))
                if j != i
            )
            if not bad_nearby:
                results.append((int(t), float(conf), "standalone"))

    if not results:
        return None, 0.0

    # Vote: pick most common value, boost conf when variants agree
    val_votes = Ctr(r[0] for r in results)
    best_val = val_votes.most_common(1)[0][0]
    best_confs = [r[1] for r in results if r[0] == best_val]
    ratio = len(best_confs) / len(results)
    base_conf = sum(best_confs) / len(best_confs)
    final_conf = min(base_conf + 0.08 * (ratio - 1.0 / max(len(results), 1)), 1.0)
    return best_val, round(final_conf, 4)


def extract_date(boxes):
    for _, text, conf in boxes:
        m = DATE_PAT.search(text)
        if m:
            return text.strip()
    return None


def extract_time(boxes):
    for _, text, conf in boxes:
        m = TIME_PAT.search(text)
        if m and float(conf) > 0.55:
            return text.strip()
    return None


def boxes_to_list(raw_boxes):
    return [
        {
            "bbox": [[int(c) for c in pt] for pt in b[0]],
            "text": b[1],
            "conf": round(float(b[2]), 4),
        }
        for b in raw_boxes
    ]


def crop_region(img, region):
    """Crop image to normalized region coordinates (0.0–1.0)."""
    h, w = img.shape[:2]
    x1 = max(0, int(region["x"] * w))
    y1 = max(0, int(region["y"] * h))
    x2 = min(w, int((region["x"] + region["w"]) * w))
    y2 = min(h, int((region["y"] + region["h"]) * h))
    if x2 <= x1 or y2 <= y1:
        return img
    return img[y1:y2, x1:x2]


# ── Phase 10: Tesseract seven-segment support ─────────────────────────────────

def _tesseract_available():
    try:
        import pytesseract
        pytesseract.get_tesseract_version()
        return True
    except Exception:
        return False


def run_tesseract_ssd(crop_gray):
    """
    Run Tesseract with ssd traineddata for seven-segment display digits.
    Falls back to digit-only whitelist if ssd lang not installed.
    """
    try:
        import pytesseract
        # Scale up small regions for better recognition
        h, w = crop_gray.shape
        if max(h, w) < 120:
            scale = 120 / max(h, w)
            crop_gray = cv2.resize(
                crop_gray,
                (int(w * scale), int(h * scale)),
                interpolation=cv2.INTER_CUBIC,
            )

        # Try ssd traineddata first (Tesseract seven-segment model)
        try:
            config = "--psm 7 --oem 1 -l ssd"
            raw = pytesseract.image_to_string(crop_gray, config=config).strip()
        except Exception:
            # ssd not installed — fall back to digits + dot whitelist
            config = r"--psm 7 --oem 3 -c tessedit_char_whitelist=0123456789."
            raw = pytesseract.image_to_string(crop_gray, config=config).strip()

        m = re.search(r"\d+\.?\d*", raw)
        val = m.group(0) if m else None
        conf = 0.90 if val else 0.10
        return val, conf, raw
    except ImportError:
        return None, 0.0, ""


def run_region_ocr(img, reader, regions):
    """Run OCR on each specified region and return field-specific results."""
    results = {}
    for region in regions:
        field = region["fieldName"]
        engine = region.get("ocrEngine", "easyocr")
        crop = crop_region(img, region)
        variant = region.get("preprocessVariant", "original")
        processed = preprocess(crop, variant)

        # ── Phase 10: Tesseract path for seven-segment fields ─────────────────
        if engine == "tesseract-ssd":
            gray = cv2.cvtColor(processed, cv2.COLOR_BGR2GRAY)
            val, conf, raw_text = run_tesseract_ssd(gray)
            results[field] = {
                "value": val,
                "confidence": conf,
                "rawText": raw_text,
                "engine": "tesseract-ssd",
            }
            continue

        # ── EasyOCR path (default) ─────────────────────────────────────────────
        raw = reader.readtext(processed)
        tuples = [(b[0], b[1], b[2]) for b in raw]
        raw_texts = " | ".join(b[1] for b in raw)

        if "counter" in field:
            val, conf = extract_counter(tuples)
            results[field] = {"value": val, "confidence": conf, "rawText": raw_texts, "engine": "easyocr"}
        elif "date" in field:
            val = extract_date(tuples)
            results[field] = {"value": val, "confidence": 0.9 if val else 0.0, "rawText": raw_texts, "engine": "easyocr"}
        elif "time" in field:
            val = extract_time(tuples)
            results[field] = {"value": val, "confidence": 0.9 if val else 0.0, "rawText": raw_texts, "engine": "easyocr"}
        elif "cycle" in field:
            digits = [b[1].strip() for b in raw if re.search(r"\d+\.?\d*", b[1])]
            val = digits[0] if digits else None
            results[field] = {"value": val, "confidence": 0.7 if val else 0.0, "rawText": raw_texts, "engine": "easyocr"}
        else:
            val = raw[0][1] if raw else None
            conf = float(raw[0][2]) if raw else 0.0
            results[field] = {"value": val, "confidence": conf, "rawText": raw_texts, "engine": "easyocr"}

    return results


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("photo_path")
    parser.add_argument("--regions", default=None, help="JSON array of region objects")
    args = parser.parse_args()

    photo_path = args.photo_path
    t0 = time.time()

    img = cv2.imread(photo_path)
    if img is None:
        print(json.dumps({"error": f"Cannot read image: {photo_path}"}), flush=True)
        sys.exit(1)

    # Resize to max 1600px wide
    h, w = img.shape[:2]
    if w > 1600:
        img = cv2.resize(img, (1600, int(h * 1600 / w)))

    reader = easyocr.Reader(["en"], gpu=False, verbose=False)

    # ── Region-based OCR (Phase 6) ─────────────────────────────────────────────
    if args.regions:
        try:
            regions = json.loads(args.regions)
        except json.JSONDecodeError as e:
            print(json.dumps({"error": f"Invalid --regions JSON: {e}"}), flush=True)
            sys.exit(1)

        region_results = run_region_ocr(img, reader, regions)
        processing_ms = int((time.time() - t0) * 1000)

        counter_data = region_results.get("counter", {})
        date_data = region_results.get("date", {})
        time_data = region_results.get("time", {})
        cycle_data = region_results.get("cycle_time", {})

        output = {
            "counter": counter_data.get("value"),
            "counterConfidence": counter_data.get("confidence", 0.0),
            "date": date_data.get("value"),
            "time": time_data.get("value"),
            "cycleTimeSec": cycle_data.get("value"),
            "regions": region_results,
            "rawBoxes": {},
            "extractedFields": {
                "counter": counter_data.get("value"),
                "counterConfidence": counter_data.get("confidence", 0.0),
                "date": date_data.get("value"),
                "time": time_data.get("value"),
                "cycleTimeSec": cycle_data.get("value"),
            },
            "engineUsed": "easyocr+regions",
            "processingMs": processing_ms,
            "variantResults": {},
        }
        print(json.dumps(output), flush=True)
        return

    # ── Full-image OCR (Phase 4 fallback) ─────────────────────────────────────
    variants = ["original", "inverted", "clahe"]
    variant_results = {}
    all_boxes = {}

    for variant in variants:
        processed = preprocess(img, variant)
        raw = reader.readtext(processed)
        box_list = boxes_to_list(raw)
        all_boxes[variant] = box_list
        tuples = [(b["bbox"], b["text"], b["conf"]) for b in box_list]
        cval, cconf = extract_counter(tuples)
        variant_results[variant] = {"counter": cval, "counterConf": cconf}

    # Cross-variant voting for counter
    all_counts = [
        (v["counter"], v["counterConf"])
        for v in variant_results.values()
        if v["counter"] is not None
    ]

    if all_counts:
        val_votes = Ctr(r[0] for r in all_counts)
        best_val = val_votes.most_common(1)[0][0]
        best_confs = [r[1] for r in all_counts if r[0] == best_val]
        ratio = len(best_confs) / len(all_counts)
        base_conf = sum(best_confs) / len(best_confs)
        final_conf = min(base_conf + 0.12 * (ratio - 0.34), 1.0)
        counter = best_val
        counter_confidence = round(final_conf, 4)
    else:
        counter = None
        counter_confidence = 0.0

    orig_tuples = [(b["bbox"], b["text"], b["conf"]) for b in all_boxes.get("original", [])]
    date_val = extract_date(orig_tuples)
    time_val = extract_time(orig_tuples)

    processing_ms = int((time.time() - t0) * 1000)

    output = {
        "counter": counter,
        "counterConfidence": counter_confidence,
        "date": date_val,
        "time": time_val,
        "cycleTimeSec": None,
        "rawBoxes": all_boxes,
        "extractedFields": {
            "counter": counter,
            "counterConfidence": counter_confidence,
            "date": date_val,
            "time": time_val,
        },
        "engineUsed": "easyocr",
        "processingMs": processing_ms,
        "variantResults": variant_results,
    }

    print(json.dumps(output), flush=True)


if __name__ == "__main__":
    main()