import streamlit as st
from PIL import Image
import pytesseract
import pandas as pd
import re
import openai


# Setup OpenAI API key (replace with your OpenAI API key)
openai.api_key = "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"


# Function to extract text from the image
def extract_text(image):
    """
    Extract text from the image using Tesseract.
    """
    return pytesseract.image_to_string(image)


# Function to parse and clean the extracted text
def clean_and_parse_extracted_text(raw_text):
    """
    Parse and clean the raw text to extract structured data.
    """
    lines = raw_text.split("\n")
    lines = [line.strip() for line in lines if line.strip()]

    data = []
    for line in lines:
        match = re.match(
            r"^(.*?)(\d+(\.\d+)?)(\s*-?\s*\d+(\.\d+)?\s*-?\s*\d+(\.\d+)?)?\s*([a-zA-Z/%]+)?\s*(H|L|Normal)?$",
            line,
        )
        if match:
            component = match.group(1).strip()
            value = float(match.group(2))
            range_match = match.group(4)
            if range_match:
                ranges = re.findall(r"[\d.]+", range_match)
                min_val = float(ranges[0]) if len(ranges) > 0 else None
                max_val = float(ranges[1]) if len(ranges) > 1 else None
            else:
                min_val = None
                max_val = None
            unit = match.group(7)
            flag = "Normal"  # Default flag

            if min_val is not None and max_val is not None:
                if value < min_val:
                    flag = "L"
                elif value > max_val:
                    flag = "H"

            if flag != "Normal":
                data.append([component, value, min_val, max_val, unit, flag])

    df = pd.DataFrame(data, columns=["Component", "Your Value", "Min", "Max", "Units", "Flag"])
    correction_map = {
        "emoglobin": "Hemoglobin",
        "ematocrit": "Hematocrit",
        "% Platelet Count": "Platelet Count",
        "ymphocyte %": "Lymphocyte %",
        "L Differential Type Automated": "Differential Type",
    }
    df["Component"] = df["Component"].replace(correction_map)

    return df


# Function to generate AI-powered recommendations using OpenAI GPT
def generate_medical_recommendation(test_results):
    """
    Generate medical recommendations using OpenAI GPT model based on abnormal test results.
    """
    # Create a structured input for the model
    prompt = f"Given the following blood test results: {test_results}, provide medical recommendations for a patient."

    response = openai.Completion.create(
        model="gpt-4",  # Use GPT-4 for medical-based responses
        prompt=prompt,
        max_tokens=150
    )

    return response.choices[0].text.strip()


# Function to display results
def display_results(df):
    """
    Display the flagged abnormalities and medical recommendations in a table format.
    """
    # Generate a summary of the abnormal test results for the AI model
    abnormal_results = df[df['Flag'] != 'Normal'].to_string(index=False)
    recommendation = generate_medical_recommendation(abnormal_results)
    
    # Display the DataFrame and the recommendation
    st.subheader("Flagged Abnormalities")
    st.dataframe(df, use_container_width=True)
    
    st.subheader("Medical Recommendations from AI")
    st.write(recommendation)


# Streamlit app
st.title("Blood Report Analyzer with AI Recommendations")
st.write("Upload an image of a blood test report to analyze and get AI-powered recommendations.")

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

if uploaded_file is not None:
    try:
        # Load the image
        image = Image.open(uploaded_file)

        # Display the uploaded image
        st.image(image, caption="Uploaded Image", use_container_width=True)

        # Extract text from the image
        extracted_text = extract_text(image)

        # Parse the extracted text into a structured format
        parsed_data = clean_and_parse_extracted_text(extracted_text)

        # Display the structured data with AI recommendations
        display_results(parsed_data)

    except Exception as e:
        st.error(f"An error occurred: {e}")