import torch
import zipfile
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms
from torchvision.models import ResNet18_Weights
from torch.utils.data import Dataset, DataLoader, ConcatDataset, Subset
import numpy as np
import random
import os
import tarfile
import pickle
from sklearn.metrics import accuracy_score
import copy
import pandas as pd

# ---------------------
# 1. Utility Functions
# ---------------------
def set_seed(seed=42):
    """Set random seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

set_seed()

# Device configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ---------------------
# 2. Custom Dataset Class
# ---------------------
class TarPTHDataset(Dataset):
    """
    Custom Dataset to load data from .tar.pth files.
    Assumes each .tar.pth file contains a dictionary with 'data' and 'labels' keys.
    - 'data': Tensor of shape (N, C, H, W)
    - 'labels': Tensor of shape (N,)
    """

    def __init__(self, tar_pth_path, transform=None):
        self.transform = transform
        self.data = []
        self.labels = []

        # Extract data from .tar.pth file
        with open(tar_pth_path, 'rb') as f:
            try:
                # Load the archive
                archive = torch.load(f, map_location='cpu')
            except Exception as e:
                print(f"Error loading {tar_pth_path}: {e}")
                raise e

        # Check the type and contents of the archive
        if isinstance(archive, dict):
            print(f"Loaded archive from {tar_pth_path} with keys: {list(archive.keys())}")
            self.data = archive.get('data', [])
            self.labels = archive.get('labels', [])
        else:
            print(f"Loaded archive from {tar_pth_path} is not a dict. Type: {type(archive)}")
            self.data = archive
            self.labels = [0] * len(self.data)  # Dummy labels if not provided

        # Print lengths for debugging
        print(f"Number of data samples: {len(self.data)}")
        print(f"Number of labels: {len(self.labels)}")

        # Assertions to ensure data integrity
        assert len(self.data) > 0, f"No data found in {tar_pth_path}."
        assert len(self.labels) == 0 or len(self.labels) == len(self.data), \
            f"Number of labels ({len(self.labels)}) does not match number of data samples ({len(self.data)})."

        # Convert to tensors if not already
        if not isinstance(self.data, torch.Tensor):
            self.data = torch.tensor(self.data)
        if not isinstance(self.labels, torch.Tensor):
            self.labels = torch.tensor(self.labels)

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

    def __getitem__(self, idx):
        image = self.data[idx]
        if len(self.labels) > 0:
            label = self.labels[idx]
        else:
            raise ValueError("Labels are empty. Cannot retrieve label.")

        if self.transform:
            image = self.transform(image)

        return image, label

# ---------------------
# 3. Feature Extractor
# ---------------------
class FeatureExtractor(nn.Module):
    """Pre-trained CNN for feature extraction."""

    def __init__(self, weights=ResNet18_Weights.IMAGENET1K_V1):
        super(FeatureExtractor, self).__init__()
        self.model = models.resnet18(weights=weights)
        # Remove the final classification layer
        self.features = nn.Sequential(*list(self.model.children())[:-1])

    def forward(self, x):
        with torch.no_grad():
            x = self.features(x)
            x = x.view(x.size(0), -1)  # Flatten
        return x

# ---------------------
# 4. Classifier with Knowledge Distillation
# ---------------------
class Classifier(nn.Module):
    """Classifier that can incorporate knowledge distillation from a previous model."""

    def __init__(self, input_dim, num_classes, previous_model=None, alpha=0.5):
        super(Classifier, self).__init__()
        self.fc = nn.Linear(input_dim, num_classes)
        self.previous_model = previous_model
        self.alpha = alpha  # Weight for distillation loss

    def forward(self, x):
        return self.fc(x)

# ---------------------
# 5. Memory Buffer for Experience Replay
# ---------------------
class MemoryBuffer:
    """Experience Replay Buffer to store representative samples from previous domains."""

    def __init__(self, capacity):
        self.capacity = capacity
        self.buffer = []

    def add_samples(self, samples, labels):
        for x, y in zip(samples, labels):
            if len(self.buffer) >= self.capacity:
                self.buffer.pop(0)  # Remove oldest samples
            self.buffer.append((x, y))

    def get_samples(self, batch_size):
        if len(self.buffer) == 0:
            return None, None
        selected = random.sample(self.buffer, min(batch_size, len(self.buffer)))
        images, labels = zip(*selected)
        return torch.stack(images), torch.tensor(labels)

# ---------------------
# 6. Random Mixup (RandMix)
# ---------------------
def randmix(images, alpha=0.2):
    """
    Applies mixup augmentation.
    Args:
        images: Batch of images.
        alpha: Mixup parameter.
    Returns:
        mixed_images: Augmented images.
    """
    lam = np.random.beta(alpha, alpha)
    index = torch.randperm(images.size(0))
    mixed_images = lam * images + (1 - lam) * images[index]
    return mixed_images, lam, index

# ---------------------
# 7. Pseudo-Labeling (T2PL)
# ---------------------
def generate_pseudo_labels(model, feature_extractor, dataloader, threshold=0.5):
    """
    Generates pseudo-labels for unlabeled data using high-confidence predictions.
    Args:
        model: Current classifier model.
        feature_extractor: Feature extractor model.
        dataloader: Dataloader for unlabeled data.
        threshold: Confidence threshold to select top 50% predictions.
    Returns:
        pseudo_images: List of images with high-confidence predictions.
        pseudo_labels: Corresponding pseudo-labels.
        confidences: Confidence scores of selected predictions.
    """
    model.eval()
    feature_extractor.eval()
    pseudo_images = []
    pseudo_labels = []
    confidences = []

    all_confidences = []
    all_preds = []
    all_images = []

    with torch.no_grad():
        for images, _ in dataloader:
            images = images.to(device)
            features = feature_extractor(images)
            outputs = model(features)
            probs = nn.functional.softmax(outputs, dim=1)
            max_conf, preds = torch.max(probs, dim=1)
            all_confidences.append(max_conf.cpu())
            all_preds.append(preds.cpu())
            all_images.append(images.cpu())

    all_confidences = torch.cat(all_confidences)
    all_preds = torch.cat(all_preds)
    all_images = torch.cat(all_images)

    # Select top 50% confident samples
    num_select = all_confidences.size(0) // 2
    top_conf, top_indices = torch.topk(all_confidences, num_select)
    selected_images = all_images[top_indices]
    selected_labels = all_preds[top_indices]

    return selected_images, selected_labels, top_conf

# ---------------------
# 8. Prototype Contrastive Alignment (PCA)
# ---------------------
class PrototypeContrastiveLoss(nn.Module):
    """
    Prototype Contrastive Loss to align class prototypes across domains.
    Encourages similarity between class prototypes and data samples,
    while separating different class prototypes.
    """

    def __init__(self, temperature=0.07):
        super(PrototypeContrastiveLoss, self).__init__()
        self.temperature = temperature
        self.cross_entropy = nn.CrossEntropyLoss()

    def forward(self, features, prototypes, labels):
        """
        Args:
            features: Feature representations of samples (batch_size, feature_dim)
            prototypes: Class prototypes (num_classes, feature_dim)
            labels: Ground truth labels (batch_size)
        Returns:
            loss: Contrastive loss value
        """
        # Normalize features and prototypes
        features = nn.functional.normalize(features, dim=1)
        prototypes = nn.functional.normalize(prototypes, dim=1)

        # Compute similarity
        logits = torch.matmul(features, prototypes.t()) / self.temperature
        loss = self.cross_entropy(logits, labels)
        return loss

# ---------------------
# 9. Training Function
# ---------------------
def train_model(
    model,
    feature_extractor,
    optimizer,
    criterion,
    dataloader,
    memory_buffer,
    previous_model=None,
    alpha=0.5,
    randmix_alpha=0.2,
    use_mixup=False
):
    """
    Trains the classifier model on the given data.
    Args:
        model: Current classifier model.
        feature_extractor: Feature extractor model.
        optimizer: Optimizer.
        criterion: Loss function (CrossEntropyLoss).
        dataloader: Dataloader for current training data.
        memory_buffer: Experience replay buffer.
        previous_model: Previous classifier model for knowledge distillation.
        alpha: Weight for distillation loss.
        randmix_alpha: Parameter for mixup.
        use_mixup: Whether to apply mixup.
    """
    model.train()
    feature_extractor.eval()  # Freeze feature extractor

    for epoch in range(NUM_EPOCHS):
        running_loss = 0.0
        for images, labels in dataloader:
            images = images.to(device)
            labels = labels.to(device)

            with torch.no_grad():
                features = feature_extractor(images)

            if use_mixup:
                # Apply mixup augmentation
                mixed_images, lam, index = randmix(images, alpha=randmix_alpha)
                mixed_features = feature_extractor(mixed_images.to(device))

                outputs = model(mixed_features)
                loss = lam * criterion(outputs, labels) + (1 - lam) * criterion(outputs, labels[index].to(device))
            else:
                outputs = model(features)
                loss = criterion(outputs, labels)

            # Knowledge Distillation Loss
            if previous_model is not None:
                with torch.no_grad():
                    previous_features = feature_extractor(images)
                    previous_outputs = previous_model(previous_features)
                distillation_loss = nn.KLDivLoss()(nn.functional.log_softmax(outputs / 2, dim=1),
                                                  nn.functional.softmax(previous_outputs / 2, dim=1))
                loss = (1 - alpha) * loss + alpha * distillation_loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        avg_loss = running_loss / len(dataloader)
        print(f"Epoch [{epoch+1}/{NUM_EPOCHS}] Loss: {avg_loss:.4f}")

    # After training, add samples to memory
    # Use a subset of the current batch to add to memory
    if use_mixup:
        memory_buffer.add_samples(mixed_images.cpu(), labels.cpu())
    else:
        memory_buffer.add_samples(images.cpu(), labels.cpu())

# ---------------------
# 10. Evaluation Function
# ---------------------
def evaluate_model(model, feature_extractor, dataloader):
    """
    Evaluates the model on the given data.
    Args:
        model: Classifier model.
        feature_extractor: Feature extractor model.
        dataloader: Dataloader for evaluation data.
    Returns:
        accuracy: Classification accuracy in percentage.
    """
    model.eval()
    feature_extractor.eval()
    all_preds = []
    all_labels = []
    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            labels = labels.to(device)
            features = feature_extractor(images)
            outputs = model(features)
            _, preds = torch.max(outputs, 1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    accuracy = accuracy_score(all_labels, all_preds) * 100
    return accuracy

# ---------------------
# 11. Prototype Computation
# ---------------------
def compute_prototypes(feature_extractor, classifier, memory_buffer, num_classes=10):
    """
    Computes class prototypes from the memory buffer.
    Args:
        feature_extractor: Feature extractor model.
        classifier: Current classifier model.
        memory_buffer: Experience replay buffer.
        num_classes: Number of classes.
    Returns:
        prototypes: Tensor of shape (num_classes, feature_dim)
    """
    feature_extractor.eval()
    classifier.eval()
    class_features = {i: [] for i in range(num_classes)}

    with torch.no_grad():
        for img, label in memory_buffer.buffer:
            img = img.to(device).unsqueeze(0)
            label = label.item()
            feat = feature_extractor(img)
            class_features[label].append(feat.cpu())

    prototypes = []
    for i in range(num_classes):
        if len(class_features[i]) > 0:
            # Compute mean feature for the class
            class_feat = torch.cat(class_features[i], dim=0)
            mean_feat = class_feat.mean(dim=0)
            prototypes.append(mean_feat)
        else:
            # If no samples for the class, use zeros
            prototypes.append(torch.zeros(512))

    prototypes = torch.stack(prototypes).to(device)
    return prototypes

# ---------------------
# 12. Continual Learning Pipeline
# ---------------------
def continual_learning_pipeline(
    train_tar_pth_paths,
    heldout_tar_pth_paths,
    num_classes=10,
    memory_size=500,
    pseudo_label_threshold=0.5,
    alpha=0.5,  # Weight for knowledge distillation
    randmix_alpha=0.2,
    use_mixup=True
):
    """
    Main pipeline to perform continual learning on sequential domains.
    Args:
        train_tar_pth_paths: List of paths to training .tar.pth files (D1 to D20).
        heldout_tar_pth_paths: List of paths to held-out .tar.pth files (D̂1 to D̂20).
        num_classes: Number of classes.
        memory_size: Capacity of experience replay buffer.
        pseudo_label_threshold: Threshold for selecting high-confidence pseudo-labels.
        alpha: Weight for knowledge distillation loss.
        randmix_alpha: Parameter for mixup.
        use_mixup: Whether to apply mixup.
    Returns:
        performance_matrix: Numpy array containing accuracies (models x datasets).
    """
    # Define transformations
    transform = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
    ])

    # Initialize feature extractor
    feature_extractor = FeatureExtractor().to(device)

    # Initialize memory buffer
    memory_buffer = MemoryBuffer(memory_size)

    # Initialize variables
    previous_model = None
    performance_matrix = np.zeros((len(train_tar_pth_paths), len(heldout_tar_pth_paths)))

    # Initialize classifier (to be updated in the loop)
    classifier = None

    # Iterate through each training dataset
    for i, train_path in enumerate(train_tar_pth_paths):
        print(f"\n=== Training on Dataset D{i+1} ===")

        # Load current training dataset
        try:
            current_dataset = TarPTHDataset(train_path, transform=transform)
            print(f"Loaded Dataset D{i+1} with {len(current_dataset)} samples.")
        except AssertionError as ae:
            print(f"Assertion Error while loading {train_path}: {ae}")
            continue
        except Exception as e:
            print(f"Unexpected error while loading {train_path}: {e}")
            continue

        # Create DataLoader for current dataset
        current_loader = DataLoader(current_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)

        # Initialize or update classifier
        if i == 0:
            # Initial classifier
            classifier = Classifier(input_dim=512, num_classes=num_classes).to(device)
        else:
            # Subsequent classifiers with knowledge distillation
            classifier = Classifier(input_dim=512, num_classes=num_classes,
                                    previous_model=copy.deepcopy(previous_model),
                                    alpha=alpha).to(device)

        # Define optimizer and loss
        optimizer = optim.Adam(classifier.parameters(), lr=LEARNING_RATE)
        criterion = nn.CrossEntropyLoss()

        if i == 0:
            # Dataset 1 is labeled: Direct supervised training
            train_model(
                model=classifier,
                feature_extractor=feature_extractor,
                optimizer=optimizer,
                criterion=criterion,
                dataloader=current_loader,
                memory_buffer=memory_buffer,
                previous_model=previous_model,
                alpha=alpha,
                randmix_alpha=randmix_alpha,
                use_mixup=use_mixup
            )
        else:
            # For datasets 2 to N: Unlabeled, use pseudo-labeling
            # Generate pseudo-labels
            pseudo_images, pseudo_labels, confidences = generate_pseudo_labels(
                model=previous_model,
                feature_extractor=feature_extractor,
                dataloader=current_loader,
                threshold=pseudo_label_threshold
            )

            print(f"Generated {len(pseudo_images)} pseudo-labeled samples.")

            # Create pseudo-labeled dataset
            if len(pseudo_images) > 0:
                pseudo_dataset = torch.utils.data.TensorDataset(pseudo_images, pseudo_labels)
                pseudo_loader = DataLoader(pseudo_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
                combined_dataset = ConcatDataset([current_dataset, pseudo_dataset])
                combined_loader = DataLoader(combined_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
                print(f"Combined dataset has {len(combined_dataset)} samples.")
            else:
                combined_loader = current_loader
                print("No pseudo-labeled samples generated. Using current dataset only.")

            # Include experience replay from memory
            replay_images, replay_labels = memory_buffer.get_samples(BATCH_SIZE)
            if replay_images is not None:
                replay_dataset = torch.utils.data.TensorDataset(replay_images, replay_labels)
                final_dataset = ConcatDataset([combined_dataset, replay_dataset])
                final_loader = DataLoader(final_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
                print(f"Final dataset with replay has {len(final_dataset)} samples.")
            else:
                final_loader = combined_loader
                print("Memory buffer is empty. Using combined dataset only.")

            if i < len(train_tar_pth_paths) - 1:
                # If more datasets remain, continue normally
                train_model(
                    model=classifier,
                    feature_extractor=feature_extractor,
                    optimizer=optimizer,
                    criterion=criterion,
                    dataloader=final_loader,
                    memory_buffer=memory_buffer,
                    previous_model=previous_model,
                    alpha=alpha,
                    randmix_alpha=randmix_alpha,
                    use_mixup=use_mixup
                )
            else:
                # If it's the last dataset and you want to apply Prototype Contrastive Alignment
                # Initialize contrastive loss
                pca_criterion = PrototypeContrastiveLoss()

                # Define a combined training loop with additional PCA loss
                classifier.train()
                feature_extractor.eval()

                for epoch in range(NUM_EPOCHS):
                    running_loss = 0.0
                    for images, labels in final_loader:
                        images = images.to(device)
                        labels = labels.to(device)

                        with torch.no_grad():
                            features = feature_extractor(images)

                        if use_mixup:
                            # Apply mixup augmentation
                            mixed_images, lam, index = randmix(images, alpha=randmix_alpha)
                            mixed_features = feature_extractor(mixed_images.to(device))

                            outputs = classifier(mixed_features)
                            loss = lam * criterion(outputs, labels) + (1 - lam) * criterion(outputs, labels[index].to(device))
                        else:
                            outputs = classifier(features)
                            loss = criterion(outputs, labels)

                        # Knowledge Distillation Loss
                        if previous_model is not None:
                            with torch.no_grad():
                                previous_features = feature_extractor(images)
                                previous_outputs = previous_model(previous_features)
                            distillation_loss = nn.KLDivLoss()(nn.functional.log_softmax(outputs / 2, dim=1),
                                                              nn.functional.softmax(previous_outputs / 2, dim=1))
                            loss = (1 - alpha) * loss + alpha * distillation_loss

                        # Prototype Contrastive Loss
                        prototypes = compute_prototypes(feature_extractor, classifier, memory_buffer, num_classes)
                        contrastive_loss = pca_criterion(features, prototypes, labels)

                        # Total loss
                        total_loss = loss + contrastive_loss

                        optimizer.zero_grad()
                        total_loss.backward()
                        optimizer.step()

                        running_loss += total_loss.item()

                    avg_loss = running_loss / len(final_loader)
                    print(f"PCA Epoch [{epoch+1}/{NUM_EPOCHS}] Loss: {avg_loss:.4f}")

            # Add current samples to memory
            if i == 0:
                # For the first dataset, add the actual labeled samples
                memory_buffer.add_samples(current_dataset.data.cpu(), current_dataset.labels.cpu())
                print(f"Added {len(current_dataset)} samples to memory buffer.")
            else:
                if len(pseudo_images) > 0:
                    memory_buffer.add_samples(pseudo_images.cpu(), pseudo_labels.cpu())
                    print(f"Added {len(pseudo_images)} pseudo-labeled samples to memory buffer.")
                else:
                    print("No pseudo-labeled samples to add to memory buffer.")

        # Update previous model
        previous_model = copy.deepcopy(classifier)

        # Evaluate the current model on all relevant held-out datasets
        for j in range(i + 1):
            heldout_dataset = TarPTHDataset(heldout_tar_pth_paths[j], transform=transform)
            heldout_loader = DataLoader(heldout_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)
            accuracy = evaluate_model(classifier, feature_extractor, heldout_loader)
            performance_matrix[i, j] = accuracy
            print(f"Model f{i+1} Accuracy on D̂{j+1}: {accuracy:.2f}%")

    return performance_matrix

# ---------------------
# 13. Preparing Datasets
# ---------------------
def prepare_datasets(train_folder, heldout_folder, num_datasets=20):
    """
    Prepares lists of paths to training and held-out datasets.
    Args:
        train_folder: Directory containing training .tar.pth files (D1 to D20).
        heldout_folder: Directory containing held-out .tar.pth files (D̂1 to D̂20).
        num_datasets: Number of datasets (default 20).
    Returns:
        train_tar_pth_paths: List of training dataset paths.
        heldout_tar_pth_paths: List of held-out dataset paths.
    """
    train_tar_pth_paths = []
    heldout_tar_pth_paths = []

    for i in range(1, num_datasets + 1):
        train_path = os.path.join(train_folder, f"trian_data/{i}_train_data.tar.pth")
        heldout_path = os.path.join(heldout_folder, f"eval_data{i}_eval_data.tar.pth")  # Using 'Dhat' to denote held-out
        if not os.path.exists(train_path):
            raise FileNotFoundError(f"Training dataset {train_path} not found.")
        if not os.path.exists(heldout_path):
            raise FileNotFoundError(f"Held-out dataset {heldout_path} not found.")
        train_tar_pth_paths.append(train_path)
        heldout_tar_pth_paths.append(heldout_path)

    return train_tar_pth_paths, heldout_tar_pth_paths

# ---------------------
# 14. Hyperparameters
# ---------------------
BATCH_SIZE = 64
LEARNING_RATE = 1e-3
NUM_EPOCHS = 10
MEMORY_SIZE = 500
PSEUDO_LABEL_THRESHOLD = 0.5
ALPHA = 0.5  # Weight for knowledge distillation
RANDMIX_ALPHA = 0.2
USE_MIXUP = True

# ---------------------
# 15. Main Execution
# ---------------------
if __name__ == "__main__":

    # Paths to the directories containing training and held-out .tar.pth files
    # Replace 'path_to_train_folder' and 'path_to_heldout_folder' with actual paths
     #   # Paths to zip files
    train_zip_path = "/content/train_data.zip"
    heldout_zip_path = "/content/eval_data.zip"

    # # # Directory where you want to extract the files
    # extracted_train_folder = "path_t_exto_extracted_train_folder"
    # extracted_heldout_folder = "path_toracted_heldout_folder"

    extracted_train_folder = "./train_data"

    extracted_heldout_folder = "./heldout_data"


    # # Extract train folder
    if not os.path.exists(extracted_train_folder):
        with zipfile.ZipFile(train_zip_path, 'r') as zip_ref:
            zip_ref.extractall(extracted_train_folder)

    # # Extract heldout folder
    if not os.path.exists(extracted_heldout_folder):
        with zipfile.ZipFile(heldout_zip_path, 'r') as zip_ref:
            zip_ref.extractall(extracted_heldout_folder)
    train_folder = "/content/train_data"      # Example: "./datasets/train/"
    heldout_folder = "/content/heldout_data"  # Example: "./datasets/heldout/"

    # Validate and prepare dataset paths
    try:
        train_tar_pth_paths, heldout_tar_pth_paths = prepare_datasets(train_folder, heldout_folder)
        print("All datasets loaded and validated successfully.")
    except FileNotFoundError as e:
        print(e)
        exit(1)

    # Optionally, perform additional validation
    # validate_datasets(train_tar_pth_paths, heldout_tar_pth_paths)

    # Run continual learning pipeline
    performance = continual_learning_pipeline(
        train_tar_pth_paths=train_tar_pth_paths,
        heldout_tar_pth_paths=heldout_tar_pth_paths,
        num_classes=10,
        memory_size=MEMORY_SIZE,
        pseudo_label_threshold=PSEUDO_LABEL_THRESHOLD,
        alpha=ALPHA,
        randmix_alpha=RANDMIX_ALPHA,
        use_mixup=USE_MIXUP
    )

    # Display performance matrix
    df = pd.DataFrame(
        performance,
        index=[f"f{i+1}" for i in range(len(train_tar_pth_paths))],
        columns=[f"Dhat{j+1}" for j in range(len(heldout_tar_pth_paths))]
    )
    print("\n=== Performance Matrix (Accuracy %) ===")
    print(df)

    # Save performance matrix to CSV
    df.to_csv("performance_matrix.csv")
    print("\nPerformance matrix saved to 'performance_matrix.csv'")  tell what this code does