"""
utils.py — Reusable Computer Vision and Geometry Utilities for Meter OCR Pipeline
Optimized for Hugging Face Spaces & Standalone Inference.
"""
from __future__ import annotations

import math
from typing import Any

import cv2
import numpy as np
from PIL import Image


def rotate_image(img_bgr: np.ndarray, angle: int) -> np.ndarray:
    """
    Rotate BGR image by 90, 180, or 270 degrees clockwise.
    Returns original image unchanged if angle == 0.
    """
    if angle == 90:
        return cv2.rotate(img_bgr, cv2.ROTATE_90_CLOCKWISE)
    if angle == 180:
        return cv2.rotate(img_bgr, cv2.ROTATE_180)
    if angle == 270:
        return cv2.rotate(img_bgr, cv2.ROTATE_90_COUNTERCLOCKWISE)
    return img_bgr


def apply_prep(
    img_bgr: np.ndarray,
    prep: str,
    clahe_clip: float = 2.0,
    clahe_grid: tuple[int, int] = (8, 8),
) -> np.ndarray:
    """
    Apply image contrast enhancement:
    - 'clahe': Contrast Limited Adaptive Histogram Equalization on L channel (LAB color space)
    - 'histeq': Global Histogram Equalization on Y channel (YCrCb color space)
    - 'orig' or unknown: returns original image without modification
    """
    if prep == "clahe":
        lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=clahe_grid)
        l_clahe = clahe.apply(l)
        return cv2.cvtColor(cv2.merge([l_clahe, a, b]), cv2.COLOR_LAB2BGR)

    if prep == "histeq":
        ycrcb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2YCrCb)
        y, cr, cb = cv2.split(ycrcb)
        y_eq = cv2.equalizeHist(y)
        return cv2.cvtColor(cv2.merge([y_eq, cr, cb]), cv2.COLOR_YCrCb2BGR)

    return img_bgr


def remap_point(x: float, y: float, angle: int, w: int, h: int) -> tuple[float, float]:
    """
    Map coordinate (x, y) from a rotated image frame back to original image dimensions (w, h).
    """
    if angle == 90:
        return y, h - x
    if angle == 180:
        return w - x, h - y
    if angle == 270:
        return w - y, x
    return x, y


def remap_bbox(bbox: list[float], angle: int, w: int, h: int) -> list[float]:
    """
    Map bounding box [x1, y1, x2, y2] from rotated frame back to original image coordinates.
    Maintains [min_x, min_y, max_x, max_y] canonical ordering.
    """
    x1, y1, x2, y2 = bbox
    p1_x, p1_y = remap_point(x1, y1, angle, w, h)
    p2_x, p2_y = remap_point(x2, y2, angle, w, h)
    return [
        min(p1_x, p2_x),
        min(p1_y, p2_y),
        max(p1_x, p2_x),
        max(p1_y, p2_y),
    ]


def inverse_remap_bbox(bbox: list[float], angle: int, w: int, h: int) -> list[float]:
    """
    Map bounding box [x1, y1, x2, y2] from original image dimensions to rotated image frame.
    Useful for drawing detections directly onto the selected orientation image.
    """
    x1, y1, x2, y2 = bbox
    if angle == 0:
        return bbox
    if angle == 90:
        return [h - y2, x1, h - y1, x2]
    if angle == 180:
        return [w - x2, h - y2, w - x1, h - y1]
    if angle == 270:
        return [y1, w - x2, y2, w - x1]
    return bbox


def iou(box1: list[float], box2: list[float]) -> float:
    """Calculate Intersection over Union (IoU) between two bounding boxes [x1, y1, x2, y2]."""
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])

    intersection = max(0.0, x2 - x1) * max(0.0, y2 - y1)
    area1 = max(0.0, box1[2] - box1[0]) * max(0.0, box1[3] - box1[1])
    area2 = max(0.0, box2[2] - box2[0]) * max(0.0, box2[3] - box2[1])
    union = area1 + area2 - intersection

    return (intersection / union) if union > 0.0 else 0.0


def dedup_detections(dets: list[dict[str, Any]], thresh: float = 0.45) -> list[dict[str, Any]]:
    """Deduplicate overlapping bounding boxes, preserving higher-confidence candidates."""
    kept: list[dict[str, Any]] = []
    for d in sorted(dets, key=lambda x: x["confidence"], reverse=True):
        if not any(iou(d["bbox"], k["bbox"]) > thresh for k in kept):
            kept.append(d)
    return sorted(kept, key=lambda x: x["center_x"])


def red_ratio(img_bgr: np.ndarray, bbox: list[float]) -> float:
    """Calculate ratio of red pixels within a bounding box (indicating decimal/fraction digits)."""
    x1, y1, x2, y2 = [int(v) for v in bbox]
    pad_x = max(1, int((x2 - x1) * 0.15))
    pad_y = max(1, int((y2 - y1) * 0.15))

    crop = img_bgr[
        max(0, y1 + pad_y): min(img_bgr.shape[0], y2 - pad_y),
        max(0, x1 + pad_x): min(img_bgr.shape[1], x2 - pad_x),
    ]
    if crop.size == 0:
        return 0.0

    hsv = cv2.cvtColor(crop, cv2.COLOR_BGR2HSV)
    mask1 = cv2.inRange(hsv, (0, 40, 40), (12, 255, 255))
    mask2 = cv2.inRange(hsv, (165, 40, 40), (180, 255, 255))
    red_mask = mask1 | mask2

    return float(np.mean(red_mask > 0))


def draw_annotated_image(
    image_rgb: np.ndarray,
    digits: list[dict[str, Any]],
    best_angle: int = 0,
    prep: str = "orig",
    conf_reliable: float = 0.60,
    show_on_rotated: bool = True,
) -> np.ndarray:
    """
    Render bounding boxes, digit labels, and confidence tags on the image.
    If show_on_rotated is True, renders on the best-detected rotation and contrast view.
    Otherwise renders on the original orientation.
    """
    if image_rgb is None or not image_rgb.size:
        return image_rgb

    h, w = image_rgb.shape[:2]
    bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)

    if show_on_rotated and best_angle != 0:
        canvas_bgr = rotate_image(bgr, best_angle)
        if prep == "clahe":
            canvas_bgr = apply_prep(canvas_bgr, "clahe")
        elif prep == "histeq":
            canvas_bgr = apply_prep(canvas_bgr, "histeq")
    else:
        canvas_bgr = bgr.copy()

    canvas = cv2.cvtColor(canvas_bgr, cv2.COLOR_BGR2RGB)
    cur_h, cur_w = canvas.shape[:2]

    line_thickness = max(2, int(round(min(cur_w, cur_h) / 320)))
    font_scale = max(0.45, min(cur_w, cur_h) / 850.0)

    for d in digits:
        bbox = d.get("bbox", [])
        if len(bbox) != 4:
            continue

        if show_on_rotated and best_angle != 0:
            rx1, ry1, rx2, ry2 = inverse_remap_bbox(bbox, best_angle, w, h)
        else:
            rx1, ry1, rx2, ry2 = bbox

        ix1, iy1 = max(0, int(round(rx1))), max(0, int(round(ry1)))
        ix2, iy2 = min(cur_w - 1, int(round(rx2))), min(cur_h - 1, int(round(ry2)))

        conf = d.get("confidence", 0.0)
        digit_val = d.get("digit", "?")
        is_reliable = d.get("reliable", conf >= conf_reliable)

        # Colors: Green for reliable (>= 0.60), Amber/Orange for warning (< 0.60)
        box_color = (46, 204, 113) if is_reliable else (255, 140, 0)

        # Draw bounding box
        cv2.rectangle(canvas, (ix1, iy1), (ix2, iy2), box_color, line_thickness)

        # Label tag
        label = f"{digit_val} ({conf:.2f})"
        (tw, th), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1)

        tag_top = max(0, iy1 - th - 6)
        tag_bottom = iy1
        tag_right = min(cur_w, ix1 + tw + 8)

        # Label background
        cv2.rectangle(canvas, (ix1, tag_top), (tag_right, tag_bottom), box_color, -1)
        # Text label (black on vibrant tag for clear contrast)
        cv2.putText(
            canvas,
            label,
            (ix1 + 4, tag_bottom - 4),
            cv2.FONT_HERSHEY_SIMPLEX,
            font_scale,
            (0, 0, 0),
            max(1, int(line_thickness / 2)),
            cv2.LINE_AA,
        )

    return canvas