import cv2
import easyocr
import numpy as np
import json
import sys
import os
import re

# ─────────────────────────────────────────────
#  CONFIGURATION
# ─────────────────────────────────────────────
OCR_LANGUAGES     = ['en']
MIN_OCR_CONF      = 0.4     # raised to filter junk text
IC_LABELS         = ['ic', 'transistor', 'clock', 'display']
PADDING           = 10

# Junk patterns to filter out from OCR results
JUNK_PATTERNS = [
    r'^[^a-zA-Z0-9]+$',   # only symbols
    r'^\d{1,2}$',          # single/double digit only (too short to be useful)
    r'^[a-zA-Z]{1}$',      # single letter
]

# ─────────────────────────────────────────────
#  INITIALIZE READER
# ─────────────────────────────────────────────
print("[->] Loading EasyOCR model...")
reader = easyocr.Reader(OCR_LANGUAGES, gpu=True)
print("[OK] EasyOCR ready")


# ─────────────────────────────────────────────
#  FILTER JUNK OCR TEXT
# ─────────────────────────────────────────────
def is_junk(text: str) -> bool:
    for pattern in JUNK_PATTERNS:
        if re.match(pattern, text):
            return True
    return False


# ─────────────────────────────────────────────
#  PREPROCESS CHIP PATCH
# ─────────────────────────────────────────────
def preprocess_patch(patch: np.ndarray) -> np.ndarray:
    h, w = patch.shape[:2]
    # Only upscale if patch is small
    scale = 3 if max(h, w) < 100 else 2
    upscaled  = cv2.resize(patch, (w * scale, h * scale),
                            interpolation=cv2.INTER_CUBIC)
    kernel    = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
    sharpened = cv2.filter2D(upscaled, -1, kernel)
    denoised  = cv2.fastNlMeansDenoisingColored(sharpened, h=10)
    return denoised


# ─────────────────────────────────────────────
#  RUN OCR ON A SINGLE PATCH
# ─────────────────────────────────────────────
def read_text_from_patch(patch: np.ndarray) -> list:
    processed = preprocess_patch(patch)
    results   = reader.readtext(processed)

    texts = []
    for (_, text, conf) in results:
        text = text.strip()
        if conf >= MIN_OCR_CONF and len(text) >= 2 and not is_junk(text):
            texts.append((text, round(conf, 3)))

    return texts


# ─────────────────────────────────────────────
#  RUN OCR ON ALL IC DETECTIONS
# ─────────────────────────────────────────────
def run_ocr_on_detections(image_path: str, detections: list) -> list:
    img = cv2.imread(image_path)
    if img is None:
        print(f"[X] Could not load image: {image_path}")
        return detections

    ih, iw = img.shape[:2]
    updated  = []

    ic_count = sum(1 for d in detections if d['label'] in IC_LABELS)
    print(f"\n[->] Running OCR on {ic_count} IC/chip regions...")

    for det in detections:
        label = det['label']

        if label not in IC_LABELS:
            det['ocr_text']   = []
            det['part_number']= "N/A"
            updated.append(det)
            continue

        x1, y1, x2, y2 = det['bbox']
        x1p = max(0,  x1 - PADDING)
        y1p = max(0,  y1 - PADDING)
        x2p = min(iw, x2 + PADDING)
        y2p = min(ih, y2 + PADDING)

        patch = img[y1p:y2p, x1p:x2p]

        if patch.size == 0:
            det['ocr_text']    = []
            det['part_number'] = "unknown"
            updated.append(det)
            continue

        texts    = read_text_from_patch(patch)
        combined = " ".join(t for t, c in texts).strip()

        det['ocr_text']    = texts
        det['part_number'] = combined if combined else "unknown"

        if texts:
            print(f"   [{label}] @ ({x1},{y1}) → '{combined}'")
        else:
            print(f"   [{label}] @ ({x1},{y1}) → (no text detected)")

        updated.append(det)

    return updated


# ─────────────────────────────────────────────
#  PRINT OCR SUMMARY
# ─────────────────────────────────────────────
def print_ocr_summary(detections: list):
    ic_dets = [d for d in detections if d['label'] in IC_LABELS]
    identified = [d for d in ic_dets if d.get('part_number', 'unknown') not in ('unknown', 'N/A', '')]

    print(f"\n-- OCR Summary ---------------------------")
    for det in ic_dets:
        label = det['label']
        part  = det.get('part_number', 'unknown')
        conf  = det['confidence']
        hits  = len(det.get('ocr_text', []))
        print(f"   {label:<15} | part: {part:<30} | conf: {conf:.0%} | ocr hits: {hits}")
    print(f"\n   Total ICs/chips : {len(ic_dets)}")
    print(f"   Text identified : {len(identified)}")
    print(f"------------------------------------------\n")


# ─────────────────────────────────────────────
#  ENTRY POINT
# ─────────────────────────────────────────────
if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python ocr.py <image_path> <results_json>")
        print("Example: python ocr.py sample5.jpg sample5_results.json")
        sys.exit(1)

    image_path   = sys.argv[1]
    results_json = sys.argv[2]

    with open(results_json) as f:
        data = json.load(f)

    detections = data.get("components", [])
    for d in detections:
        d['bbox'] = tuple(d['bbox'])

    print(f"[OK] Loaded {len(detections)} detections from {results_json}")

    updated = run_ocr_on_detections(image_path, detections)
    print_ocr_summary(updated)

    # Save updated JSON
    base     = os.path.splitext(results_json)[0]
    out_path = base + "_ocr.json"
    out_data = {
        "total_components": len(updated),
        "components": [
            {**d,
             "bbox":     list(d["bbox"]),
             "ocr_text": [[t, c] for t, c in d.get("ocr_text", [])]}
            for d in updated
        ]
    }
    with open(out_path, "w") as f:
        json.dump(out_data, f, indent=2)
    print(f"[OK] Updated results saved: {out_path}")