import tensorflow as tf
import random
from PIL import Image
from tensorflow import keras
import numpy as np
import os
import logging
from tensorflow.keras.preprocessing import image as keras_image
from huggingface_hub import from_pretrained_keras
from openai import AzureOpenAI
import gradio as gr
from zipfile import ZipFile

logging.basicConfig(level=logging.INFO)

class DiseaseDetectionApp:
    def __init__(self):
      
        self.class_names = [
            "Normal",
            "PNEUMONIA",
        ]

        keras.utils.set_random_seed(42)
        self.model = from_pretrained_keras("ryefoxlime/PneumoniaDetection")
        self.client=AzureOpenAI()
        

    def predict_disease(self, image_path):
        """
        Predict the disease present in the X-Ray image.

        Args:
        - image_data: PIL image data

        Returns:
        - predicted_disease: string
        """
        try:
            # Load the image file, resizing it to the dimensions expected by the model
            img = keras_image.load_img(image_path, target_size=(224, 224))  # Adjust target_size according to your model's expected input dimensions

            # Convert the image to a numpy array
            img_array = keras_image.img_to_array(img)

            # Add an additional dimension to the array: (1, height, width, channels)
            img_array = tf.expand_dims(img_array, 0)  # Model expects a batch of images, but we're only passing a single image

            # Make predictions
            predictions =self.model.predict(img_array)

            # Extract the predicted class and confidence
            predict_class = self.class_names[np.argmax(predictions[0])]
            # confidence = round(100 * np.max(predictions[0]), 2)

            return predict_class

        except Exception as e:
            logging.error(f"Error predicting disease: {str(e)}")
            return None

    def classify_disease(self,image_path):
      
          disease_name=self.predict_disease(image_path)
          print(disease_name)
          if disease_name=="PNEUMONIA":
            conversation = [
              {"role": "system", "content": "You are a medical assistant"},
              {"role": "user", "content": f""" your task describe(classify) about the given disease as summary only in 3 lines.
                        ```{disease_name}```
                """}
                ]
            # Generate completion using ChatGPT model
            response = self.client.chat.completions.create(
                model="ChatGPT",
                messages=conversation,
                temperature=0,
                max_tokens=1000
            )
            # Get the generated topics message
            
            result = response.choices[0].message.content
            return disease_name,result

          elif disease_name=="Normal":
            result="No problem in your xray image"
            return disease_name,result
          
          else:
            logging.error("Error classify_disease disease")
            return "Something went wrong"

    def unzip_image_data(self):
        """
        Unzips an image dataset into a specified directory.

        Returns:
            str: The path to the directory containing the extracted image files.
        """
        try:
          with ZipFile("image_dataset.zip","r") as extract:
            directory_path="dataset_image"
            extract.extractall(f"{directory_path}")
            return f"{directory_path}"

        except Exception as e:
                logging.error(f"An error occurred during extraction: {e}")
                return ""

    def example_images(self):
          """
          Unzips the image dataset and generates a list of paths to the individual image files and use image for showing example

          Returns:
              List[str]: A list of file paths to each image in the dataset.
          """
          image_dataset_folder = self.unzip_image_data()
          image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp']
          image_count = len([name for name in os.listdir(image_dataset_folder) if os.path.isfile(os.path.join(image_dataset_folder, name)) and os.path.splitext(name)[1].lower() in image_extensions])
          example=[]
          for i in range(image_count):
            for name in os.listdir(image_dataset_folder):
                path=(os.path.join(os.path.dirname(image_dataset_folder),os.path.join(image_dataset_folder,name)))
                example.append(path)
          print()
          return example
          

    def gradio_interface(self):

      with gr.Blocks(theme='JohnSmith9982/small_and_pretty') as demo:
        gr.HTML("""<center><h1>Pneumonia Disease Detection</h1></center>""")

        exam_img=self.example_images()
        with gr.Row():
          input_image =gr.Image(type="filepath",sources="upload")
          with gr.Column():
            output=gr.Label(label="Disease Name")
        with gr.Row():
          classify_disease_=gr.Textbox(label="About disease")
        with gr.Row():
          button =gr.Button(value="Detect The Disease")

        button.click(self.classify_disease,[input_image],[output,classify_disease_])
        gr.Examples(
        examples=exam_img,
        inputs=[input_image],
        outputs=[output,classify_disease_],
        fn=self.classify_disease,
        cache_examples=False)
    
        
      demo.launch(debug=True)

if __name__ == "__main__":
    app = DiseaseDetectionApp()
    result=app.gradio_interface()