import re
import cv2
import pytesseract
import pandas as pd
import requests
from io import BytesIO
from PIL import Image
import numpy as np

pytesseract.pytesseract.tesseract_cmd = '/opt/homebrew/bin/tesseract'

entity_unit_map = {
    'width': {'centimetre', 'cm', 'foot', 'ft', 'inch', 'in', 'metre', 'm', 'millimetre', 'mm', 'yard', 'yd'},
    'depth': {'centimetre', 'cm', 'foot', 'ft', 'inch', 'in', 'metre', 'm', 'millimetre', 'mm', 'yard', 'yd'},
    'height': {'centimetre', 'cm', 'foot', 'ft', 'inch', 'in', 'metre', 'm', 'millimetre', 'mm', 'yard', 'yd'},
    'item_weight': {'gram', 'g', 'kilogram', 'kg', 'microgram', 'µg', 'milligram', 'mg', 'ounce', 'oz', 'pound', 'lb', 'ton', 't'},
    'maximum_weight_recommendation': {'gram', 'g', 'kilogram', 'kg', 'microgram', 'µg', 'milligram', 'mg', 'ounce', 'oz', 'pound', 'lb', 'ton', 't'},
    'voltage': {'kilovolt', 'kV', 'millivolt', 'mV', 'volt', 'V'},
    'wattage': {'kilowatt', 'kW', 'watt', 'W'},
    'item_volume': {'centilitre', 'cL', 'cubic foot', 'cubic inch', 'cup', 'decilitre', 'dL', 'fluid ounce', 'fl oz', 'gallon', 'gal', 'imperial gallon', 'litre', 'L', 'microlitre', 'μL', 'millilitre', 'mL', 'pint', 'pt', 'quart', 'qt'}

}

abbreviation_map = {
    'cm': 'centimetre', 'm': 'metre', 'mm': 'millimetre', 'ft': 'foot', 'in': 'inch', 'yd': 'yard',
    'kg': 'kilogram', 'g': 'gram', 'mg': 'milligram', 'µg': 'microgram', 'oz': 'ounce', 'lb': 'pound', 't': 'ton',
    'kV': 'kilovolt', 'mV': 'millivolt', 'V': 'volt',
    'kW': 'kilowatt', 'W': 'watt',
    'L': 'litre', 'mL': 'millilitre', 'cL': 'centilitre', 'dL': 'decilitre', 'μL': 'microlitre',
    'gal': 'gallon', 'fl oz': 'fluid ounce', 'pt': 'pint', 'qt': 'quart', 'cup': 'cup',
    'cubic foot': 'cubic foot', 'cubic inch': 'cubic inch', 'imperial gallon': 'imperial gallon'
}

entity_keywords = {
    'width': ['width', 'breadth','wide'],
    'depth': ['depth'],
    'height': ['height','lenght','len'],
    'item_weight': ['weight','wt'],
    'maximum_weight_recommendation': ['maximum weight', 'recommended weight','max weight','max wt','maximum wt'],
    'voltage': ['voltage'],
    'wattage': ['wattage'],
    'item_volume': ['capacity', 'volume']
}


def extract_numerical_values(text, entity):
    units = entity_unit_map[entity]
    pattern = r'(\d+(\.\d+)?)\s*(' + '|'.join(units) + r')'
    matches = re.findall(pattern, text, re.IGNORECASE)

    if matches:
        standardized_matches = [
            (float(match[0]), abbreviation_map.get(match[2].lower(), abbreviation_map.get(match[2], match[2]))) for match in matches
        ]
        return standardized_matches
    return 'void'

def extract_dimensions_pattern(text):
    pattern = r'(\d+)x(\d+)x(\d+)\s*(mm|cm|in|ft|m)'
    match = re.search(pattern, text, re.IGNORECASE)

    if match:
        length_value = float(match.group(1))
        width_value = float(match.group(2))
        height_value = float(match.group(3))
        unit = abbreviation_map.get(match.group(4).lower(), match.group(4))

        return {
            'length': (length_value, unit),
            'width': (width_value, unit),
            'height': (height_value, unit)
        }
    return {}

def extract_entity_values(text, query_entities):
    entity_matches = {}

    for entity in query_entities:
        entity_matches[entity] = 'void'

    dimension_matches = extract_dimensions_pattern(text)
    if dimension_matches:
        for entity in query_entities:
            if entity in dimension_matches:
                entity_matches[entity] = dimension_matches[entity]

    for entity, keywords in entity_keywords.items():
        if entity in query_entities and entity_matches[entity] == 'void':
            for keyword in keywords:
                if keyword.lower() in text.lower():
                    entity_matches[entity] = extract_numerical_values(text, entity)

    for entity in query_entities:
        if entity_matches[entity] == 'void':
            unit_matches = extract_numerical_values(text, entity)
            if unit_matches != 'void':
                entity_matches[entity] = unit_matches

    return entity_matches

def ocr_extract_entities(image_path, query_entities):
    response = requests.get(image_path)
    img = Image.open(BytesIO(response.content))
    img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    ret, thresh1 = cv2.threshold(gray, 0, 255, cv2.THRESH_OTSU | cv2.THRESH_BINARY_INV)

    rect_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (18, 18))
    dilation = cv2.dilate(thresh1, rect_kernel, iterations=1)

    contours, hierarchy = cv2.findContours(dilation, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    im2 = img.copy()

    extracted_text = ""

    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        cropped = im2[y:y + h, x:x + w]
        text = pytesseract.image_to_string(cropped)
        extracted_text += text + "\n"

    extracted_entities = extract_entity_values(extracted_text, query_entities)

    return extracted_entities

input_file = '/Users/kartik/Downloads/student_resource 3/test2.csv'
output_file = 'output12345.csv'

df = pd.read_csv(input_file)

results = []

partition_size = 10000
partition_index = 1
partition_results = []

for index, row in df.iterrows():
    image_link = row['image_link']
    entity_name = row['entity_name']

    query_entities = [entity_name]
    extracted_entities = ocr_extract_entities(image_link, query_entities)

    entity_value = extracted_entities.get(entity_name, 'void')
    
    # Check if entity_value is a list or float
    if entity_value != 'void' and entity_value != []:
        if isinstance(entity_value, list):
            numerical_value = f"{entity_value[0][0]:,.2f}"
            unit = entity_value[0][1].lower()
            prediction = f"{numerical_value} {unit}"
        elif isinstance(entity_value, float) or isinstance(entity_value, int):
            # Handle case when entity_value is a float or int, not a list
            prediction = f"{entity_value:,.2f}"
        else:
            prediction = ''
    else:
        prediction = ''

    results.append({
        'index': index,
        'prediction': prediction
    })
    partition_results.append({
        'index': index,
        'prediction': prediction
    })

    
    try:
        if (index + 1) % 50 == 0:
            print(f"Parsed {index + 1} rows")
    except Exception as e:
        pass

    if (index + 1) % partition_size == 0:
      partition_df = pd.DataFrame(partition_results)
      partition_output_file = f'output_partition_{partition_index}.csv'
      partition_df.to_csv(partition_output_file, index=False)
      partition_index += 1
      partition_results = []

# Save the final partition if there are remaining rows
if partition_results:
    partition_df = pd.DataFrame(partition_results)
    partition_output_file = f'output_partition_{partition_index}.csv'
    partition_df.to_csv(partition_output_file, index=False)

output_df = pd.DataFrame(results)
output_df.to_csv(output_file, index=False)

i made this function and i am passing a csv with around 1,00,000 entries 
it generates output for 100 rows in a min as per this it will require 16 hours
how do i increase it's speed so that it does this in like 30-20 mins 
i have a m1 mac