import torch
import gradio as gr
import timm
from torchvision import transforms as T
from PIL import Image
import numpy as np
import cv2
from pytorch_grad_cam import GradCAMPlusPlus
from pytorch_grad_cam.utils.image import show_cam_on_image

# ==========================================
# 1. Configuration & Setup
# ==========================================
class_names = ['Anger', 'Fear', 'Joy', 'Neutral']

model_name = "resnext101_32x8d"
weights_path = "autism_best_model.pth"
im_size = 224
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]

# Hugging Face Free Tier uses CPU
device = torch.device("cpu")

# ==========================================
# 2. Model Loading
# ==========================================
def load_model():
    print("Loading Model...")
    model = timm.create_model(
        model_name, pretrained=False, num_classes=len(class_names))
    state_dict = torch.load(weights_path, map_location=device)
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    return model

model = load_model()

# ==========================================
# 3. Preprocessing
# ==========================================
transform = T.Compose([
    T.Resize((im_size, im_size)),
    T.ToTensor(),
    T.Normalize(mean=mean, std=std)
])

# ==========================================
# 4. Inference & Grad-CAM Function
# ==========================================
# 🟢 Fixed: Added !important to force pure black text so Dark Mode doesn't hide it
waiting_text = """
<div style='background-color: #F8F9FA !important; padding: 15px; border-radius: 8px; border: 2px solid #DEE2E6; text-align: center;'>
    <h3 style='margin: 0; color: #000000 !important; font-weight: bold;'>⏳ Waiting for upload...</h3>
</div>
"""

def predict(image):
    if image is None:
        return None, waiting_text, None

    # Prepare Image
    image = image.convert("RGB")
    input_tensor = transform(image).unsqueeze(0).to(device)

    # 1. Classification (Inside no_grad)
    with torch.no_grad():
        outputs = model(input_tensor)
        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)

    confidences = {class_names[i]: float(probabilities[i]) for i in range(len(class_names))}
    top_class = max(confidences, key=confidences.get)
    top_score = confidences[top_class]
    
    # 🟢 Fixed: Forced pure black text on light blue background
    message = f"""
    <div style='background-color: #E0F2FE !important; padding: 12px; border-radius: 8px; border: 2px solid #BAE6FD; text-align: center;'>
        <h3 style='margin: 0; color: #000000 !important; font-weight: bold;'>
            🧠 The model is <span style='color: #0369A1 !important;'>{top_score*100:.1f}%</span> confident that the primary expression is <span style='color: #0369A1 !important;'>{top_class}</span>.
        </h3>
    </div>
    """

    # 2. Grad-CAM Generation (Outside no_grad because Grad-CAM requires gradients)
    target_layers = [model.layer4[-1].conv3]
    cam = GradCAMPlusPlus(model=model, target_layers=target_layers, use_cuda=False) # CPU mode
    
    # Generate heatmap
    grayscale_cam = cam(input_tensor=input_tensor)[0, :]
    
    # Overlay heatmap on original image
    rgb_img = np.array(image.resize((im_size, im_size))) / 255.0
    cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=0.4)

    return confidences, message, cam_image

# ==========================================
# 5. Modern Gradio UI (Blocks)
# ==========================================
theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"]
)

custom_css = """
.gradio-container { background-color: #f8fafc; }
.header-box { text-align: center; padding: 2rem; background: linear-gradient(135deg, #4f46e5 0%, #3b82f6 100%); border-radius: 15px; color: white; box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1); margin-bottom: 20px; }
.header-title { font-size: 2.5rem; font-weight: 800; margin-bottom: 0.5rem; }
.header-subtitle { font-size: 1.1rem; opacity: 0.9; }
"""

with gr.Blocks(theme=theme, css=custom_css) as demo:

    gr.HTML("""
    <div class="header-box">
        <div class="header-title">✨ Facial Emotion Recognition in Autistic Children</div>
        <div class="header-subtitle">Explainable deep learning model trained on the Dataset to analyze facial emotions of autistic children.</div>
    </div>
    """)

    with gr.Row():
        # Left Column: Inputs
        with gr.Column(scale=1):
            
            # 🟢 Fixed: Forced Pure Black Text on White Background
            gr.HTML("""
            <div style='background-color: #E0F2FF !important; padding: 12px; border-radius: 8px; border: 2px solid #E2E8F0; margin-bottom: 10px;'>
                <h3 style='margin: 0 0 5px 0; color: #000000 !important; font-weight: bold;'>📸 Upload Image</h3>
                <p style='margin: 0; color: #000000 !important;'>Upload a clear photo of a child's face.</p>
            </div>
            """)
            
            input_image = gr.Image(type="pil", label="Input Image", elem_classes="image-box")

            with gr.Row():
                clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                submit_btn = gr.Button("🚀 Analyze Expression", variant="primary")

        # Right Column: Outputs (Parallel to Input)
        with gr.Column(scale=1):
            
            # 🟢 Fixed: Forced Pure Black Text on White Background
            gr.HTML("""
            <div style='background-color: #E0F2FE !important; padding: 12px; border-radius: 8px; border: 2px solid #E2E8F0; margin-bottom: 10px;'>
                <h3 style='margin: 0; color: #000000 !important; font-weight: bold;'>📊 Result Analysis & Grad-CAM Visualization</h3>
            </div>
            """)
            
            # Message Block
            output_message = gr.HTML(waiting_text)
            
            # Parallel Display for Labels and Grad-CAM image
            with gr.Row():
                output_label = gr.Label(num_top_classes=4, label="Confidence Scores")
                output_cam = gr.Image(type="numpy", label="Grad-CAM Focus Map")

    with gr.Accordion("ℹ️ About this Model", open=False):
        gr.Markdown("""
        * **Architecture:** ResNeXt-101 (32x8d)
        * **Input Resolution:** 224x224 pixels
        * **Visualizer:** Grad-CAM (Highlights the impactful region)
        * **Note:** This tool is for research demonstration purposes and is not a clinical diagnostic tool.
        """)

    # Button Functionality
    submit_btn.click(
        fn=predict,
        inputs=[input_image],
        outputs=[output_label, output_message, output_cam]
    )

    clear_btn.click(
        fn=lambda: (None, waiting_text, None),
        inputs=[],
        outputs=[input_image, output_label, output_message, output_cam]
    )

if __name__ == "__main__":
    demo.launch()