import gradio as gr
from transformers import pipeline
from PIL import Image
import json
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
import easyocr
import re

# Initialize EasyOCR reader
reader = easyocr.Reader(["en"], gpu=False)

# Load BioGPT model for recommendations
bio_gpt = pipeline("text-generation", model="microsoft/BioGPT")

# Load reference ranges from dataset.json
def load_reference_ranges(file_path="dataset.json"):
    with open(file_path, "r") as file:
        reference_ranges = json.load(file)
    return reference_ranges

reference_ranges = load_reference_ranges()

# Parameter name mapping
parameter_mapping = {
    "hgb": "hemoglobin",
    "hemoglobin": "hemoglobin",
    "rbc": "rbc",
    "wbc": "wbc",
    "plt": "platelet",
    "platelets": "platelet",
    "hematocrit": "hematocrit",
    "mcv": "mcv",
    "mch": "mch",
    "mchc": "mchc",
    # Add more mappings as needed
}

# Extract text from uploaded image using EasyOCR
def extract_text_from_image(image_path):
    try:
        # Read text from image
        result = reader.readtext(image_path, detail=0)  # Extract only the text
        extracted_text = " ".join(result)
        print("Extracted Text:", extracted_text)  # Debugging
        return extracted_text
    except Exception as e:
        return f"Error extracting text: {e}"

# Extract value using regex
def extract_value(text, param):
    match = re.search(rf"{param}[:\s]*([\d.]+)", text, re.IGNORECASE)
    if match:
        return float(match.group(1))
    return None

# Analyze extracted text and compare against reference ranges
def analyze_blood_report(text):
    abnormalities = []
    analysis = "Blood Test Analysis Results:\n\n"

    for key, standard_param in parameter_mapping.items():
        if key in text.lower():
            ranges = reference_ranges[standard_param]
            value = extract_value(text, key)
            if value is not None:
                if value < ranges["low"]:
                    abnormalities.append(f"{standard_param.capitalize()} is LOW ({value} {ranges['unit']}).")
                elif value > ranges["high"]:
                    abnormalities.append(f"{standard_param.capitalize()} is HIGH ({value} {ranges['unit']}).")
                else:
                    analysis += f"{standard_param.capitalize()} is NORMAL ({value} {ranges['unit']}).\n"
            else:
                analysis += f"{standard_param.capitalize()} could not be analyzed.\n"

    if abnormalities:
        analysis += "\nAbnormalities Detected:\n" + "\n".join(abnormalities) + "\n"
    else:
        analysis += "\nNo abnormalities detected.\n"

    return analysis, abnormalities

# Generate recommendations using BioGPT
def get_recommendations(abnormalities):
    if not abnormalities:
        return "No recommendations needed."
    query = " ".join(abnormalities) + " Provide medical recommendations."
    recommendations = bio_gpt(query, max_length=100, num_return_sequences=1)[0]["generated_text"]
    return recommendations

# Create a PDF report
def create_pdf_report(content, output_path="blood_test_report.pdf"):
    c = canvas.Canvas(output_path, pagesize=letter)
    c.drawString(100, 750, "Blood Test 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

# Main function to process blood test image
def process_blood_test(image_path):
    # Step 1: Extract text
    extracted_text = extract_text_from_image(image_path)
    if "Error" in extracted_text:
        return extracted_text, None

    # Step 2: Analyze extracted text
    analysis, abnormalities = analyze_blood_report(extracted_text)

    # Step 3: Generate recommendations
    recommendations = get_recommendations(abnormalities)

    # Step 4: Combine results and create PDF
    full_report = analysis + "\nRecommendations:\n" + recommendations
    pdf_path = create_pdf_report(full_report)

    return full_report, pdf_path

# Gradio Interface
interface = gr.Interface(
    fn=process_blood_test,
    inputs=gr.Image(type="filepath", label="Upload Blood Test Report Image (PNG, JPG, JPEG)"),
    outputs=[
        gr.Textbox(label="Analysis and Recommendations"),
        gr.File(label="Download PDF Report"),
    ],
    title="AI Blood Test Analyzer",
    description=(
        "Upload a blood test report image (PNG, JPG, JPEG), and the app will analyze the values, flag abnormalities, "
        "and provide recommendations using BioGPT. You can also download a PDF report."
    ),
    theme="compact",
    css="""
        body {
            font-family: 'Arial', sans-serif;
            background-color: #f9f9f9;
        }
        .gradio-container {
            color: #333;
            max-width: 800px;
            margin: 0 auto;
        }
        .gradio-container .label {
            font-weight: bold;
            font-size: 18px;
        }
        .gradio-container .output {
            background-color: #eef;
            padding: 10px;
            border-radius: 10px;
        }
    """,
    allow_flagging="never"
)

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