import streamlit as st
import easyocr
from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
from PIL import Image
import numpy as np
import cv2 

# Initialize EasyOCR for text extraction
reader = easyocr.Reader(['en'])

# Load ClinicalBERT for NER (Named Entity Recognition)
model_name = "emilyalsentzer/Bio_ClinicalBERT"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name)
nlp = pipeline("ner", model=model, tokenizer=tokenizer)

# Mock function for validating test results
def validate_medical_value(test_name, test_value):
    reference_ranges = {
        "Blood Pressure": (90, 120),  # Example range: systolic
        "Glucose": (70, 100),  # Normal range for glucose in mg/dL
    }
    
    if test_name in reference_ranges:
        min_range, max_range = reference_ranges[test_name]
        if test_value < min_range:
            return f"{test_name} is low."
        elif test_value > max_range:
            return f"{test_name} is high."
        else:
            return f"{test_name} is within the normal range."
    else:
        return f"{test_name} range is unknown."

# Function to extract text from image
def extract_text_from_image(image):
    result = reader.readtext(image)
    text = ' '.join([item[1] for item in result])
    return text

# Function to check image clarity (basic approach using edge detection)
def check_image_clarity(image):
    # Convert image to grayscale
    gray_image = image.convert("L")
    image_np = np.array(gray_image)
    
    # Simple edge detection (Laplacian method)
    laplacian_var = cv2.Laplacian(image_np, cv2.CV_64F).var()
    
    # Threshold for determining clarity
    return laplacian_var > 100  # This threshold can be adjusted

# Function to extract medical entities (test results) from text
def extract_medical_values(text):
    entities = nlp(text)
    medical_data = {}
    
    for entity in entities:
        if entity['entity_group'] in ['LAB_RESULT', 'DISEASE', 'MEDICATION']:  # Modify based on actual labels
            medical_data[entity['word']] = entity['score']
    
    return medical_data

# Function to analyze medical report
def analyze_report(image):
    # Step 1: Check image clarity
    if not check_image_clarity(image):
        return "Image is unclear, please upload a clearer image."
    
    # Step 2: Extract text from the medical report image
    text = extract_text_from_image(image)
    
    # Step 3: Extract medical values using ClinicalBERT
    medical_values = extract_medical_values(text)
    
    analysis_results = []
    
    for test_name, test_value in medical_values.items():
        # Step 4: Validate test value against normal range
        validation_result = validate_medical_value(test_name, test_value)
        analysis_results.append(validation_result)
    
    return analysis_results

# Streamlit UI
st.title("Medical Report Analyzer")
st.write("Upload your medical report image to analyze the test results.")

# Upload the image
uploaded_image = st.file_uploader("Upload Image", type=["jpg", "png", "jpeg"])

if uploaded_image is not None:
    # Open the uploaded image
    image = Image.open(uploaded_image)
    
    # Show image for verification
    st.image(image, caption="Uploaded Medical Report", use_column_width=True)
    
    # Analyze the report
    analysis_results = analyze_report(image)
    
    # Display results
    if isinstance(analysis_results, list):
        for result in analysis_results:
            st.write(result)
    else:
        st.write(analysis_results)