import os
import gradio as gr
from PIL import Image
import torch
from transformers import ViTForImageClassification, ViTImageProcessor
from datasets import load_dataset, DownloadConfig
import matplotlib.pyplot as plt
import numpy as np
import cv2
import requests

# Mistral AI API configuration
MISTRAL_API_KEY = "eoiBrPQzLjwNgOFgD7I4A4XF3TJOgBet"
MISTRAL_API_URL = "https://api.mistral.ai/v1/chat/completions"

def get_mistral_completion(prompt):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {MISTRAL_API_KEY}"
    }
    
    data = {
        "model": "mistral-medium",
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.7
    }
    
    response = requests.post(MISTRAL_API_URL, headers=headers, json=data)
    if response.status_code == 200:
        return response.json()['choices'][0]['message']['content']
    else:
        return "Error getting AI explanation. Please try again."

# Model and processor configuration
model_name_or_path = "google/vit-base-patch16-224-in21k"
processor = ViTImageProcessor.from_pretrained(model_name_or_path)

# Load dataset
dataset_path = "pawlo2013/chest_xray"
download_config = DownloadConfig(max_retries=10)
train_dataset = load_dataset(dataset_path, split="train", download_config=download_config)
class_names = train_dataset.features["label"].names

# Load ViT model
model = ViTForImageClassification.from_pretrained(
    "./models",
    num_labels=len(class_names),
    id2label={str(i): label for i, label in enumerate(class_names)},
    label2id={label: i for i, label in enumerate(class_names)},
)

model.eval()

def get_ai_explanation(diagnosis, probabilities):
    if diagnosis == "normal":
        prompt = f"""Given a chest X-ray analysis showing NORMAL results with {probabilities['normal']:.2%} confidence:
        1. Explain what this means, please remember that NORMAL ['normal'] means this user does not have any Pneumonia. 
        2. Suggest when they should still consider consulting a doctor even though this user does not have any penumonia as per the test result.
        3. List key symptoms that would warrant medical attention. Always identify yourself as PneumoInsight Bot.
        Keep the tone professional yet reassuring."""
    else:
        prompt = f"""Given a chest X-ray analysis showing {diagnosis} pneumonia with {probabilities[diagnosis]:.2%} confidence:
        1. Explain what {diagnosis} pneumonia is .  Always identify yourself as PneumoInsight Bot.
        2. List immediate steps the patient should take
        3. Provide care recommendations
        4. Mention warning signs to watch for
        Keep the tone informative and caring but emphasize the importance of professional medical consultation."""

    return get_mistral_completion(prompt)

def classify_and_visualize(img, device="cpu", discard_ratio=0.9, head_fusion="mean"):
    img = img.convert("RGB")
    processed_input = processor(images=img, return_tensors="pt").to(device)
    processed_input = processed_input["pixel_values"].to(device)

    with torch.no_grad():
        outputs = model(processed_input, output_attentions=True)
        logits = outputs.logits
        probabilities = torch.softmax(logits, dim=1)[0].tolist()
        prediction = torch.argmax(logits, dim=-1).item()
        predicted_class = class_names[prediction]

    result = {class_name: prob for class_name, prob in zip(class_names, probabilities)}

    # Generate attention heatmap
    heatmap_img = show_final_layer_attention_maps(
        outputs, processed_input, device, discard_ratio, head_fusion
    )

    return {"probabilities": result, "heatmap": heatmap_img}

def show_final_layer_attention_maps(
    outputs,
    processed_input,
    device,
    discard_ratio=0.6,
    head_fusion="max",
    only_last_layer=False,
):
    with torch.no_grad():
        image = processed_input.squeeze(0)
        image = image - image.min()
        image = image / image.max()

        result = torch.eye(outputs.attentions[0].size(-1)).to(device)
        if only_last_layer:
            attention_list = outputs.attentions[-1].unsqueeze(0).to(device)
        else:
            attention_list = outputs.attentions

        for attention in attention_list:
            if head_fusion == "mean":
                attention_heads_fused = attention.mean(axis=1)
            elif head_fusion == "max":
                attention_heads_fused = attention.max(axis=1)[0]
            elif head_fusion == "min":
                attention_heads_fused = attention.min(axis=1)[0]

            flat = attention_heads_fused.view(attention_heads_fused.size(0), -1)
            _, indices = flat.topk(int(flat.size(-1) * discard_ratio), -1, False)
            indices = indices[indices != 0]
            flat[0, indices] = 0

            I = torch.eye(attention_heads_fused.size(-1)).to(device)
            a = (attention_heads_fused + 1.0 * I) / 2
            a = a / a.sum(dim=-1)
            result = torch.matmul(a, result)

        mask = result[0, 0, 1:]
        width = int(mask.size(-1) ** 0.5)
        mask = mask.reshape(width, width).cpu().numpy()
        mask = mask / np.max(mask)

        mask = cv2.resize(mask, (224, 224))
        mask = (mask - np.min(mask)) / (np.max(mask) - np.min(mask))
        heatmap = plt.cm.jet(mask)[:, :, :3]

        showed_img = image.permute(1, 2, 0).detach().cpu().numpy()
        showed_img = (showed_img - np.min(showed_img)) / (
            np.max(showed_img) - np.min(showed_img)
        )
        superimposed_img = heatmap * 0.4 + showed_img * 0.6

        superimposed_img_pil = Image.fromarray(
            (superimposed_img * 255).astype(np.uint8)
        )
        return superimposed_img_pil

def load_examples_from_folder(folder_path):
    examples = []
    if os.path.exists(folder_path):
        for file in os.listdir(folder_path):
            if file.endswith((".png", ".jpg", ".jpeg")):
                examples.append(os.path.join(folder_path, file))
    return examples

def create_interface():
    # Custom CSS
    custom_css = """
    .logo-container { text-align: center; margin-bottom: 20px; }
    .logo-container img { max-width: 300px; }
    .welcome-message { text-align: center; margin: 20px 0; padding: 20px; background-color: #f5f5f5; border-radius: 10px; }
    .model-explanation { margin: 20px 0; padding: 20px; background-color: #f0f7ff; border-radius: 10px; }
    .pneumonia-info { margin: 20px 0; padding: 20px; background-color: #fff5f5; border-radius: 10px; }
    .disclaimer { margin-top: 20px; padding: 20px; background-color: #f5f5f5; border-radius: 10px; font-size: 0.9em; }
    """

    # HTML Components
    logo_html = """
    <div class="logo-container">
        <img src="file/logo.png" alt="PneumoInsight Logo">
    </div>
    """

    welcome_message = """
    <div class="welcome-message">
        <h1>🫁 Welcome to PneumoInsight</h1>
        <p>PneumoInsight is a side project of EarlyMed—an initiative by our team at VIT-AP University dedicated to empowering you with early health insights. 
        Leveraging AI for early detection, our mission is simple: "Early Detection, Smarter Decision." 
        This project is one of our key efforts to help you stay informed before visiting a doctor.</p>
    </div>
    """

    model_explanation = """
    <div class="model-explanation">
        <h2>How Our Model Works</h2>
        <p>Our system uses a Vision Transformer (ViT) model to analyze chest X-ray images. The attention heatmap visualizes 
        areas the AI focuses on while making its diagnosis, helping make the decision-making process more transparent. 
        The warmer colors (red/yellow) indicate areas of higher attention.</p>
        <p>Credits: The attention heatmap visualization is implemented using the attention rollout technique by 
        <a href="https://github.com/jacobgil/vit-explain" target="_blank">jacobgil</a>.</p>
    </div>
    """

    pneumonia_info = """
    <div class="pneumonia-info">
        <h2>Understanding Pneumonia</h2>
        <p>Pneumonia is an infection that inflames the air sacs in one or both lungs. Common symptoms include:</p>
        <ul>
            <li>Chest pain when breathing or coughing</li>
            <li>Cough with phlegm or pus</li>
            <li>Fatigue and difficulty breathing</li>
            <li>Fever, sweating, and shaking chills</li>
        </ul>
        <p>Prevention tips:</p>
        <ul>
            <li>Get vaccinated</li>
            <li>Practice good hygiene</li>
            <li>Don't smoke</li>
            <li>Maintain a strong immune system</li>
        </ul>
    </div>
    """

    disclaimer = """
    <div class="disclaimer">
        <h3>Disclaimer</h3>
        <p>This tool is for educational purposes only and should not be used as a substitute for professional medical advice, 
        diagnosis, or treatment. Always seek the advice of your physician or other qualified health provider.</p>
        <p>Created by the team at VIT-AP University. View the source code on 
        <a href="https://github.com/Mahatir-Ahmed-Tusher/PneumoInsight" target="_blank">GitHub</a>.</p>
    </div>
    """

    def enhanced_classification(img):
        if img is None:
            return None, None, "Please upload an image to proceed."
        
        result = classify_and_visualize(img)
        probabilities = result["probabilities"]
        heatmap = result["heatmap"]
        
        # Get the predicted class
        predicted_class = max(probabilities.items(), key=lambda x: x[1])[0]
        
        # Get AI explanation
        ai_explanation = get_ai_explanation(predicted_class, probabilities)
        
        return probabilities, heatmap, ai_explanation

    # Create the Gradio interface
    iface = gr.Interface(
        fn=enhanced_classification,
        inputs=gr.Image(type="pil", label="Upload Chest X-Ray Image"),
        outputs=[
            gr.Label(label="Diagnosis Probabilities"),
            gr.Image(label="Attention Heatmap"),
            gr.Textbox(label="AI Analysis and Recommendations", lines=10)
        ],
        css=custom_css,
        examples=load_examples_from_folder("./Examples"),
        cache_examples=False,
        article=model_explanation + pneumonia_info + disclaimer,
        description=welcome_message,
        title=logo_html,
        theme=gr.themes.Soft()
    )
    
    return iface

# Launch the app
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(debug=True)