"""
🫁 Multi-Class Chest X-Ray Detection with Adaptive Sparse Training
4-Class Screening: Normal, Tuberculosis, Pneumonia, COVID-19

Mission:
    This open research tool is being built to help humanity –
    especially patients and clinicians in low-resource settings –
    by providing energy-efficient, explainable AI support for chest
    X-ray screening. It is a digital second opinion, NOT a replacement
    for radiologists or doctors.
"""

import gradio as gr
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import numpy as np
import cv2
import matplotlib.pyplot as plt
from pathlib import Path
import io

# ============================================================================
# Model Setup
# ============================================================================

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_efficientnet_model():
    """
    Build EfficientNet-B2 and load your working 4-class best.pt checkpoint.

    We intentionally keep this simple and very close to the version you
    already confirmed is working, to avoid shape-mismatch issues.
    """
    # Base architecture: EfficientNet-B2
    model = models.efficientnet_b2(weights=None)
    in_features = model.classifier[1].in_features
    model.classifier[1] = nn.Linear(in_features, 4)  # 4 classes

    # Where we expect your weights to live
    candidate_paths = [
        Path("checkpoints/best.pt"),  # HF Space path (from your screenshot)
        Path("best.pt"),              # fallback for local runs
    ]

    last_error = None
    for ckpt_path in candidate_paths:
        if not ckpt_path.exists():
            print(f"⚠️  Checkpoint not found at {ckpt_path}")
            continue

        try:
            print(f"🔍 Loading weights from: {ckpt_path}")
            state = torch.load(ckpt_path, map_location=device)

            # If it comes from a training script with wrappers
            if isinstance(state, dict):
                if "model_state_dict" in state:
                    state = state["model_state_dict"]
                elif "state_dict" in state:
                    state = state["state_dict"]

            # This is the same idea as your original working call
            missing, unexpected = model.load_state_dict(state, strict=False)
            if missing or unexpected:
                print(f"   ⚠️ Non-critical keys - missing: {missing}, unexpected: {unexpected}")
            print(f"✅ Model weights successfully loaded from {ckpt_path}")
            model.to(device)
            model.eval()
            return model
        except Exception as e:
            print(f"❌ Could not load from {ckpt_path}: {e}")
            last_error = e

    raise RuntimeError(
        "Could not load EfficientNet-B2 4-class weights from any known path.\n"
        f"Last error: {last_error}"
    )


model = load_efficientnet_model()

# Classes
CLASSES = ["Normal", "Tuberculosis", "Pneumonia", "COVID-19"]
CLASS_COLORS = {
    "Normal": "#2ecc71",       # Green
    "Tuberculosis": "#e74c3c",  # Red
    "Pneumonia": "#f39c12",     # Orange
    "COVID-19": "#9b59b6",      # Purple
}

# Image preprocessing
transform = transforms.Compose(
    [
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            [0.485, 0.456, 0.406],
            [0.229, 0.224, 0.225],
        ),
    ]
)

# ============================================================================
# Grad-CAM Implementation
# ============================================================================


class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None

        def save_gradient(grad):
            self.gradients = grad

        def save_activation(module, input, output):
            self.activations = output.detach()

        # Forward hook: store activations
        target_layer.register_forward_hook(save_activation)
        # Backward hook: store gradients
        target_layer.register_full_backward_hook(
            lambda m, grad_in, grad_out: save_gradient(grad_out[0])
        )

    def generate(self, input_image, target_class=None):
        output = self.model(input_image)

        if target_class is None:
            target_class = output.argmax(dim=1)

        self.model.zero_grad()
        one_hot = torch.zeros_like(output)
        one_hot[0][target_class] = 1
        output.backward(gradient=one_hot, retain_graph=True)

        if self.gradients is None or self.activations is None:
            return None, output

        # Global average pooling over gradients
        weights = self.gradients.mean(dim=(2, 3), keepdim=True)
        cam = (weights * self.activations).sum(dim=1, keepdim=True)
        cam = torch.relu(cam)
        cam = cam.squeeze().cpu().numpy()
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)

        return cam, output


# Setup Grad-CAM on the last feature layer
target_layer = model.features[-1]
grad_cam = GradCAM(model, target_layer)

# ============================================================================
# Prediction & Visualization
# ============================================================================


def predict_chest_xray(image, show_gradcam=True):
    """
    Predict disease class from chest X-ray with Grad-CAM visualization.

    Returns:
        - class probabilities dict
        - annotated original image
        - Grad-CAM heatmap image
        - overlay image
        - markdown clinical interpretation
    """
    if image is None:
        return None, None, None, None, "Please upload a chest X-ray."

    # Convert to PIL if needed
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image).convert("RGB")
    else:
        image = image.convert("RGB")

    # Keep original for visualization
    original_img = image.copy()

    # Preprocess
    input_tensor = transform(image).unsqueeze(0).to(device)

    # Forward + optional Grad-CAM
    with torch.set_grad_enabled(show_gradcam):
        if show_gradcam:
            cam, output = grad_cam.generate(input_tensor)
        else:
            cam = None
            output = model(input_tensor)

    # Probabilities
    probs = torch.softmax(output, dim=1)[0].cpu().detach().numpy()
    prob_sum = float(np.sum(probs))
    if not (0.99 <= prob_sum <= 1.01):
        print(f"⚠️ Probability sum is {prob_sum:.4f}, expected ~1.0 – check model weights.")

    pred_class = int(output.argmax(dim=1).item())
    pred_label = CLASSES[pred_class]
    confidence = float(probs[pred_class] * 100.0)

    # Ensure values between 0–100
    results = {
        CLASSES[i]: float(min(100.0, max(0.0, probs[i] * 100.0)))
        for i in range(len(CLASSES))
    }

    # Visualizations
    original_pil = create_original_display(original_img, pred_label, confidence)

    if cam is not None and show_gradcam:
        gradcam_viz = create_gradcam_visualization(
            original_img, cam, pred_label, confidence
        )
        overlay_viz = create_overlay_visualization(original_img, cam)
    else:
        gradcam_viz = None
        overlay_viz = None

    # Interpretation text
    interpretation = create_interpretation(pred_label, confidence, results)

    return results, original_pil, gradcam_viz, overlay_viz, interpretation


def create_original_display(image, pred_label, confidence):
    """Create annotated original image."""
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.imshow(image)
    ax.axis("off")

    color = CLASS_COLORS[pred_label]
    title = f"Prediction: {pred_label}\nConfidence: {confidence:.1f}%"
    ax.set_title(title, fontsize=16, fontweight="bold", color=color, pad=20)

    plt.tight_layout()
    buf = io.BytesIO()
    plt.savefig(
        buf,
        format="png",
        dpi=150,
        bbox_inches="tight",
        facecolor="white",
    )
    plt.close()
    buf.seek(0)

    return Image.open(buf)


def create_gradcam_visualization(image, cam, pred_label, confidence):
    """Create Grad-CAM heatmap."""
    img_array = np.array(image.resize((224, 224)))
    cam_resized = cv2.resize(cam, (224, 224))

    heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)

    fig, ax = plt.subplots(figsize=(8, 8))
    ax.imshow(heatmap)
    ax.axis("off")
    ax.set_title(
        "Attention Heatmap\n(Areas the model focuses on)",
        fontsize=14,
        fontweight="bold",
        pad=20,
    )

    plt.tight_layout()
    buf = io.BytesIO()
    plt.savefig(
        buf,
        format="png",
        dpi=150,
        bbox_inches="tight",
        facecolor="white",
    )
    plt.close()
    buf.seek(0)

    return Image.open(buf)


def create_overlay_visualization(image, cam):
    """Overlay original image and Grad-CAM heatmap."""
    img_array = np.array(image.resize((224, 224))) / 255.0
    cam_resized = cv2.resize(cam, (224, 224))

    heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB) / 255.0

    overlay = img_array * 0.5 + heatmap * 0.5
    overlay = np.clip(overlay, 0, 1)

    fig, ax = plt.subplots(figsize=(8, 8))
    ax.imshow(overlay)
    ax.axis("off")
    ax.set_title(
        "Explainable AI Visualization\n(Original + Heatmap)",
        fontsize=14,
        fontweight="bold",
        pad=20,
    )

    plt.tight_layout()
    buf = io.BytesIO()
    plt.savefig(
        buf,
        format="png",
        dpi=150,
        bbox_inches="tight",
        facecolor="white",
    )
    plt.close()
    buf.seek(0)

    return Image.open(buf)


def create_interpretation(pred_label, confidence, results):
    """
    Clinical-style interpretation text with strong global-health framing
    and strict medical disclaimer.
    """

    interpretation = f"""
## 🫁 AI Chest X-Ray Screening – Global Health Edition
This tool is part of an open effort to **support clinicians and patients worldwide**,
especially in places where radiologists are scarce.

---

## 🔬 Analysis Summary
**Predicted class:** **{pred_label}**  
**Model confidence:** **{confidence:.1f}%**

### Probability Breakdown
- 🟢 Normal: **{results['Normal']:.1f}%**
- 🔴 Tuberculosis: **{results['Tuberculosis']:.1f}%**
- 🟠 Pneumonia: **{results['Pneumonia']:.1f}%**
- 🟣 COVID-19: **{results['COVID-19']:.1f}%**

---
"""

    # Disease-specific details
    if pred_label == "Tuberculosis":
        if confidence >= 85:
            interpretation += """
### ⚠️ High-Confidence Tuberculosis Pattern Detected

The AI model has found features strongly suggestive of **pulmonary tuberculosis (TB)**.

**Suggested next steps for a clinical team (NOT automatic orders):**
1. Correlate with symptoms:
   - Cough > 2 weeks
   - Night sweats, fever
   - Weight loss
   - Hemoptysis (coughing blood)
2. Order **confirmatory TB tests**:
   - Sputum smear / culture
   - GeneXpert MTB/RIF or TB-PCR
3. Consider **isolation** and **contact screening** if TB is suspected.
4. Evaluate HIV status and comorbidities according to local guidelines.

➡️ This system is designed to **support TB programs** in low-resource settings,
where early triage can save lives.
"""
        else:
            interpretation += """
### ⚠️ Possible Tuberculosis Features

The model sees **TB-like patterns**, but confidence is moderate.

**Recommended clinical follow-up (not automatic diagnosis):**
- Detailed history and physical examination
- Evaluate TB risk factors and symptoms
- Consider sputum-based TB testing
- Repeat imaging or CT if clinically indicated
"""

    elif pred_label == "Pneumonia":
        if confidence >= 85:
            interpretation += """
### ⚠️ High-Confidence Pneumonia Pattern

The model detects findings consistent with **pneumonia**.

**Clinical team may consider:**
- Distinguishing bacterial vs viral pneumonia
- Correlating with:
  - Fever, cough, sputum
  - Pleuritic chest pain
  - Shortness of breath
- Laboratory tests (WBC, CRP, cultures)
- Guideline-based antibiotic or supportive therapy if confirmed

This tool aims to **prioritize patients** for rapid review, especially
where waiting times are long.
"""
        else:
            interpretation += """
### ⚠️ Possible Pneumonia

The chest X-ray may show **subtle or early pneumonia-like changes**.

**Clinical suggestions:**
- Evaluate symptoms and vital signs
- Consider repeat imaging or further labs
- Use local pneumonia treatment guidelines if diagnosis is confirmed
"""

    elif pred_label == "COVID-19":
        if confidence >= 85:
            interpretation += """
### ⚠️ High-Confidence COVID-19 Pneumonia Pattern

The AI sees a pattern often associated with **COVID-19 pneumonia**.

**Clinical next steps typically include:**
- **SARS-CoV-2 testing** (RT-PCR or antigen)
- Isolation and infection prevention
- Monitoring oxygen saturation (SpO2)
- Urgent care if:
  - SpO2 < 94%
  - Respiratory distress
  - Persistent chest pain or confusion

Imaging alone **cannot confirm COVID-19**. Lab testing + clinical judgment are essential.
"""
        else:
            interpretation += """
### ⚠️ Possible COVID-19 Pattern

There are features that *could* be compatible with COVID-19, but the AI is not very certain.

**Clinical suggestions:**
- Follow local COVID-19 testing protocols
- Use symptoms and exposure history
- Monitor for deterioration and hypoxia
"""

    else:  # Normal
        if confidence >= 85:
            interpretation += """
### ✅ High-Confidence "No Major Abnormality" Pattern

The model does **not** see strong evidence of TB, pneumonia, or COVID-19.

This may support a **normal chest X-ray**, but:

- Early disease can be radiographically subtle
- Some lung or cardiac diseases are **not detectable** here
- Symptoms always override AI reassurance

If a patient is symptomatic, clinical review is still required.
"""
        else:
            interpretation += """
### ⚠️ Likely Normal, But With Low Confidence

The model leans toward a **normal** study, but uncertainty is higher than usual.

- If the patient is unwell, treat this as **inconclusive**
- Consider follow-up imaging or alternative diagnostics
"""

    interpretation += """
---

## 🌍 Built to Help Humanity

This AI system is being developed to:

- Support **front-line clinicians** in low-resource and high-burden regions  
- Provide an **energy-efficient (Adaptive Sparse Training)** screening assistant  
- Help triage patients when **radiologists are not immediately available**

It is **open research**, not a commercial product, and we welcome
feedback from clinicians, researchers, and public health teams.

---

## ⚠️ Critical Medical Disclaimer

- This is a **screening and research tool only** – **NOT** an FDA/CE approved device.
- It does **not** replace radiologists, pulmonologists, or infectious disease experts.
- All decisions about diagnosis and treatment must be made by qualified clinicians.
- Gold-standard confirmation remains:
  - **TB** – sputum tests, culture, GeneXpert, TB-PCR
  - **Pneumonia** – full clinical assessment + labs/imaging
  - **COVID-19** – RT-PCR / validated antigen testing

If there is any doubt, always follow local clinical guidelines and consult a specialist.
"""

    return interpretation

# ============================================================================
# Gradio Interface
# ============================================================================

custom_css = """
#main-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    padding: 20px;
}
#title {
    text-align: center;
    color: white;
    font-size: 2.5em;
    font-weight: 800;
    margin-bottom: 10px;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.35);
}
#subtitle {
    text-align: center;
    color: #f5f5ff;
    font-size: 1.1em;
    margin-bottom: 12px;
}
#mission {
    text-align: center;
    color: #ffffff;
    font-size: 0.95em;
    margin-bottom: 24px;
    padding: 14px 18px;
    background: rgba(0,0,0,0.15);
    border-radius: 12px;
    backdrop-filter: blur(12px);
}
#stats {
    text-align: center;
    color: #fff;
    font-size: 0.95em;
    margin-bottom: 30px;
    padding: 12px 16px;
    background: rgba(255,255,255,0.08);
    border-radius: 10px;
}
.gradio-container {
    font-family: "Inter", system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
}
#upload-box {
    border: 3px dashed #667eea;
    border-radius: 15px;
    padding: 20px;
    background: rgba(255,255,255,0.97);
}
#results-box {
    background: white;
    border-radius: 15px;
    padding: 20px;
    box-shadow: 0 4px 12px rgba(0,0,0,0.12);
}
.output-image {
    border-radius: 10px;
    box-shadow: 0 2px 6px rgba(0,0,0,0.15);
}
footer {
    text-align: center;
    margin-top: 30px;
    color: white;
    font-size: 0.9em;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.HTML(
        """
        <div id="main-container">
            <div id="title">🫁 Global Chest X-Ray Screening AI</div>
            <div id="subtitle">
                4-Class detection • Explainable AI • Adaptive Sparse Training
            </div>
            <div id="mission">
                <b>Mission:</b> Support clinicians and patients worldwide – especially in
                low-resource, high-burden regions – by providing an energy-efficient AI
                assistant for chest X-ray screening. This is a <b>second opinion</b> tool,
                not a replacement for human experts.
            </div>
            <div id="stats">
                <b>Trained on 4 classes:</b> Normal • Tuberculosis • Pneumonia • COVID-19<br/>
                <b>Energy-efficient:</b> Adaptive Sparse Training (AST) – ~89% compute savings (research setting)<br/>
                <b>Use case:</b> Triage & screening support for TB, pneumonia, and COVID-19 programs
            </div>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1, elem_id="upload-box"):
            gr.Markdown("## 📤 Upload Chest X-Ray")
            image_input = gr.Image(
                type="pil",
                label="Upload X-Ray Image (PA or AP view)",
                elem_classes="output-image",
            )

            show_gradcam = gr.Checkbox(
                value=True,
                label="Enable Grad-CAM (Explainable AI)",
                info="Shows which lung regions the model is focusing on.",
            )

            analyze_btn = gr.Button("🔬 Analyze X-Ray", variant="primary", size="lg")

            gr.Markdown(
                """
### 📋 Supported Images
- Chest X-rays (PA or AP view)
- PNG / JPG / JPEG
- Grayscale or RGB

### 💡 Designed For
- TB & pneumonia screening programs
- Remote / low-resource clinics
- Educational and research use

> ⚠️ Always combine AI output with clinical judgment and lab tests.
                """
            )

        with gr.Column(scale=2, elem_id="results-box"):
            gr.Markdown("## 📊 AI Analysis Results")

            with gr.Row():
                prob_output = gr.Label(
                    label="Prediction Confidence (per class)",
                    num_top_classes=4,
                )

            with gr.Tabs():
                with gr.Tab("Original (Annotated)"):
                    original_output = gr.Image(
                        label="Annotated X-Ray",
                        elem_classes="output-image",
                    )

                with gr.Tab("Grad-CAM Heatmap"):
                    gradcam_output = gr.Image(
                        label="Model Attention Heatmap",
                        elem_classes="output-image",
                    )

                with gr.Tab("Overlay"):
                    overlay_output = gr.Image(
                        label="Explainable AI Overlay",
                        elem_classes="output-image",
                    )

            interpretation_output = gr.Markdown(label="Clinical-Style Interpretation")

    gr.Markdown("## 📁 Example X-Rays (for testing only – not real patients)")
    gr.Examples(
        examples=[
            ["examples/normal.png"],
            ["examples/tb.png"],
            ["examples/pneumonia.png"],
            ["examples/covid.png"],
        ],
        inputs=image_input,
        label="Click an example to load it into the app",
    )

    analyze_btn.click(
        fn=predict_chest_xray,
        inputs=[image_input, show_gradcam],
        outputs=[
            prob_output,
            original_output,
            gradcam_output,
            overlay_output,
            interpretation_output,
        ],
    )

    gr.HTML(
        """
        <footer>
            <p>
                <b>🫁 Global Chest X-Ray Screening with Adaptive Sparse Training</b><br/>
                Built as open research to support clinicians and public health teams worldwide.<br/>
                Not a medical device • Not for autonomous diagnosis or treatment decisions.
            </p>
            <p style="font-size: 0.8em; margin-top: 12px;">
                ⚠️ <b>MEDICAL DISCLAIMER:</b> This tool is for research and educational use only.
                All findings must be confirmed by qualified medical professionals using
                appropriate clinical and laboratory standards.
            </p>
        </footer>
        """
    )

if __name__ == "__main__":
    demo.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )