import gradio as gr
import pytesseract
from PIL import Image
import re


def process_meter_image(image):
    # Convert to PIL Image if needed
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # Extract text using OCR
    text = pytesseract.image_to_string(image)
    
    # Parse meter information
    info = {
        'brand': '',
        'model': '',
        'type': '',
        'reading': '',
        'voltage': '',
        'current': '',
        'frequency': '',
        'calibration': '',
        'property': '',
        'year': '',
        'serial': ''
    }
    
    # Extract information using regex patterns
    patterns = {
        'brand': r'HEXING',
        'model': r'HXE\d+',
        'type': r'\dPhase \dWire Meter',
        'reading': r'(\d+)\s*kWh',
        'voltage': r'(\d+)V',
        'frequency': r'(\d+)Hz',
        'calibration': r'(\d+)imp/kWh',
        'property': r'Property of ([^\n]+)',
        'year': r'20\d{2}',
        'serial': r'No\.(\d+)'
    }
    
    for key, pattern in patterns.items():
        match = re.search(pattern, text)
        if match:
            info[key] = match.group(1) if len(match.groups()) > 0 else match.group(0)
    
    # Format output
    output = f"""
    Meter Information:
    ------------------
    Brand: {info['brand']}
    Model: {info['model']}
    Type: {info['type']}
    Current Reading: {info['reading']}
    Voltage Rating: {info['voltage']}
    Frequency: {info['frequency']}
    Calibration: {info['calibration']}
    Property of: {info['property']}
    Manufacturing Year: {info['year']}
    Serial Number: {info['serial']}
    """
    
    return output

# Create Gradio interface
iface = gr.Interface(
    fn=process_meter_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(label="Extracted Information"),
    title="Electricity Meter OCR",
    description="Upload an image of an electricity meter to extract information",
    examples=[["meter_image.jpg"]]
)

# Launch the app
iface.launch()