"""
Chest X-Ray Classifier with Grad-CAM Explainability
"""

import traceback
import sys

print("=" * 50)
print("STEP 1: Starting imports...")
sys.stdout.flush()

try:
    import torch
    print("  torch imported OK")
    import numpy as np
    print("  numpy imported OK")
    import gradio as gr
    print("  gradio imported OK")
    from PIL import Image
    print("  PIL imported OK")
    from torchvision import transforms
    print("  torchvision transforms imported OK")
    import timm
    print("  timm imported OK")
    from pytorch_grad_cam import GradCAM
    print("  GradCAM imported OK")
    from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
    print("  ClassifierOutputTarget imported OK")
    from pytorch_grad_cam.utils.image import show_cam_on_image
    print("  show_cam_on_image imported OK")
except Exception as e:
    print(f"IMPORT ERROR: {e}")
    traceback.print_exc()
    sys.stdout.flush()
    raise

sys.stdout.flush()

# CONFIG
CHEXPERT_LABELS = [
    "No Finding",
    "Enlarged Cardiomediastinum",
    "Cardiomegaly",
    "Lung Opacity",
    "Lung Lesion",
    "Edema",
    "Consolidation",
    "Pneumonia",
    "Atelectasis",
    "Pneumothorax",
    "Pleural Effusion",
    "Pleural Other",
    "Fracture",
    "Support Devices",
]

DEVICE = torch.device("cpu")

# PREPROCESS
preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])


def load_model():
    """Load DenseNet-121."""
    print("STEP 2: Loading model...")
    sys.stdout.flush()
    model = timm.create_model("densenet121", pretrained=True, num_classes=14)
    model = model.to(DEVICE)
    model.eval()
    print("STEP 3: Model loaded OK")
    sys.stdout.flush()
    return model


def classify_and_explain(image):
    """Run classification and Grad-CAM."""
    try:
        if image is None:
            return None, "Please upload a chest X-ray image."

        print("INFERENCE: Starting...")
        sys.stdout.flush()

        image_rgb = image.convert("RGB")
        input_tensor = preprocess(image_rgb).unsqueeze(0).to(DEVICE)
        rgb_display = np.array(image_rgb.resize((224, 224))).astype(np.float32) / 255.0

        with torch.no_grad():
            logits = model(input_tensor)
            probabilities = torch.sigmoid(logits).cpu().numpy()[0]

        top_idx = int(np.argmax(probabilities))
        top_label = CHEXPERT_LABELS[top_idx]
        top_conf = float(probabilities[top_idx])

        print(f"INFERENCE: Top prediction = {top_label} ({top_conf:.1%})")
        sys.stdout.flush()

        # Grad-CAM
        print("INFERENCE: Computing Grad-CAM...")
        sys.stdout.flush()
        cam = GradCAM(model=model, target_layers=target_layers)
        targets = [ClassifierOutputTarget(top_idx)]
        grayscale_cam = cam(input_tensor=input_tensor, targets=targets)[0]

        print("INFERENCE: Creating overlay...")
        sys.stdout.flush()
        heatmap_overlay = show_cam_on_image(rgb_display, grayscale_cam, use_rgb=True)

        print("INFERENCE: Done!")
        sys.stdout.flush()

        # Results
        sorted_indices = np.argsort(probabilities)[::-1]
        result_md = f"## Top Prediction: **{top_label}** ({top_conf:.1%})\n\n"
        result_md += "### All Condition Probabilities\n\n"
        result_md += "| Rank | Condition | Probability | Confidence |\n"
        result_md += "|:----:|-----------|:-----------:|:----------:|\n"

        for rank, idx in enumerate(sorted_indices, 1):
            label = CHEXPERT_LABELS[idx]
            prob = float(probabilities[idx])
            bar_len = int(prob * 20)
            bar = "X" * bar_len + "-" * (20 - bar_len)
            emoji = "HIGH" if prob > 0.5 else "MED" if prob > 0.3 else "LOW"
            result_md += f"| {rank} | {label} | {prob:.1%} | {emoji} `{bar}` |\n"

        result_md += "\n---\n"
        result_md += f"**Grad-CAM target:** {top_label} - red regions indicate areas the model focused on.\n"

        return heatmap_overlay, result_md

    except Exception as e:
        error_msg = f"Error during inference: {e}\n{traceback.format_exc()}"
        print(error_msg)
        sys.stdout.flush()
        return None, f"**Error:** {e}"


# BUILD UI
print("STEP 4: Building Gradio UI...")
sys.stdout.flush()

with gr.Blocks(
    title="Chest X-Ray Classifier with Grad-CAM",
    theme=gr.themes.Soft(),
) as demo:

    gr.Markdown(
        """
        # Chest X-Ray Classifier with Grad-CAM Explainability

        Upload a chest X-ray to classify **14 thoracic conditions** and see
        **where the model is looking** via Grad-CAM heatmaps.

        > **Disclaimer:** This is an educational tool only. NOT a medical device.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="Upload Chest X-Ray", height=400)
            classify_btn = gr.Button("Classify and Explain", variant="primary", size="lg")
        with gr.Column(scale=1):
            output_image = gr.Image(label="Grad-CAM Heatmap Overlay", height=400)
            output_text = gr.Markdown(value="*Upload an image and click Classify to see results.*")

    with gr.Accordion("How Grad-CAM Works", open=False):
        gr.Markdown(
            """
            **Grad-CAM** highlights which parts of the image the model focused on.

            - **Red** = High attention
            - **Yellow** = Moderate attention
            - **Green/Blue** = Low attention

            Reference: Selvaraju et al., "Grad-CAM", ICCV 2017.
            """
        )

    classify_btn.click(fn=classify_and_explain, inputs=[input_image], outputs=[output_image, output_text])

# LOAD MODEL AND LAUNCH
print("STEP 5: Loading model at startup...")
sys.stdout.flush()

model = load_model()
target_layers = [model.features[-1]]

print("STEP 6: Launching server...")
sys.stdout.flush()

demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)