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

# 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}")