import easyocr
import numpy as np
import cv2
import re
import logging
from datetime import datetime
import os

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Initialize EasyOCR
easyocr_reader = easyocr.Reader(['en'], gpu=False)

# Directory for debug images
DEBUG_DIR = "debug_images"
os.makedirs(DEBUG_DIR, exist_ok=True)

def save_debug_image(img, filename_suffix, prefix=""):
    """Save image to debug directory with timestamp."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    filename = os.path.join(DEBUG_DIR, f"{prefix}{timestamp}_{filename_suffix}.png")
    if len(img.shape) == 3:
        cv2.imwrite(filename, img)
    else:
        cv2.imwrite(filename, img)
    logging.info(f"Saved debug image: {filename}")

def estimate_brightness(img):
    """Estimate image brightness."""
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return np.mean(gray)

def preprocess_image(img):
    """Preprocess image for OCR."""
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    denoised = cv2.bilateralFilter(gray, 5, 8, 8)
    save_debug_image(denoised, "01_preprocess_bilateral")
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    enhanced = clahe.apply(denoised)
    save_debug_image(enhanced, "02_preprocess_clahe")
    return enhanced

def correct_rotation(img):
    """Correct image rotation."""
    try:
        edges = cv2.Canny(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), 50, 150)
        lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=50, minLineLength=30, maxLineGap=10)
        if lines is not None:
            angles = [np.arctan2(line[0][3] - line[0][1], line[0][2] - line[0][0]) * 180 / np.pi for line in lines]
            angle = np.median(angles)
            if abs(angle) > 2:
                h, w = img.shape[:2]
                center = (w // 2, h // 2)
                M = cv2.getRotationMatrix2D(center, angle, 1.0)
                img = cv2.warpAffine(img, M, (w, h))
                save_debug_image(img, "00_rotated_image")
                logging.info(f"Applied rotation: {angle:.2f} degrees")
        return img
    except Exception as e:
        logging.error(f"Rotation correction failed: {str(e)}")
        return img

def detect_roi(img):
    """Detect region of interest (display)."""
    try:
        save_debug_image(img, "03_original")
        preprocessed = preprocess_image(img)
        brightness_map = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        block_size = max(9, min(31, int(img.shape[0] / 25) * 2 + 1))
        thresh = cv2.adaptiveThreshold(preprocessed, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                       cv2.THRESH_BINARY_INV, block_size, 2)
        save_debug_image(thresh, "04_roi_threshold")
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if contours:
            img_area = img.shape[0] * img.shape[1]
            valid_contours = []
            for c in contours:
                area = cv2.contourArea(c)
                x, y, w, h = cv2.boundingRect(c)
                roi_brightness = np.mean(brightness_map[y:y+h, x:x+w])
                aspect_ratio = w / h
                if (50 < area < (img_area * 0.95) and 
                    0.2 <= aspect_ratio <= 30.0 and w > 30 and h > 10 and roi_brightness > 30):
                    valid_contours.append((c, roi_brightness))
                    logging.debug(f"Contour: Area={area}, Aspect={aspect_ratio:.2f}, Brightness={roi_brightness:.2f}")
            
            if valid_contours:
                contour, _ = max(valid_contours, key=lambda x: x[1])
                x, y, w, h = cv2.boundingRect(contour)
                padding = 200
                x, y = max(0, x - padding), max(0, y - padding)
                w, h = min(w + 2 * padding, img.shape[1] - x), min(h + 2 * padding, img.shape[0] - y)
                roi_img = img[y:y+h, x:x+w]
                save_debug_image(roi_img, "05_detected_roi")
                logging.info(f"Detected ROI: ({x}, {y}, {w}, {h})")
                return roi_img, (x, y, w, h)
        
        logging.info("No ROI found, using full image.")
        save_debug_image(img, "05_no_roi_fallback")
        return img, None
    except Exception as e:
        logging.error(f"ROI detection failed: {str(e)}")
        save_debug_image(img, "05_roi_error_fallback")
        return img, None

def detect_segments(digit_img, brightness):
    """Detect seven-segment digits."""
    h, w = digit_img.shape
    if h < 5 or w < 3:
        return None

    segments = {
        'top': (int(w*0.1), int(w*0.9), 0, int(h*0.25)),
        'middle': (int(w*0.1), int(w*0.9), int(h*0.45), int(h*0.55)),
        'bottom': (int(w*0.1), int(w*0.9), int(h*0.75), h),
        'left_top': (0, int(w*0.3), int(h*0.1), int(h*0.5)),
        'left_bottom': (0, int(w*0.3), int(h*0.5), int(h*0.9)),
        'right_top': (int(w*0.7), w, int(h*0.1), int(h*0.5)),
        'right_bottom': (int(w*0.7), w, int(h*0.5), int(h*0.9))
    }

    segment_presence = {}
    for name, (x1, x2, y1, y2) in segments.items():
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(w, x2), min(h, y2)
        region = digit_img[y1:y2, x1:x2]
        if region.size == 0:
            segment_presence[name] = False
            continue
        pixel_count = np.sum(region == 255)
        total_pixels = region.size
        segment_presence[name] = pixel_count / total_pixels > (0.1 if brightness < 80 else 0.25)

    digit_patterns = {
        '0': ('top', 'bottom', 'left_top', 'left_bottom', 'right_top', 'right_bottom'),
        '1': ('right_top', 'right_bottom'),
        '2': ('top', 'middle', 'bottom', 'left_bottom', 'right_top'),
        '3': ('top', 'middle', 'bottom', 'right_top', 'right_bottom'),
        '4': ('middle', 'left_top', 'right_top', 'right_bottom'),
        '5': ('top', 'middle', 'bottom', 'left_top', 'right_bottom'),
        '6': ('top', 'middle', 'bottom', 'left_top', 'left_bottom', 'right_bottom'),
        '7': ('top', 'right_top', 'right_bottom'),
        '8': ('top', 'middle', 'bottom', 'left_top', 'left_bottom', 'right_top', 'right_bottom'),
        '9': ('top', 'middle', 'bottom', 'left_top', 'right_top', 'right_bottom')
    }

    best_match = None
    max_score = -1
    for digit, pattern in digit_patterns.items():
        matches = sum(1 for segment in pattern if segment_presence.get(segment, False))
        non_matches_penalty = sum(1 for segment in segment_presence if segment not in pattern and segment_presence[segment])
        score = matches - 0.1 * non_matches_penalty
        if matches >= len(pattern) * 0.55:
            score += 1.0
        if score > max_score:
            max_score = score
            best_match = digit
    
    logging.debug(f"Segment presence: {segment_presence}, Digit: {best_match}")
    return best_match

def custom_seven_segment_ocr(img, roi_bbox):
    """Perform OCR for seven-segment displays."""
    try:
        preprocessed = preprocess_image(img)
        brightness = estimate_brightness(img)
        _, thresh = cv2.threshold(preprocessed, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        save_debug_image(thresh, "06_roi_thresh_digits")
        results = easyocr_reader.readtext(thresh, detail=1, paragraph=False, 
                                          contrast_ths=0.05, adjust_contrast=1.2, 
                                          text_threshold=0.15, mag_ratio=4.0, 
                                          allowlist='0123456789.', batch_size=2, y_ths=0.3)
        
        logging.info(f"EasyOCR results: {results}")
        if not results:
            logging.info("No digits found.")
            return None

        digits_info = []
        for (bbox, text, conf) in results:
            (x1, y1), (x2, y2), (x3, y3), (x4, y4) = bbox
            h_bbox = max(y1, y2, y3, y4) - min(y1, y2, y3, y4)
            if (text.isdigit() or text == '.') and h_bbox > 4:
                x_min, x_max = int(min(x1, x4)), int(max(x2, x3))
                y_min, y_max = int(min(y1, y2)), int(max(y3, y4))
                digits_info.append((x_min, x_max, y_min, y_max, text, conf))

        digits_info.sort(key=lambda x: x[0])
        recognized_text = ""
        for idx, (x_min, x_max, y_min, y_max, easyocr_char, easyocr_conf) in enumerate(digits_info):
            x_min, y_min = max(0, x_min), max(0, y_min)
            x_max, y_max = min(thresh.shape[1], x_max), min(thresh.shape[0], y_max)
            if x_max <= x_min or y_max <= y_min:
                continue
            digit_img_crop = thresh[y_min:y_max, x_min:x_max]
            save_debug_image(digit_img_crop, f"07_digit_crop_{idx}_{easyocr_char}")
            if easyocr_conf > 0.8 or easyocr_char == '.':
                recognized_text += easyocr_char
            else:
                digit_from_segments = detect_segments(digit_img_crop, brightness)
                recognized_text += digit_from_segments if digit_from_segments else easyocr_char
        
        logging.info(f"Recognized text: {recognized_text}")
        text = re.sub(r"[^\d\.]", "", recognized_text)
        if text.count('.') > 1:
            text = text.replace('.', '', text.count('.') - 1)
        if text and re.fullmatch(r"^\d*\.?\d*$", text):
            text = text.strip('.')
            if text == '':
                return None
            return text.lstrip('0') or '0'
        logging.info(f"Text '{recognized_text}' failed validation.")
        return None
    except Exception as e:
        logging.error(f"Seven-segment OCR failed: {str(e)}")
        return None

def extract_weight_from_image(pil_img):
    """Extract weight from a digital scale image."""
    try:
        img = np.array(pil_img)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        save_debug_image(img, "00_input_image")
        img = correct_rotation(img)
        brightness = estimate_brightness(img)
        conf_threshold = 0.6 if brightness > 150 else (0.4 if brightness > 80 else 0.2)

        roi_img, roi_bbox = detect_roi(img)
        if roi_bbox:
            conf_threshold *= 1.05 if (roi_bbox[2] * roi_bbox[3]) > (img.shape[0] * img.shape[1] * 0.5) else 1.0

        custom_result = custom_seven_segment_ocr(roi_img, roi_bbox)
        if custom_result and custom_result != '0':
            try:
                weight = float(custom_result)
                if 0.00001 <= weight <= 10000:
                    logging.info(f"Custom OCR: {custom_result}, Confidence: 90.0%")
                    return custom_result, 90.0
                logging.warning(f"Custom OCR {custom_result} out of range.")
            except ValueError:
                logging.warning(f"Custom OCR '{custom_result}' invalid number.")

        logging.info("Custom OCR failed, using EasyOCR fallback.")
        preprocessed_roi = preprocess_image(roi_img)
        final_roi = cv2.adaptiveThreshold(preprocessed_roi, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                          cv2.THRESH_BINARY_INV, max(9, min(31, int(roi_img.shape[0] / 25) * 2 + 1)), 2)
        save_debug_image(final_roi, "08_fallback_thresh")

        results = easyocr_reader.readtext(final_roi, detail=1, paragraph=False, 
                                          contrast_ths=0.05, adjust_contrast=1.2, 
                                          text_threshold=0.15, mag_ratio=4.0, 
                                          allowlist='0123456789. kglb', batch_size=2, y_ths=0.3)

        if not results:
            logging.info("First EasyOCR pass failed, trying fallback.")
            results = easyocr_reader.readtext(final_roi, detail=1, paragraph=False, 
                                              contrast_ths=0.02, adjust_contrast=1.5, 
                                              text_threshold=0.1, mag_ratio=5.0, 
                                              allowlist='0123456789. kglb', batch_size=2, y_ths=0.3)
            save_debug_image(final_roi, "08_fallback_thresh_fallback")

        logging.info(f"EasyOCR results: {results}")
        candidates = []
        unit = None
        for (bbox, text, conf) in results:
            if 'kg' in text.lower():
                unit = 'kg'
                continue
            elif 'g' in text.lower():
                unit = 'g'
                continue
            elif 'lb' in text.lower():
                unit = 'lb'
                continue
            text = re.sub(r"[^\d\.]", "", text)
            if text.count('.') > 1:
                text = text.replace('.', '', text.count('.') - 1)
            text = text.strip('.')
            if re.fullmatch(r"^\d*\.?\d*$", text):
                try:
                    weight = float(text)
                    if unit == 'g':
                        weight /= 1000
                    elif unit == 'lb':
                        weight *= 0.453592
                    range_score = 1.5 if 0.00001 <= weight <= 10000 else 0.5
                    digit_count = len(text.replace('.', ''))
                    digit_score = 1.4 if 1 <= digit_count <= 8 else 0.6
                    score = conf * range_score * digit_score
                    if roi_bbox:
                        x_roi, y_roi, w_roi, h_roi = roi_bbox
                        roi_area = w_roi * h_roi
                        x_min, y_min = int(min(b[0] for b in bbox)), int(min(b[1] for b in bbox))
                        x_max, y_max = int(max(b[0] for b in bbox)), int(max(b[1] for b in bbox))
                        bbox_area = (x_max - x_min) * (y_max - y_min)
                        if roi_area > 0 and bbox_area / roi_area < 0.02:
                            score *= 0.4
                    candidates.append((text, conf, score, unit))
                    logging.info(f"Candidate: '{text}', Unit: {unit or 'none'}, Conf: {conf}, Score: {score}")
                except ValueError:
                    logging.warning(f"Could not convert '{text}' to float.")

        if not candidates and not roi_bbox:
            logging.info("No candidates, trying full image.")
            preprocessed_full = preprocess_image(img)
            final_full = cv2.adaptiveThreshold(preprocessed_full, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                               cv2.THRESH_BINARY_INV, max(9, min(31, int(img.shape[0] / 25) * 2 + 1)), 2)
            save_debug_image(final_full, "08_fallback_full")
            results = easyocr_reader.readtext(final_full, detail=1, paragraph=False, 
                                              contrast_ths=0.05, adjust_contrast=1.5, 
                                              text_threshold=0.15, mag_ratio=4.0, 
                                              allowlist='0123456789. kglb', batch_size=2, y_ths=0.3)
            logging.info(f"Full image EasyOCR: {results}")
            for (bbox, text, conf) in results:
                if 'kg' in text.lower():
                    unit = 'kg'
                    continue
                elif 'g' in text.lower():
                    unit = 'g'
                    continue
                elif 'lb' in text.lower():
                    unit = 'lb'
                    continue
                text = re.sub(r"[^\d\.]", "", text)
                if text.count('.') > 1:
                    text = text.replace('.', '', text.count('.') - 1)
                text = text.strip('.')
                if re.fullmatch(r"^\d*\.?\d*$", text):
                    try:
                        weight = float(text)
                        if unit == 'g':
                            weight /= 1000
                        elif unit == 'lb':
                            weight *= 0.453592
                        range_score = 1.2 if 0.00001 <= weight <= 10000 else 0.4
                        digit_count = len(text.replace('.', ''))
                        digit_score = 1.2 if 1 <= digit_count <= 8 else 0.5
                        score = conf * range_score * digit_score * 0.7
                        candidates.append((text, conf, score, unit))
                        logging.info(f"Full image candidate: '{text}', Unit: {unit or 'none'}, Conf: {conf}, Score: {score}")
                    except ValueError:
                        logging.warning(f"Could not convert '{text}' to float (full image).")

        if not candidates:
            logging.info("No valid weight detected.")
            return "Not detected", 0.0

        best_weight, best_conf, best_score, best_unit = max(candidates, key=lambda x: x[2])
        if "." in best_weight:
            int_part, dec_part = best_weight.split(".")
            int_part = int_part.lstrip("0") or "0"
            dec_part = dec_part.rstrip('0')
            best_weight = f"{int_part}.{dec_part}" if dec_part else int_part
        else:
            best_weight = best_weight.lstrip('0') or "0"

        try:
            final_weight = float(best_weight)
            if final_weight < 0.00001 or final_weight > 10000:
                best_conf *= 0.4
            elif final_weight == 0 and best_conf < 0.95:
                best_conf *= 0.5
        except ValueError:
            pass

        logging.info(f"Final weight: {best_weight} kg, Confidence: {round(best_conf * 100, 2)}%, Unit: {best_unit or 'none'}")
        return best_weight, round(best_conf * 100, 2)
    except Exception as e:
        logging.error(f"Weight extraction failed: {str(e)}")
        return "Not detected", 0.0