# Use TensorFlow backend for better compatibility with existing models
import os
os.environ["KERAS_BACKEND"] = "tensorflow"

import gradio as gr
import tensorflow as tf
from tensorflow import keras
import numpy as np
from PIL import Image
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
from huggingface_hub import hf_hub_download

# Class names for the 5 skin conditions
CLASS_NAMES = [
    'Atopic Dermatitis',
    'Eczema', 
    'Psoriasis',
    'Seborrheic Keratoses',
    'Tinea Ringworm Candidiasis'
]

# Class descriptions
CLASS_DESCRIPTIONS = {
    'Atopic Dermatitis': 'A chronic inflammatory skin condition causing dry, itchy patches',
    'Eczema': 'Inflammatory skin condition causing red, itchy, and inflamed patches',
    'Psoriasis': 'Autoimmune condition causing thick, scaly patches on the skin',
    'Seborrheic Keratoses': 'Common benign (non-cancerous) skin growths',
    'Tinea Ringworm Candidiasis': 'Fungal skin infections causing circular, scaly patches'
}

class DermaAIModel:
    def __init__(self):
        self.model = None
        self.load_model()
    
    def load_model(self):
        """Load the DermaAI model from Hugging Face Hub using traditional approach"""
        try:
            print("🔄 Loading DermaAI model from Hugging Face...")
            
            # Get authentication token from environment variable only (secure)
            hf_token = os.getenv("HF_TOKEN")
            
            # Download model file using HuggingFace Hub
            if hf_token:
                print("🔐 Using authentication token...")
                model_path = hf_hub_download(
                    repo_id="Siraja704/DermaAI",
                    filename="DermaAI.keras",
                    token=hf_token,
                    cache_dir="./model_cache"
                )
            else:
                print("📂 Trying without authentication...")
                model_path = hf_hub_download(
                    repo_id="Siraja704/DermaAI",
                    filename="DermaAI.keras",
                    cache_dir="./model_cache"
                )
            
            # Load the model using TensorFlow/Keras
            print(f"📁 Loading model from: {model_path}")
            self.model = keras.models.load_model(model_path)
            print("✅ Model loaded successfully!")
            
        except Exception as e:
            error_msg = str(e)
            print(f"❌ Error loading model: {e}")
            if "401" in error_msg or "gated" in error_msg.lower() or "restricted" in error_msg.lower():
                print("\n🔐 AUTHENTICATION ERROR:")
                print("- The model repository is private/gated")
                print("- Please add your HF_TOKEN to the Space secrets")
                print("- Go to Space Settings > Repository secrets")
                print("- Add: HF_TOKEN = your_huggingface_token")
                print("- Make sure the token has access to Siraja704/DermaAI\n")
            raise e
    
    def predict(self, image):
        """Make prediction on the input image"""
        if self.model is None:
            return {"error": "Model not loaded"}
        
        try:
            # Preprocess image
            if image is None:
                return {"error": "No image provided"}
            
            # Convert to RGB if necessary
            if image.mode != 'RGB':
                image = image.convert('RGB')
            
            # Resize to model input size
            image_resized = image.resize((224, 224))
            
            # Convert to numpy array and preprocess
            image_array = np.array(image_resized)
            image_array = preprocess_input(image_array)
            image_array = np.expand_dims(image_array, axis=0)
            
            # Make prediction
            predictions = self.model.predict(image_array, verbose=0)
            
            # Get results
            predicted_class_idx = np.argmax(predictions[0])
            confidence = float(predictions[0][predicted_class_idx])
            
            # Prepare results dictionary for Gradio
            results = {}
            for i, class_name in enumerate(CLASS_NAMES):
                results[class_name] = float(predictions[0][i])
            
            return results
            
        except Exception as e:
            print(f"❌ Error during prediction: {e}")
            return {"error": f"Prediction failed: {str(e)}"}

# Initialize model
print("🚀 Initializing DermaAI...")
derma_model = DermaAIModel()

def predict_skin_condition(image):
    """Wrapper function for Gradio interface"""
    if image is None:
        return {"error": "Please upload an image"}
    
    return derma_model.predict(image)

def get_medical_advice(image):
    """Provide medical advice based on prediction"""
    if image is None:
        return "Please upload an image first."
    
    results = derma_model.predict(image)
    
    if "error" in results:
        return results["error"]
    
    # Find the top prediction
    top_prediction = max(results, key=results.get)
    confidence = results[top_prediction] * 100
    
    # Generate advice based on confidence
    advice = f"**Predicted Condition:** {top_prediction}\n\n"
    advice += f"**Confidence:** {confidence:.1f}%\n\n"
    advice += f"**Description:** {CLASS_DESCRIPTIONS.get(top_prediction, 'No description available')}\n\n"
    
    if confidence < 30:
        advice += "⚠️ **Low Confidence Warning:** The AI model has low confidence in this prediction. Please retake the image with better lighting and focus, or consult a healthcare professional."
    elif confidence < 60:
        advice += "🔍 **Moderate Confidence:** This is a preliminary assessment. Consider consulting with a healthcare professional for accurate diagnosis."
    else:
        advice += "✅ **High Confidence:** The model shows high confidence, but this is still a preliminary assessment."
    
    advice += "\n\n🏥 **Important Medical Disclaimer:** This AI tool is for educational purposes only and should not replace professional medical diagnosis. Always consult qualified healthcare professionals for proper medical evaluation and treatment."
    
    return advice

# Custom CSS for better styling with dark/light mode support
custom_css = """
/* Simple blue theme that works in both light and dark modes */
.gradio-container {
    font-family: 'Inter', 'Segoe UI', 'Roboto', sans-serif;
    max-width: 1200px;
    margin: 0 auto;
}

/* All text blue for universal visibility */
.gradio-container * {
    color: #2196f3 !important;
}

/* Main heading - darker blue */
.main-title h1 {
    color: #1565c0 !important;
    font-size: 2.5rem !important;
    font-weight: 700 !important;
    text-align: center !important;
}

.main-title p {
    color: #1976d2 !important;
    font-size: 1.2rem !important;
    text-align: center !important;
}

/* Section headings - medium blue */
.gradio-container h1,
.gradio-container h2,
.gradio-container h3,
.gradio-container h4,
.gradio-container h5,
.gradio-container h6,
.section-header {
    color: #1976d2 !important;
    font-weight: 600 !important;
}

.section-header {
    font-size: 1.4rem !important;
    margin-bottom: 15px !important;
}

/* Regular text - lighter blue */
.gradio-container p,
.gradio-container span,
.gradio-container div,
.gradio-container li,
.gradio-container label,
.gradio-container button {
    color: #2196f3 !important;
}

/* Special components */
.gradio-container .gr-markdown,
.gradio-container .gr-markdown *,
.gradio-container .gr-label-text,
.gradio-container .gr-input-label {
    color: #2196f3 !important;
}

/* Icons and SVG elements */
.gradio-container svg {
    fill: #2196f3 !important;
    stroke: #2196f3 !important;
}

/* Links */
.gradio-container a {
    color: #1565c0 !important;
    text-decoration: underline;
}

/* Simple styling without backgrounds */
.medical-disclaimer,
.info-box,
.about-section {
    border-radius: 8px;
    padding: 15px;
    margin: 10px 0;
    border: 1px solid #2196f3;
}

.conditions-grid {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(300px, 1fr));
    gap: 15px;
    margin-top: 15px;
}

.condition-item {
    padding: 15px;
    border-radius: 8px;
    border: 1px solid #2196f3;
}
"""

# Create Gradio interface
with gr.Blocks(
    css=custom_css,
    title="DermaAI - Skin Disease Classification",
    theme=gr.themes.Base().set(
        button_primary_background_fill="*primary_500",
        button_primary_background_fill_hover="*primary_600",
        button_primary_text_color="white",
        block_background_fill="*background_fill_primary",
        body_background_fill="*background_fill_primary",
    )
) as demo:
    
    gr.HTML("""
    <div class="main-title">
        <h1>🏥 DermaAI - Skin Disease Classification</h1>
        <p>AI-powered skin condition analysis using deep learning</p>
    </div>
    """)
    
    gr.HTML("""
    <div class="medical-disclaimer">
        <h3>⚕️ Important Medical Disclaimer</h3>
        <p><strong>This AI tool is for educational and research purposes only.</strong> 
        It should not be used as a substitute for professional medical diagnosis or treatment. 
        Always consult with qualified healthcare professionals for proper medical evaluation.</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.HTML("<h3 class='section-header'>📸 Upload Skin Image</h3>")
            input_image = gr.Image(
                type="pil",
                label="Upload a clear image of the skin condition",
                height=400
            )
            
            gr.HTML("""
            <div class="info-box">
                <h4>📋 Image Guidelines:</h4>
                <ul>
                    <li>Use good lighting and focus</li>
                    <li>Ensure the affected area is clearly visible</li>
                    <li>Avoid blurry or dark images</li>
                    <li>JPG, PNG formats supported</li>
                </ul>
            </div>
            """)
        
        with gr.Column(scale=1):
            gr.HTML("<h3 class='section-header'>🔍 Analysis Results</h3>")
            
            prediction_output = gr.Label(
                label="Prediction Confidence Scores",
                num_top_classes=5
            )
            
            medical_advice = gr.Markdown(
                label="Medical Assessment",
                value="Upload an image to see the analysis..."
            )
    
    gr.HTML("""
    <div class="info-box">
        <h3>🩺 Supported Skin Conditions</h3>
        <div class="conditions-grid">
            <div class="condition-item"><strong>Atopic Dermatitis:</strong> Chronic inflammatory skin condition</div>
            <div class="condition-item"><strong>Eczema:</strong> Red, itchy, inflamed skin patches</div>
            <div class="condition-item"><strong>Psoriasis:</strong> Thick, scaly skin patches</div>
            <div class="condition-item"><strong>Seborrheic Keratoses:</strong> Benign skin growths</div>
            <div class="condition-item"><strong>Tinea Ringworm Candidiasis:</strong> Fungal skin infections</div>
        </div>
    </div>
    """)
    
    # Set up the interface interactions
    input_image.change(
        fn=predict_skin_condition,
        inputs=input_image,
        outputs=prediction_output
    )
    
    input_image.change(
        fn=get_medical_advice,
        inputs=input_image,
        outputs=medical_advice
    )
    
    gr.HTML("""
    <div class="about-section">
        <h3>📚 About DermaAI</h3>
        <p>DermaAI is built using EfficientNetV2 architecture and trained on dermatological images. 
        The model analyzes skin images and provides confidence scores for 5 different skin conditions.</p>
        <p><strong>Model:</strong> <a href="https://huggingface.co/Siraja704/DermaAI" target="_blank">Siraja704/DermaAI</a></p>
        <p><strong>Framework:</strong> TensorFlow/Keras | <strong>Architecture:</strong> EfficientNetV2</p>
    </div>
    """)

# Launch the interface
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )