import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
from PIL import Image
import streamlit as st
from groq import Groq

# Set Groq API key in environment variable
os.environ['GROQ_API_KEY'] = "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"
GROQ_API_KEY = os.getenv('GROQ_API_KEY')

# Initialize Groq client
client = Groq(api_key=GROQ_API_KEY)

# Load the classification model (make sure the model is trained and saved)
classification_model = load_model('classification_model.h5')  # Model for normal/abnormal classification

# Function to load the image and process it
def load_image(image_file):
    img = Image.open(image_file)
    img = img.resize((224, 224))  # Resize image to match model input
    img_array = np.array(img) / 255.0  # Normalize the image
    return np.expand_dims(img_array, axis=0)

# Function for AI-based knowledge generation using Groq API
def generate_ai_insights(organ_name):
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": f"Explain the diseases and treatments related to {organ_name}.",
            }
        ],
        model="llama-3.3-70b-versatile",
    )
    return chat_completion.choices[0].message.content

# Streamlit UI
st.title('Medical Image Classification and Insights')
st.sidebar.title("Menu")

uploaded_image = st.sidebar.file_uploader("Upload X-ray or MRI Image", type=["jpg", "png", "jpeg"])

if uploaded_image is not None:
    image = load_image(uploaded_image)
    
    # Classify normal or abnormal
    prediction = classification_model.predict(image)
    if prediction[0] > 0.5:
        classification_result = "Normal"
    else:
        classification_result = "Abnormal"

    st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
    st.write(f"Image Classification: {classification_result}")
    
    # Recognize the organ (You can expand the model to predict organ type)
    organ_name = "Lung"  # Placeholder
    st.write(f"Recognized Organ: {organ_name}")
    
    # Get AI insights for the recognized organ
    ai_insights = generate_ai_insights(organ_name)
    st.write("AI-Based Insights on Organ Diseases and Treatments:")
    st.write(ai_insights)