# Install the ultralytics package for YOLOv8 (if not already installed)
!pip install ultralytics

# Import necessary libraries
import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from tqdm.notebook import tqdm
from torchvision import transforms
from PIL import Image
import torch.nn as nn

# Suppress warnings from YOLO and transformers
import warnings
warnings.filterwarnings("ignore")

# # Mount Google Drive if needed
# from google.colab import drive
# drive.mount('/content/drive')

# Load the CSV file
csv_file = '/kaggle/input/scenario32/scenario32_dev.csv'  # Replace with your actual path
df = pd.read_csv(csv_file)

# Import YOLOv8 model
from ultralytics import YOLO

# Suppress YOLO logging output
import logging
logging.getLogger('ultralytics').setLevel(logging.ERROR)

# Load the pre-trained YOLOv8 segmentation model
yolo_model = YOLO('yolov8l-seg.pt')  # Using 'n'ano model for speed and low resource consumption

# Create a custom dataset class
class BeamDataset(Dataset):
    def __init__(self, dataframe, transform=None, label_mapping=None):
        self.data = dataframe.reset_index(drop=True)
        self.image_paths = self.data['scenario32/'+'unit1_rgb']
        self.labels = self.data['unit1_beam']
        self.transform = transform
        self.label_mapping = label_mapping

        # Map labels to indices
        self.labels = self.labels.map(self.label_mapping).astype(int)

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

    def __getitem__(self, idx):
        img_path = self.image_paths.iloc[idx]

        # Load image
        image = Image.open(img_path).convert('RGB')

        # Perform YOLOv8 segmentation without verbose output
        results = yolo_model(img_path, verbose=False)
        # Get the segmentation mask
        masks = results[0].masks.data.cpu().numpy()  # Shape: (num_objects, H, W)

        # If no objects detected, create an empty mask
        if len(masks) == 0:
            seg_mask = np.zeros((image.height, image.width), dtype=np.uint8)
        else:
            # Combine all masks into a single mask
            seg_mask = np.sum(masks, axis=0)
            seg_mask = np.clip(seg_mask, 0, 1).astype(np.uint8)
            # Resize to match image size if necessary
            seg_mask = Image.fromarray(seg_mask * 255).resize(image.size)

        # Convert segmentation mask to tensor
        if self.transform:
            image = self.transform(image)
            seg_mask = self.transform(seg_mask.convert('RGB'))
        else:
            transform = transforms.ToTensor()
            image = transform(image)
            seg_mask = transform(seg_mask.convert('RGB'))

        label = self.labels.iloc[idx]
        return {'image': image, 'seg_mask': seg_mask, 'label': label}

# Define the transforms
from transformers import ViTImageProcessor

# Use a smaller model for practical GPU memory usage
model_name = 'google/vit-large-patch16-224-in21k'
image_processor = ViTImageProcessor.from_pretrained(model_name)

common_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
])

# Create label mapping
all_labels = df['unit1_beam'].unique()
label_mapping = {label: idx for idx, label in enumerate(sorted(all_labels))}
num_labels = len(label_mapping)

# Split the dataset into training and testing sets
from sklearn.model_selection import train_test_split

train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)

train_dataset = BeamDataset(train_df, transform=common_transforms, label_mapping=label_mapping)
test_dataset = BeamDataset(test_df, transform=common_transforms, label_mapping=label_mapping)

train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=8, shuffle=False, num_workers=0)

# Load the ViT model
from transformers import ViTModel

# Load pre-trained ViT models for images and segmentation masks
vit_image = ViTModel.from_pretrained(model_name)
vit_segmentation = ViTModel.from_pretrained(model_name)

# Freeze ViT models if desired
# for param in vit_image.parameters():
#     param.requires_grad = False
# for param in vit_segmentation.parameters():
#     param.requires_grad = False

# Define the combined model
class BeamPredictionModel(nn.Module):
    def __init__(self, num_labels):
        super(BeamPredictionModel, self).__init__()
        self.vit_image = vit_image
        self.vit_segmentation = vit_segmentation
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(vit_image.config.hidden_size * 2, num_labels)

    def forward(self, image, seg_mask):
        # Pass through ViT models
        image_outputs = self.vit_image(image)
        seg_outputs = self.vit_segmentation(seg_mask)

        # Concatenate embeddings
        combined_representation = torch.cat(
            (image_outputs.pooler_output, seg_outputs.pooler_output), dim=1
        )

        # Classification
        x = self.dropout(combined_representation)
        logits = self.classifier(x)
        return logits

# Initialize the model
num_labels = len(label_mapping)

model = BeamPredictionModel(num_labels)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

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

# Use mixed precision
scaler = torch.cuda.amp.GradScaler()

# Import matplotlib for plotting
import matplotlib.pyplot as plt

# Initialize lists to store metrics
train_losses = []
val_losses = []
val_top1_accuracies = []
val_top2_accuracies = []
val_top3_accuracies = []
val_dba_scores = []

# Training loop
num_epochs = 5  # Adjust the number of epochs as needed

for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}")

    for batch in progress_bar:
        images = batch['image'].to(device)
        seg_masks = batch['seg_mask'].to(device)
        labels = batch['label'].to(device)

        optimizer.zero_grad()
        with torch.cuda.amp.autocast():
            logits = model(images, seg_masks)
            loss = criterion(logits, labels)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        total_loss += loss.item()
        progress_bar.set_postfix(loss=loss.item())

    avg_loss = total_loss / len(train_loader)
    train_losses.append(avg_loss)
    print(f"Epoch [{epoch+1}/{num_epochs}], Training Loss: {avg_loss:.4f}")

    # Evaluate on the test set
    model.eval()
    total_correct_top1 = 0
    total_correct_top2 = 0
    total_correct_top3 = 0
    total_samples = 0
    total_dba = 0
    total_val_loss = 0

    with torch.no_grad():
        for batch in test_loader:
            images = batch['image'].to(device)
            seg_masks = batch['seg_mask'].to(device)
            labels = batch['label'].to(device)

            logits = model(images, seg_masks)
            val_loss = criterion(logits, labels)
            total_val_loss += val_loss.item()

            # Get top-k predictions
            _, preds_topk = logits.topk(3, dim=1)
            correct_topk = preds_topk.eq(labels.view(-1, 1))  # Shape: (batch_size, 3)

            # Compute Top-1 accuracy
            total_correct_top1 += correct_topk[:, 0].sum().item()

            # Compute Top-2 accuracy
            total_correct_top2 += correct_topk[:, :2].any(dim=1).sum().item()

            # Compute Top-3 accuracy
            total_correct_top3 += correct_topk.any(dim=1).sum().item()

            # Compute DBA score
            probs = torch.softmax(logits, dim=1)
            batch_size = labels.size(0)
            for i in range(batch_size):
                prob = probs[i]
                label = labels[i]
                sorted_probs, indices = prob.sort(descending=True)
                rank = (indices == label).nonzero(as_tuple=True)[0].item() + 1  # Rank starts from 1
                total_dba += 1.0 / rank

            total_samples += labels.size(0)

    avg_val_loss = total_val_loss / len(test_loader)
    val_losses.append(avg_val_loss)

    top1_accuracy = total_correct_top1 / total_samples
    top2_accuracy = total_correct_top2 / total_samples
    top3_accuracy = total_correct_top3 / total_samples
    avg_dba = total_dba / total_samples

    val_top1_accuracies.append(top1_accuracy)
    val_top2_accuracies.append(top2_accuracy)
    val_top3_accuracies.append(top3_accuracy)
    val_dba_scores.append(avg_dba)

    print(f"Validation Loss: {avg_val_loss:.4f}")
    print(f"Validation Top-1 Accuracy: {top1_accuracy:.4f}")
    print(f"Validation Top-2 Accuracy: {top2_accuracy:.4f}")
    print(f"Validation Top-3 Accuracy: {top3_accuracy:.4f}")
    print(f"Validation DBA Score: {avg_dba:.4f}")

# Plotting Training Loss over Epochs
plt.figure(figsize=(8, 6))
plt.plot(range(1, num_epochs+1), train_losses, label='Training Loss')
plt.plot(range(1, num_epochs+1), val_losses, label='Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training and Validation Loss over Epochs')
plt.legend()
plt.grid(True)
plt.show()

# Plotting Top-1, Top-2, Top-3 Accuracy over Epochs
plt.figure(figsize=(8, 6))
plt.plot(range(1, num_epochs+1), val_top1_accuracies, label='Top-1 Accuracy')
plt.plot(range(1, num_epochs+1), val_top2_accuracies, label='Top-2 Accuracy')
plt.plot(range(1, num_epochs+1), val_top3_accuracies, label='Top-3 Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.title('Top-1, Top-2, Top-3 Accuracy over Epochs')
plt.legend()
plt.grid(True)
plt.show()

# Plotting DBA Score over Epochs
plt.figure(figsize=(8, 6))
plt.plot(range(1, num_epochs+1), val_dba_scores, label='DBA Score', marker='o')
plt.xlabel('Epoch')
plt.ylabel('DBA Score')
plt.title('DBA Score over Epochs')
plt.legend()
plt.grid(True)
plt.show()

what is error in this code