"""
Offline Sample Dataset for Reviewers
Loads pre-extracted sample data without requiring database access
"""

import os
import pickle
import torch
from torch.utils.data import Dataset
import numpy as np

class OfflineSampleDataset(Dataset):
    """
    Dataset that loads pre-extracted samples from pickle file
    No database access required
    """
    def __init__(self, sample_data_path, transform=None):
        """
        Args:
            sample_data_path: Path to sample_data.pkl file
            transform: Optional transform to apply to samples
        """
        self.sample_data_path = sample_data_path
        self.transform = transform
        
        # Load sample data
        print(f"Loading sample data from {sample_data_path}...")
        with open(sample_data_path, 'rb') as f:
            data = pickle.load(f)
        
        self.samples = data['samples']
        self.indices = data.get('indices', list(range(len(self.samples))))
        self.config = data.get('config', {})
        self.num_samples = len(self.samples)
        
        print(f"Loaded {self.num_samples} samples from offline data")
        
    def __len__(self):
        return self.num_samples
    
    def __getitem__(self, idx):
        """
        Get a sample by index
        Returns the same format as the original dataset
        """
        if idx >= self.num_samples:
            raise IndexError(f"Index {idx} out of range for dataset of size {self.num_samples}")
        
        sample = self.samples[idx]
        
        # Apply transform if provided
        if self.transform is not None:
            sample = self.transform(sample)
        
        return sample
    
    def get_config(self):
        """Get the configuration used for this dataset"""
        return self.config

