# ============================================================
# BEST FULL PIPELINE – Brain Stroke Classification
# Dataset: Tekno21 (Normal – Ischemic – Hemorrhagic)
# Model: EfficientNet (High accuracy)
# Includes: Training + Accuracy + Error Rate + Gradio UI
# ============================================================

!pip install -q datasets torch torchvision pillow gradio efficientnet_pytorch

from datasets import load_dataset
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from efficientnet_pytorch import EfficientNet
from PIL import Image
import gradio as gr
import numpy as np
from sklearn.metrics import accuracy_score

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Device:", device)

# ---------------------- LOAD DATASET ----------------------
ds = load_dataset("BTX24/tekno21-brain-stroke-dataset-multi")

labels_map = {"NORMAL": 0, "ISCHEMIC": 1, "HEMORRHAGIC": 2}

class_names = ["Normal", "Ischemic", "Hemorrhagic"]

# ---------------------- TRANSFORMS ----------------------
train_tf = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ColorJitter(0.2,0.2,0.2,0.1),
    transforms.ToTensor(),
    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
])

val_tf = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
])

# ---------------------- CUSTOM DATASET CLASS ----------------------
class BrainDataset(Dataset):
    def __init__(self, hf_data, transform):
        self.data = hf_data
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img = Image.open(self.data[idx]["image"]).convert("RGB")
        label = labels_map[self.data[idx]["label"].upper()]
        return self.transform(img), label

train_data = BrainDataset(ds["train"], train_tf)
val_data   = BrainDataset(ds["validation"], val_tf)

train_loader = DataLoader(train_data, batch_size=16, shuffle=True)
val_loader   = DataLoader(val_data, batch_size=16)

# ---------------------- MODEL ----------------------
model = EfficientNet.from_pretrained("efficientnet-b0")
model._fc = nn.Linear(model._fc.in_features, len(class_names))
model = model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-4)

# ---------------------- TRAINING ----------------------
EPOCHS = 5
best_acc = 0

for epoch in range(EPOCHS):
    model.train()
    train_correct = 0
    total = 0

    for imgs, labels in train_loader:
        imgs, labels = imgs.to(device), labels.to(device)
        optimizer.zero_grad()

        out = model(imgs)
        loss = criterion(out, labels)

        loss.backward()
        optimizer.step()

        _, preds = out.max(1)
        train_correct += preds.eq(labels).sum().item()
        total += labels.size(0)

    train_acc = 100 * train_correct / total
    print(f"Epoch {epoch+1}/{EPOCHS} – Train Acc: {train_acc:.2f}%")

    # Validation
    model.eval()
    val_preds = []
    val_true  = []

    with torch.no_grad():
        for imgs, labels in val_loader:
            imgs, labels = imgs.to(device), labels.to(device)
            out = model(imgs)
            _, preds = out.max(1)

            val_preds.extend(preds.cpu().numpy())
            val_true.extend(labels.cpu().numpy())

    val_acc = accuracy_score(val_true, val_preds) * 100
    print(f"Validation Accuracy: {val_acc:.2f}%")

    if val_acc > best_acc:
        best_acc = val_acc
        torch.save(model.state_dict(), "best_model.pth")
        print("✔ Best Model Saved")

print(" BEST ACCURACY =", best_acc)

# ---------------------- ERROR RATE ----------------------
error_rate = 100 - best_acc
print(" ERROR RATE =", error_rate, "%")

# ---------------------- LOAD BEST MODEL ----------------------
model.load_state_dict(torch.load("best_model.pth"))
model.eval()
# ---------------------- GRADIO INTERFACE ----------------------
def predict(img):
    img = val_tf(img).unsqueeze(0).to(device)
    with torch.no_grad():
        out = model(img)
        probs = torch.softmax(out[0], dim=0).cpu().numpy()
    return {class_names[i]: float(probs[i]) for i in range(len(class_names))}

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(),
    title="Brain Stroke Classifier (EfficientNet-B0)",
    description=f"Best Validation Accuracy: {best_acc:.2f}% | Error Rate: {error_rate:.2f}%"
)

iface.launch(share=True)