import gradio as gr
import torch
import torch.nn as nn
from torchvision import models
import numpy as np
import cv2
import albumentations as A
from albumentations.pytorch import ToTensorV2

# ── Model Definition ──────────────────────────────────────────
class HybridSpineClassifier(nn.Module):
    def __init__(self, num_targets=25, num_classes=3, dropout=0.3):
        super().__init__()
        backbone = models.mobilenet_v2(weights=None)
        self.features = backbone.features
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
            nn.Dropout(p=dropout),
            nn.Linear(1280, 512),
            nn.ReLU(),
            nn.Dropout(p=dropout / 2),
            nn.Linear(512, num_targets * num_classes)
        )
        self.num_targets = num_targets
        self.num_classes = num_classes

    def forward(self, x):
        x = self.features(x)
        x = self.pool(x)
        x = x.flatten(1)
        x = self.classifier(x)
        return x.view(-1, self.num_targets, self.num_classes)

# ── Constants ─────────────────────────────────────────────────
LEVELS = ['L1/L2', 'L2/L3', 'L3/L4', 'L4/L5', 'L5/S1']
CONDITIONS = [
    'Spinal Canal Stenosis',
    'Left Neural Foraminal Narrowing',
    'Right Neural Foraminal Narrowing',
    'Left Subarticular Stenosis',
    'Right Subarticular Stenosis',
]
CLASS_NAMES = ['Normal/Mild', 'Moderate', 'Severe']

val_transforms = A.Compose([
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2(),
])

# ── Load Model ────────────────────────────────────────────────
model = HybridSpineClassifier(num_targets=25, num_classes=3)
model.load_state_dict(torch.load("original_model.pth", map_location='cpu'))
model.eval()
print("Model loaded!")

# ── Inference ─────────────────────────────────────────────────
def analyze_spine(image):
    if image is None:
        return "Please upload an MRI slice."

    # Preprocess
    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    else:
        gray = image

    gray = cv2.resize(gray, (224, 224))
    img_3ch = np.stack([gray, gray, gray], axis=-1)
    tensor = val_transforms(image=img_3ch)['image'].unsqueeze(0)

    # Classify
    with torch.no_grad():
        logits = model(tensor)
        probs = torch.softmax(logits, dim=-1).squeeze().numpy()

    # Build results
    lines = []
    severe_count = 0
    moderate_count = 0

    lines.append("=" * 55)
    lines.append("  LUMBAR SPINE ANALYSIS RESULTS")
    lines.append("=" * 55)
    lines.append("")

    for ci, condition in enumerate(CONDITIONS):
        lines.append(f"  {condition}")
        lines.append(f"  {'─' * 45}")
        for li, level in enumerate(LEVELS):
            idx = ci * 5 + li
            pred_class = int(np.argmax(probs[idx]))
            confidence = float(np.max(probs[idx]))
            severity = CLASS_NAMES[pred_class]

            if severity == "Severe":
                marker = "🔴"
                severe_count += 1
            elif severity == "Moderate":
                marker = "🟡"
                moderate_count += 1
            else:
                marker = "🟢"

            lines.append(f"    {marker} {level:8s}  {severity:12s} ({confidence*100:.1f}%)")
        lines.append("")

    # Risk assessment
    if severe_count >= 3:
        risk = "HIGH"
    elif severe_count >= 1 or moderate_count >= 5:
        risk = "MODERATE"
    else:
        risk = "LOW"

    lines.append("=" * 55)
    lines.append(f"  RISK LEVEL: {risk}")
    lines.append(f"  Severe: {severe_count}/25  |  Moderate: {moderate_count}/25")
    lines.append("=" * 55)
    lines.append("")

    if risk == "HIGH":
        lines.append("  ⚠ Significant degenerative findings detected.")
        lines.append("  Specialist referral recommended.")
    elif risk == "MODERATE":
        lines.append("  Moderate degenerative changes detected.")
        lines.append("  Clinical correlation and follow-up recommended.")
    else:
        lines.append("  Predominantly normal findings.")
        lines.append("  Routine monitoring recommended.")

    lines.append("")
    lines.append("  Note: AI-assisted screening tool.")
    lines.append("  Not a substitute for professional diagnosis.")

    return "\n".join(lines)


# ── Gradio Interface ──────────────────────────────────────────
demo = gr.Interface(
    fn=analyze_spine,
    inputs=gr.Image(label="Upload MRI Slice (PNG or JPG)"),
    outputs=gr.Textbox(label="Analysis Results", lines=40),
    title="Automated Lumbar Spine Analysis",
    description="Upload a lumbar spine MRI slice to get automated severity classification across 5 conditions and 5 vertebral levels. Powered by MobileNetV2 trained on RSNA 2024 data.",
    examples=None,
    theme="default",
)

demo.launch()