import cv2
import pytesseract
import nltk
from nltk.tokenize import RegexpTokenizer
from nltk.corpus import stopwords
import string
from transformers import pipeline

# Path to your Tesseract executable (for Raspberry Pi, typically /usr/bin/tesseract)
pytesseract.pytesseract.tesseract_cmd = r'/usr/bin/tesseract'  # Update if needed

# Set up translation pipeline (using a smaller model for efficiency)
translation_pipeline = pipeline("translation", model="Helsinki-NLP/opus-mt-en-de", device=0)

# Open the webcam (camera 0 is usually the default camera)
cap = cv2.VideoCapture(0)

if not cap.isOpened():
    print("Error: Could not access the camera.")
    exit()

# Ensure that NLTK stopwords are downloaded
nltk.download('stopwords')

def translate_text(text):
    """Translate text to another language using Hugging Face's translation pipeline."""
    translated = translation_pipeline(text, max_length=400)
    return translated[0]['translation_text']

def preprocess_text(text):
    """Preprocess text by removing stopwords and punctuation."""
    tokenizer = RegexpTokenizer(r'\w+')
    stop_words = set(stopwords.words('english'))

    # Tokenize the input text
    tokens = tokenizer.tokenize(text.lower())  # Convert to lowercase and tokenize

    # Remove stopwords and punctuation
    tokens = [word for word in tokens if word not in stop_words and word not in string.punctuation]

    # Join the remaining words back into a single string
    return " ".join(tokens)

# Set the desired frame rate (lowering frame rate to reduce processing load)
frame_rate = 5  # Process every 5th frame for example
frame_count = 0

while True:
    ret, frame = cap.read()

    if not ret:
        print("Failed to grab frame!")
        break

    frame_count += 1
    if frame_count % frame_rate != 0:
        continue  # Skip processing for most frames to reduce load

    # Resize the frame for faster processing (lower resolution)
    frame = cv2.resize(frame, (640, 480))  # Adjust resolution as needed

    # Convert the frame to grayscale
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # Apply some preprocessing (thresholding)
    _, threshold_image = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)

    # Use Tesseract to get the bounding boxes of text
    boxes = pytesseract.image_to_boxes(threshold_image)

    # Draw bounding boxes around the recognized text
    h, w, _ = frame.shape  # Get the frame dimensions
    for b in boxes.splitlines():
        b = b.split()
        x, y, x2, y2 = int(b[1]), int(b[2]), int(b[3]), int(b[4])  # Coordinates of the box
        cv2.rectangle(frame, (x, h - y), (x2, h - y2), (0, 255, 0), 2)  # Draw the box in green

    # Use Tesseract to recognize text from the processed frame
    text = pytesseract.image_to_string(threshold_image)
    print(f"Recognized Text: {text}")

    # Process the recognized text to remove stopwords
    cleaned_text = preprocess_text(text)

    # Translate the text to German (or another language)
    translated_text = translate_text(cleaned_text)

    # Display the cleaned text on the console
    print(f"Cleaned Text: {cleaned_text}")
    print(f"Translated Text: {translated_text}")

    # Optionally, save the cleaned and translated text to a file
    with open("detectedtext_translated.txt", "w") as file:
        file.write(f"Recognized Text: {cleaned_text}\n")
        file.write(f"Translated Text: {translated_text}\n")

    # Optionally, display the cleaned text and translated text on the video frame
    cv2.putText(frame, cleaned_text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
    cv2.putText(frame, translated_text, (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

    # Show the processed video frame with bounding boxes
    cv2.imshow("Real-time Text Recognition and Translation", frame)

    # Break the loop if the user presses the 'q' key
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# Release the webcam and close all OpenCV windows
cap.release()
cv2.destroyAllWindows()