import torch
import tarfile
import os

def load_tar_pth_dataset(file_path):
    """Load a .tar.pth file and extract the contained .pth dataset."""
    try:
        # Open the tar file
        with tarfile.open(file_path, 'r') as archive:
            # List contents of the archive to verify the files inside
            archive_contents = archive.getnames()
            print(f"Contents of the archive: {archive_contents}")
            
            # Extract all files to a temporary folder
            archive.extractall('temp_dataset')
            
            # Assuming the tar file contains a .pth file (e.g., 'dataset.pth')
            pth_file_path = 'temp_dataset/dataset.pth'
            if os.path.exists(pth_file_path):
                # Load the .pth file containing the dataset
                dataset = torch.load(pth_file_path)
                images = dataset['images']  # Tensor of shape (N, C, H, W)
                labels = dataset['labels']  # Tensor of shape (N,)
                return images, labels
            else:
                print(f"Error: {pth_file_path} not found in the archive.")
                return None, None
    except Exception as e:
        print(f"Error loading tar file: {e}")
        return None, None


import torch
import random

class MemoryBuffer:
    def __init__(self, capacity):
        self.capacity = capacity
        self.buffer = []
    
    def add_samples(self, images, labels):
        """Add images and labels to the memory buffer."""
        for image, label in zip(images, labels):
            if len(self.buffer) >= self.capacity:
                self.buffer.pop(0)  # Remove the oldest sample if buffer is full
            self.buffer.append((image, label))
    
    def get_replay_data(self):
        """Randomly sample data from the buffer for replay."""
        if len(self.buffer) == 0:
            return [], []  # Return empty if no data
        sampled_data = random.sample(self.buffer, len(self.buffer))
        images, labels = zip(*sampled_data)
        return torch.stack(images), torch.tensor(labels)


import torch

import tarfile
import os
from torchvision import models, transforms
from torch.utils.data import DataLoader, TensorDataset
from sklearn.mixture import GaussianMixture
from sklearn.metrics import accuracy_score
import numpy as np

# Load a .tar.pth dataset from the provided tarball
# def load_tar_pth_dataset(file_path):
#     with tarfile.open(file_path, 'r') as archive:
#         archive.extractall('temp_dataset')
#         # Assuming the dataset contains a .pth file with 'images' and 'labels' tensors
#         dataset = torch.load('temp_dataset/dataset.pth')
#     images = dataset['images']  # Tensor of shape (N, C, H, W)
#     labels = dataset['labels']  # Tensor of shape (N,)
#     return images, labels

# Define a function for extracting features using ResNet50
def extract_features_resnet50(images):
    # Preprocessing pipeline for ResNet50
    preprocess = transforms.Compose([
        transforms.ToPILImage(),  # Convert tensor to PIL Image
        transforms.Resize(256),  # Resize to 256x256
        transforms.CenterCrop(224),  # Crop to 224x224
        transforms.ToTensor(),  # Convert back to tensor
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize
    ])

    # Load pre-trained ResNet50 model
    resnet50 = models.resnet50(pretrained=True)
    resnet50.eval()  # Set model to evaluation mode

    # Extract features from each image
    features = []
    for image in images:
        image = preprocess(image)
        image = image.unsqueeze(0)  # Add batch dimension
        with torch.no_grad():
            feature = resnet50(image)
        features.append(feature.squeeze().cpu().numpy())  # Remove batch dimension and convert to numpy
    return np.array(features)

# Initialize memory buffer
BUFFER_SIZE = 10000  # Set your buffer size
memory_buffer = MemoryBuffer(capacity=BUFFER_SIZE)

# Initialize internal distribution
internal_gmm = None

# Accuracy matrices to store the evaluation results
accuracy_matrix_task1 = np.zeros((10, 10))  # Task 1 accuracy matrix
accuracy_matrix_task2 = np.zeros((10, 20))  # Task 2 accuracy matrix

# Task 1: Handle D1 to D10
for i in range(1, 11):
    # Load dataset D_i from .tar.pth (i_train_data.tar.pth)
    dataset = load_tar_pth_dataset(f"{i}_train_data.tar.pth")
    images, labels = dataset
    features = extract_features_resnet50(images)
    
    if i == 1:
        # Train initial model f1 on D1
        model = LwPClassifier()
        model.train(features, labels)
        
        # Fit initial GMM
        internal_gmm = fit_gmm(features)
        
        # Update memory buffer
        memory_buffer.add_samples(images, labels)
    else:
        # Predict pseudo-labels on D_i using f_{i-1}
        pseudo_labels = model.predict(features)
        
        # Update model to f_i using pseudo-labeled D_i and memory buffer
        replay_features, replay_labels = memory_buffer.get_replay_data()
        model.update(features, pseudo_labels, replay_features, replay_labels)
        
        # Update internal GMM
        new_gmm = fit_gmm(features)
        internal_gmm = consolidate_gmm(internal_gmm, new_gmm)
        
        # Update memory buffer
        memory_buffer.add_samples(images, labels)
    
    # Evaluation on held-out datasets
    heldout_datasets = [load_tar_pth_dataset(f"hat_D{j}.tar.pth") for j in range(1, i+1)]
    accuracies = []
    for heldout in heldout_datasets:
        heldout_images, heldout_labels = heldout
        heldout_features = extract_features_resnet50(heldout_images)
        acc = accuracy_score(heldout_labels, model.predict(heldout_features))
        accuracies.append(acc)
    accuracy_matrix_task1[i-1, :i] = accuracies

# Task 2: Handle D11 to D20
model = f10  # Start with the last model from Task 1
for i in range(11, 21):
    # Load dataset D_i from .tar.pth (i_train_data.tar.pth)
    dataset = load_tar_pth_dataset(f"/content/{i}_train_data.tar.pth")
    images, labels = dataset
    features = extract_features_resnet50(images)
    
    # Predict pseudo-labels on D_i using f_{i-1}
    pseudo_labels = model.predict(features)
    
    # Align features with internal distribution
    aligned_features = align_features(features, internal_gmm)
    
    # Update model to f_i using pseudo-labeled D_i, aligned features, and memory buffer
    replay_features, replay_labels = memory_buffer.get_replay_data()
    model.update(aligned_features, pseudo_labels, replay_features, replay_labels, domain_adaptation=True)
    
    # Update internal GMM
    new_gmm = fit_gmm(aligned_features)
    internal_gmm = consolidate_gmm(internal_gmm, new_gmm)
    
    # Update memory buffer
    memory_buffer.add_samples(images, labels)
    
    # Evaluation on held-out datasets
    heldout_datasets = [load_tar_pth_dataset(f"hat_D{j}.tar.pth") for j in range(1, i+1)]
    accuracies = []
    for heldout in heldout_datasets:
        heldout_images, heldout_labels = heldout
        heldout_features = extract_features_resnet50(heldout_images)
        acc = accuracy_score(heldout_labels, model.predict(heldout_features))
        accuracies.append(acc)
    accuracy_matrix_task2[i-11, :i] = accuracies

# Save or display the accuracy matrices
def save_accuracy_matrix(accuracy_matrix, filename):
    np.savetxt(filename, accuracy_matrix, delimiter=",")

save_accuracy_matrix(accuracy_matrix_task1, "accuracy_task1.csv")
save_accuracy_matrix(accuracy_matrix_task2, "accuracy_task2.csv")
 correct this code