import os
import streamlit as st
from PIL import Image
import numpy as np
from transformers import AutoProcessor, AutoModelForPreTraining
import torch
from huggingface_hub import login

# Get the token from environment variables (set in Hugging Face Space)
token = os.getenv("HF_Token")
login(token)

# Initialize the Hugging Face model
processor = AutoProcessor.from_pretrained("google/paligemma-3b-mix-224")
model = AutoModelForPreTraining.from_pretrained("google/paligemma-3b-mix-224")

# Function to transcribe handwritten notes using Hugging Face model
def transcribe_handwriting(image):
    # Convert image to RGB and numpy array
    image = image.convert("RGB")
    image_array = np.array(image)

    # Prepare input for the model
    inputs = processor(image_array, return_tensors="pt")

    # Generate output
    with torch.no_grad():
        outputs = model.generate(**inputs, max_length=512)

    transcription = processor.decode(outputs[0], skip_special_tokens=True)
    return transcription

# Set Streamlit page configuration
st.set_page_config(
    page_title="AI and ML Handwriting Reader",
    page_icon="💊",
    layout="centered",
)

# Function to display extracted text in a fancy div
def display_extracted_text(extracted_text):
    # Custom HTML and CSS for styling
    fancy_div = f"""
    <div style="
        border: 2px solid #128c7E; 
        border-radius: 10px; 
        padding: 20px; 
        background-color: #f0f0f0; 
        box-shadow: 0 4px 10px rgba(0, 0, 0, 0.1); 
        font-family: 'Arial', sans-serif; 
        color: #333;">
        <h2 style="color: #128c7E;">Transcribed Text from Image:</h2>
        <p style="font-size: 16px;">{extracted_text}</p>
    </div>
    """
    st.markdown(fancy_div, unsafe_allow_html=True)
    
# Header
st.title("AI and ML based Handwriting Reader")

# Upload prescription image
uploaded_file = st.file_uploader("Upload Handwriting Image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Open the image using PIL
    image = Image.open(uploaded_file)

    # Display uploaded image
    st.image(image, caption="Uploaded Handwriting Image", width=400)

    with st.spinner("Transcribing handwriting..."):
        # Transcribe handwritten notes
        extracted_text = transcribe_handwriting(image)
        display_extracted_text(extracted_text)