import os
import zipfile
import numpy as np
import torch
from transformers import SegformerForImageSegmentation, ResNetForImageClassification, AdamW
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import streamlit as st
import gradio as gr

gr.load("models/nvidia/segformer-b0-finetuned-ade-512-512").launch()

# Function to extract zip files
def extract_zip(zip_file, extract_to):
    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
        zip_ref.extractall(extract_to)

# Preprocess images
def preprocess_image(image_path):
    ext = os.path.splitext(image_path)[-1].lower()

    if ext == '.npy':
        image_data = np.load(image_path)
        image_tensor = torch.tensor(image_data).float()
        if len(image_tensor.shape) == 3:
            image_tensor = image_tensor.unsqueeze(0)

    elif ext in ['.jpg', '.jpeg']:
        img = Image.open(image_path).convert('RGB').resize((224, 224))
        img_np = np.array(img)
        image_tensor = torch.tensor(img_np).permute(2, 0, 1).float()

    else:
        raise ValueError(f"Unsupported format: {ext}")

    image_tensor /= 255.0  # Normalize to [0, 1]
    return image_tensor

# Prepare dataset
def prepare_dataset(extracted_folder):
    neuronii_path = os.path.join(extracted_folder, "neuroniiimages")
    
    if not os.path.exists(neuronii_path):
        raise FileNotFoundError(f"The folder neuroniiimages does not exist in the extracted folder: {neuronii_path}")
    
    image_paths = []
    labels = []
    
    for disease_folder in ['alzheimers_dataset', 'parkinsons_dataset', 'MSjpg']:
        folder_path = os.path.join(neuronii_path, disease_folder)
        
        if not os.path.exists(folder_path):
            print(f"Folder not found: {folder_path}")
            continue  
        label = {'alzheimers_dataset': 0, 'parkinsons_dataset': 1, 'MSjpg': 2}[disease_folder]
        
        for img_file in os.listdir(folder_path):
            if img_file.endswith(('.npy', '.jpg', '.jpeg')):
                image_paths.append(os.path.join(folder_path, img_file))
                labels.append(label)
            else:
                print(f"Unsupported file: {img_file}")
    print(f"Total images loaded: {len(image_paths)}")
    return image_paths, labels

# Custom Dataset class
class CustomImageDataset(Dataset):
    def __init__(self, image_paths, labels):
        self.image_paths = image_paths
        self.labels = labels

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

    def __getitem__(self, idx):
        image = preprocess_image(self.image_paths[idx])
        label = self.labels[idx]
        return image, label

# Training function for classification
def fine_tune_classification_model(train_loader):
    model = ResNetForImageClassification.from_pretrained('microsoft/resnet-50', num_labels=3)
    model.train()
    optimizer = AdamW(model.parameters(), lr=1e-4)
    criterion = torch.nn.CrossEntropyLoss()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    
    for epoch in range(10):
        running_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(pixel_values=images).logits
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
    return running_loss / len(train_loader)

# Streamlit UI for Fine-tuning
st.title("Fine-tune ResNet for MRI/CT Scans Classification")

zip_file_url = "https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/neuroniiimages.zip"

if st.button("Start Training"):
    extraction_dir = "extracted_files"
    os.makedirs(extraction_dir, exist_ok=True)

    # Download the zip file (placeholder)
    zip_file = "neuroniiimages.zip"  # Assuming you downloaded it with this name

    # Extract zip file
    extract_zip(zip_file, extraction_dir)

    # Prepare dataset
    image_paths, labels = prepare_dataset(extraction_dir)
    dataset = CustomImageDataset(image_paths, labels)
    train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

    # Fine-tune the classification model
    final_loss = fine_tune_classification_model(train_loader)
    st.write(f"Training Complete with Final Loss: {final_loss}")

# Segmentation function (using SegFormer)
def fine_tune_segmentation_model(train_loader):
    model = SegformerForImageSegmentation.from_pretrained('nvidia/segformer-b0', num_labels=3)
    model.train()
    optimizer = AdamW(model.parameters(), lr=1e-4)
    criterion = torch.nn.CrossEntropyLoss()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    
    for epoch in range(10):
        running_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(pixel_values=images).logits
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
    return running_loss / len(train_loader)

# Add a button for segmentation training
if st.button("Start Segmentation Training"):
    # Assuming the dataset for segmentation is prepared similarly
    seg_train_loader = DataLoader(dataset, batch_size=32, shuffle=True)
    
    # Fine-tune the segmentation model
    final_loss_seg = fine_tune_segmentation_model(seg_train_loader)
    st.write(f"Segmentation Training Complete with Final Loss: {final_loss_seg}")