import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
from PIL import Image
from PyPDF2 import PdfReader
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas

# Load OCR model for extracting text from images
ocr_model = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")

# Load medical AI model (BioGPT or similar) for prescription validation
medical_model_name = "microsoft/BioGPT"
medical_tokenizer = AutoTokenizer.from_pretrained(medical_model_name)
medical_model = AutoModelForSequenceClassification.from_pretrained(medical_model_name)

# Function to extract text from images or PDFs
def extract_text(file_path):
    try:
        # Handle PDFs
        if file_path.endswith(".pdf"):
            reader = PdfReader(file_path)
            text = "".join(page.extract_text() for page in reader.pages)
            return text.strip()
        # Handle images
        result = ocr_model(Image.open(file_path))
        return result[0]['generated_text']
    except Exception as e:
        return f"Error processing the file: {e}"

# Function to validate prescription using the medical model
def validate_prescription_with_model(extracted_text):
    # Tokenize and process with the AI model
    inputs = medical_tokenizer(extracted_text, return_tensors="pt", truncation=True, padding=True)
    outputs = medical_model(**inputs)
    logits = outputs.logits
    predictions = logits.softmax(dim=-1)

    # Generate model-driven validation results
    validation_report = "🔍 Prescription Validation Results:\n"
    for i, score in enumerate(predictions[0]):
        token = medical_tokenizer.decode([i]).strip()
        if token not in ["[PAD]", "[unused1]"]:  # Ignore invalid tokens
            validation_report += f"- {token}: {score.item():.2f}\n"
    return validation_report

# Main function to handle prescription analysis
def analyze_prescription(file):
    try:
        # Step 1: Extract text
        extracted_text = extract_text(file)
        if not extracted_text.strip():
            return "No readable text found in the uploaded file.", None

        # Step 2: Validate prescription using AI model
        validation_report = validate_prescription_with_model(extracted_text)

        # Combine the extracted text and validation results
        full_report = f"Extracted Text:\n{extracted_text}\n\n{validation_report}"

        # Step 3: Generate a PDF report
        output_pdf = "prescription_validation_report.pdf"
        create_pdf_report(full_report, output_pdf)

        return full_report, output_pdf
    except Exception as e:
        return f"Error processing file: {e}", None

# Function to create a PDF report
def create_pdf_report(content, output_path):
    c = canvas.Canvas(output_path, pagesize=letter)
    c.drawString(100, 750, "Prescription Validation Report")
    c.drawString(100, 730, "------------------------------")
    y_position = 700
    for line in content.split("\n"):
        c.drawString(100, y_position, line)
        y_position -= 20
    c.save()
    return output_path

# Gradio interface
interface = gr.Interface(
    fn=analyze_prescription,
    inputs=gr.File(label="Upload your Prescription (PNG, JPG, JPEG, or PDF)"),
    outputs=[
        gr.Textbox(label="Validation Results"),
        gr.File(label="Download PDF Report")
    ],
    title="AI-Powered Prescription Validator",
    description=(
        "Upload your medical prescription in image (PNG, JPG, JPEG) or PDF format. "
        "The app will extract the text, analyze it using advanced AI models, and validate the prescription. "
        "Download a comprehensive PDF report of the validation results."
    ),
    allow_flagging="never"
)

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