import gradio as gr
import cv2
import pytesseract
from PIL import Image
import io
import base64
from datetime import datetime
import pytz
import numpy as np
import logging

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

# Configure Tesseract path (ensure it's correctly set to your Tesseract installation)
try:
    pytesseract.pytesseract.tesseract_cmd = '/usr/bin/tesseract'  # Change this to your tesseract path
    pytesseract.get_tesseract_version()  # Test Tesseract availability
    logging.info("Tesseract is available")
except Exception as e:
    logging.error(f"Tesseract not found or misconfigured: {str(e)}")

# Image Preprocessing to improve OCR accuracy
def preprocess_image(img_cv):
    """Preprocess image for OCR: enhance contrast, reduce noise, and apply adaptive thresholding."""
    try:
        # Convert to grayscale
        gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
        # Enhance contrast using CLAHE (Contrast Limited Adaptive Histogram Equalization)
        clahe = cv2.createCLAHE(clipLimit=5.0, tileGridSize=(8, 8))
        contrast = clahe.apply(gray)
        # Reduce noise with Gaussian blur
        blurred = cv2.GaussianBlur(contrast, (5, 5), 0)
        # Apply adaptive thresholding for better binary image representation
        thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 11, 2)
        # Sharpen the image to enhance details
        sharpened = cv2.filter2D(thresh, -1, np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]]))
        return sharpened
    except Exception as e:
        logging.error(f"Image preprocessing failed: {str(e)}")
        return img_cv

# Function to extract weight from the image using Tesseract OCR
def extract_weight(img):
    """Extract weight from image using Tesseract OCR."""
    try:
        if img is None:
            logging.error("No image provided for OCR")
            return "Not detected", 0.0, None
        
        # Convert PIL image to OpenCV format for processing
        img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
        # Preprocess the image (contrast, noise reduction, etc.)
        processed_img = preprocess_image(img_cv)
        
        # OCR configuration to focus on digits and decimals
        custom_config = r'--oem 3 --psm 6 -c tessedit_char_whitelist=0123456789.'
        
        # Run OCR on the processed image
        text = pytesseract.image_to_string(processed_img, config=custom_config)
        logging.info(f"OCR result: '{text}'")
        
        # Extract valid weight (only digits and decimals)
        weight = ''.join(filter(lambda x: x in '0123456789.', text.strip()))
        if weight:
            try:
                weight_float = float(weight)
                if weight_float >= 0:  # Ensure valid weight value
                    confidence = 95.0  # High confidence if weight is valid
                    logging.info(f"Weight detected: {weight} (Confidence: {confidence:.2f}%)")
                    return weight, confidence, processed_img
            except ValueError:
                logging.warning(f"Invalid number format: {weight}")
        
        logging.error("OCR failed to detect a valid weight")
        return "Not detected", 0.0, None
    except Exception as e:
        logging.error(f"OCR processing failed: {str(e)}")
        return "Not detected", 0.0, None

# Main function to process the image and return results
def process_image(img):
    """Process uploaded or captured image and extract weight."""
    if img is None:
        logging.error("No image provided")
        return "No image uploaded", None, gr.update(visible=False), gr.update(visible=False)
    
    # Get the current time in IST format
    ist_time = datetime.now(pytz.timezone("Asia/Kolkata")).strftime("%d-%m-%Y %I:%M:%S %p")
    
    # Extract weight and confidence from the image
    weight, confidence, processed_img = extract_weight(img)
    
    # If no weight detected, display the failure message
    if weight == "Not detected" or confidence < 95.0:
        logging.warning(f"Weight detection failed: {weight} (Confidence: {confidence:.2f}%)")
        return f"{weight} (Confidence: {confidence:.2f}%)", ist_time, gr.update(visible=True), gr.update(visible=False)
    
    # Convert processed image to base64 for displaying it as a snapshot
    pil_image = Image.fromarray(processed_img)
    buffered = io.BytesIO()
    pil_image.save(buffered, format="PNG")
    img_base64 = base64.b64encode(buffered.getvalue()).decode()
    
    return f"{weight} kg (Confidence: {confidence:.2f}%)", ist_time, img_base64, gr.update(visible=True)

# Gradio Interface for user input and displaying results
with gr.Blocks(title="⚖️ Auto Weight Logger") as demo:
    gr.Markdown("## ⚖️ Auto Weight Logger")
    gr.Markdown("📷 Upload or capture an image of a digital weight scale (max 5MB).")

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload / Capture Image", sources=["upload", "webcam"])
        output_weight = gr.Textbox(label="⚖️ Detected Weight (in kg)")

    with gr.Row():
        timestamp = gr.Textbox(label="🕒 Captured At (IST)")
        snapshot = gr.Image(label="📸 Snapshot Image", type="pil")

    submit = gr.Button("🔍 Detect Weight")
    submit.click(
        fn=process_image,
        inputs=image_input,
        outputs=[output_weight, timestamp, snapshot]
    )

    gr.Markdown("""
    ### Instructions
    - Upload a clear, well-lit image of a digital weight scale display (7-segment font preferred).
    - Ensure the image is < 5MB (automatically resized if larger).
    - Review the detected weight and try again if it's incorrect.
    """)

if __name__ == "__main__":
    demo.launch()