from octopi.datasets import dataset, augment, cached_datset
from monai.data import DataLoader, SmartCacheDataset, CacheDataset, Dataset
from typing import List, Optional
from octopi.utils import io as io2
from octopi.datasets import io
import torch, os, random, gc
import multiprocess as mp

class TrainLoaderManager:

    def __init__(self, 
                 config: str, 
                 target_name: str,
                 target_session_id: str = None,
                 target_user_id: str = None,
                 voxel_size: float = 10, 
                 tomo_algorithm: List[str] = ['wbp'], 
                 tomo_batch_size: int = 15
                 ): 

        # Read Copick Projectdd
        self.config = config
        self.root = io.load_copick_config(config)

        # Copick Query for Target
        self.target_name = target_name 
        self.target_session_id = target_session_id
        self.target_user_id = target_user_id

        # Copick Query For Input Tomogram
        self.voxel_size = voxel_size
        self.tomo_algorithm = tomo_algorithm

        self.reload_training_dataset = True
        self.reload_validation_dataset = True
        self.val_loader = None
        self.train_loader = None
        self.tomo_batch_size = tomo_batch_size

        # Initialize the input dimensions   
        self.nx = None
        self.ny = None
        self.nz = None

    def get_available_runIDs(self):
        """
        Identify and return a list of run IDs that have segmentations available for the target.
        
        - Iterates through all runs in the project to check for segmentations that match 
        the specified target name, session ID, and user ID.
        - Only includes runs that have at least one matching segmentation.

        Returns:
            available_runIDs (list): List of run IDs with available segmentations.
        """        
        available_runIDs = []
        runIDs = [run.name for run in self.root.runs]
        for run in runIDs:
            run = self.root.get_run(run)
            seg = run.get_segmentations(name=self.target_name, 
                                        session_id=self.target_session_id, 
                                        user_id=self.target_user_id,
                                        voxel_size=float(self.voxel_size))
            if len(seg) > 0:
                available_runIDs.append(run.name)

        # If No Segmentations are Found, Inform the User
        if len(available_runIDs) == 0:
            print(
                f"[Error] No segmentations found for the target query:\n"
                f"TargetName: {self.target_name}, UserID: {self.target_user_id}, "
                f"SessionID: {self.target_session_id}\n"
                f"Please check the target name, user ID, and session ID.\n"
            )
            exit()

        return available_runIDs

    def get_data_splits(self,
                        trainRunIDs: str = None,
                        validateRunIDs: str = None,
                        train_ratio: float = 0.8,
                        val_ratio: float = 0.2,
                        test_ratio: float = 0.0,
                        create_test_dataset: bool = False):
        """
        Split the available data into training, validation, and testing sets based on input parameters.

        Args:
            trainRunIDs (str): Predefined list of run IDs for training. If provided, it overrides splitting logic.
            validateRunIDs (str): Predefined list of run IDs for validation. If provided with trainRunIDs, no splitting occurs.
            train_ratio (float): Proportion of available data to allocate to the training set.
            val_ratio (float): Proportion of available data to allocate to the validation set.
            test_ratio (float): Proportion of available data to allocate to the test set.
            create_test_dataset (bool): Whether to create a test dataset or leave it empty.

        Returns:
            myRunIDs (dict): Dictionary containing run IDs for training, validation, and testing.
        """          

        # Option 1: Only TrainRunIDs are Provided, Split into Train, Validate and Test (Optional)
        if trainRunIDs is not None and validateRunIDs is None:
            trainRunIDs, validateRunIDs, testRunIDs = io.split_multiclass_dataset(
                trainRunIDs, train_ratio, val_ratio, test_ratio, 
                return_test_dataset = create_test_dataset
            )
        # Option 2: TrainRunIDs and ValidateRunIDs are Provided, No Need to Split
        elif trainRunIDs is not None and validateRunIDs is not None:
            testRunIDs = None
        # Option 3: Use the Entire Copick Project, Split into Train, Validate and Test
        else:
            runIDs = self.get_available_runIDs()          
            trainRunIDs, validateRunIDs, testRunIDs = io.split_multiclass_dataset(
                runIDs, train_ratio, val_ratio, test_ratio, 
                return_test_dataset = create_test_dataset
            )

        # Get Class Info from the Training Dataset
        self._get_class_info(trainRunIDs)

        # Swap if Test Runs is Larger than Validation Runs
        if create_test_dataset and len(testRunIDs) > len(validateRunIDs):
            testRunIDs, validateRunIDs = validateRunIDs, testRunIDs

        # Determine if datasets fit entirely in memory based on the batch size
        # If the validation set is smaller than the batch size, avoid reloading
        if len(validateRunIDs) < self.tomo_batch_size:
            self.reload_validation_dataset  = False

        # If the training set is smaller than the batch size, avoid reloading
        if len(trainRunIDs) < self.tomo_batch_size:
            self.reload_training_dataset = False

        # Store the split run IDs into a dictionary for easy access
        self.myRunIDs = {
            'train': trainRunIDs,
            'validate': validateRunIDs,
            'test': testRunIDs
        }

        print(f"Number of training samples: {len(trainRunIDs)}")
        print(f"Number of validation samples: {len(validateRunIDs)}")
        if testRunIDs is not None:
            print(f'Number of test samples: {len(testRunIDs)}')    

        # Define separate batch sizes
        self.train_batch_size = min( len(self.myRunIDs['train']), self.tomo_batch_size)
        self.val_batch_size   = min( len(self.myRunIDs['validate']), self.tomo_batch_size)

        # Initialize data iterators for training and validation
        self._initialize_val_iterators()
        self._initialize_train_iterators()

        return self.myRunIDs

    def _get_class_info(self, trainRunDs):

        # Fetch a segmentation to determine class names and number of classes
        for runID in trainRunDs:
            run = self.root.get_run(runID)
            seg = run.get_segmentations(name=self.target_name, 
                                        session_id=self.target_session_id, 
                                        user_id=self.target_user_id,
                                        voxel_size=float(self.voxel_size))
            if len(seg) == 0:
                continue

            # If Session ID or User ID are None, Set Them Based on the First Found Segmentation
            if self.target_session_id is None:
                self.target_session_id = seg[0].session_id
            if self.target_user_id is None:
                self.target_user_id = seg[0].user_id

            # Read Yaml Config to Get Number of Classes and Class Names
            target_config = io2.check_target_config_path(self)
            class_names = target_config['input']['labels']
            self.Nclasses = len(class_names) + 1
            self.class_names = [name for name, idx in sorted(class_names.items(), key=lambda x: x[1])]

            # We Only need to read One Segmentation to Get Class Info
            break
    
    def _get_padded_list(self, data_list, batch_size):
        # Calculate padding needed to make `data_list` a multiple of `batch_size`
        remainder = len(data_list) % batch_size
        if remainder > 0:
            # Number of additional items needed to make the length a multiple of batch size
            padding_needed = batch_size - remainder
            # Extend `data_list` with a random subset to achieve the padding
            data_list = data_list + random.sample(data_list, padding_needed)
        # Shuffle the full list
        random.shuffle(data_list)
        return data_list        
    
    def _initialize_train_iterators(self):
        # Initialize padded train and validation data lists
        self.padded_train_list = self._get_padded_list(self.myRunIDs['train'], self.train_batch_size)

        # Create iterators
        self.train_data_iter = iter(self._get_data_batches(self.padded_train_list, self.train_batch_size))

    def _initialize_val_iterators(self):     
        # Initialize padded train and validation data lists
        self.padded_val_list = self._get_padded_list(self.myRunIDs['validate'], self.val_batch_size)

        # Create iterators
        self.val_data_iter = iter(self._get_data_batches(self.padded_val_list, self.val_batch_size))        

    def _get_data_batches(self, data_list, batch_size):
        # Generator that yields batches of specified size
        for i in range(0, len(data_list), batch_size):
            yield data_list[i:i + batch_size]

    def _extract_run_ids(self, data_iter_name, initialize_method):
        # Access the instance's data iterator by name
        data_iter = getattr(self, data_iter_name)
        try:
            # Attempt to get the next batch from the iterator
            runIDs = next(data_iter)
        except StopIteration:
            # Reinitialize the iterator if exhausted
            initialize_method()
            # Update the iterator reference after reinitialization
            data_iter = getattr(self, data_iter_name)
            runIDs = next(data_iter)
        # Update the instance attribute with the new iterator state
        setattr(self, data_iter_name, data_iter)
        return runIDs
    
    def create_train_dataloaders(
        self,
        crop_size: int = 96,
        num_samples: int = 64):

        train_batch_size = 1
        val_batch_size = 1

        # If reloads are disabled and loaders already exist, reuse them
        if self.reload_frequency < 0 and (self.train_loader is not None) and (self.val_loader is not None):
            return self.train_loader, self.val_loader         

        # We Only Need to Reload the Training Dataset if the Total Number of Runs is larger than 
        # the tomo batch size
        if self.train_loader is None: 

            # Fetch the next batch of run IDs
            trainRunIDs = self._extract_run_ids('train_data_iter', self._initialize_train_iterators)
            train_files = io.load_training_data(self.root, trainRunIDs, self.voxel_size, self.tomo_algorithm, 
                                                self.target_name, self.target_session_id, self.target_user_id, 
                                                progress_update=False)
            self._check_max_label_value(train_files)

            # Create the cached dataset with non-random transforms
            train_ds = CacheDataset(data=train_files, transform=augment.get_transforms(), cache_rate=1.0)

            # Delete the training files to free memory
            train_files = None
            gc.collect()

            # I need to read (nx,ny,nz) and scale the crop size to make sure it isnt larger than nx.
            if self.nx is None: (self.nx,self.ny,self.nz) = train_ds[0]['image'].shape[1:]
            self.input_dim = io.get_input_dimensions(train_ds, crop_size)

            # Wrap the cached dataset to apply random transforms during iteration
            self.dynamic_train_dataset = dataset.DynamicDataset(
                data=train_ds, 
                transform=augment.get_random_transforms(self.input_dim, num_samples, self.Nclasses)
            )

            # Define the number of processes for the DataLoader
            n_procs = min(mp.cpu_count(), 4)

            # DataLoader remains the same
            self.train_loader = DataLoader(
                self.dynamic_train_dataset,
                batch_size=train_batch_size,
                shuffle=False,
                num_workers=n_procs,
                pin_memory=torch.cuda.is_available(),
            )

        else:
            # Fetch the next batch of run IDs
            trainRunIDs = self._extract_run_ids('train_data_iter', self._initialize_train_iterators)
            train_files = io.load_training_data(self.root, trainRunIDs, self.voxel_size, self.tomo_algorithm, 
                                                self.target_name, self.target_session_id, self.target_user_id, 
                                                progress_update=False)
            self._check_max_label_value(train_files)

            train_ds = CacheDataset(data=train_files, transform=augment.get_transforms(), cache_rate=1.0)
            self.dynamic_train_dataset.update_data(train_ds)

        # We Only Need to Reload the Validation Dataset if the Total Number of Runs is larger than 
        # the tomo batch size
        if self.val_loader is None: 

            validateRunIDs = self._extract_run_ids('val_data_iter', self._initialize_val_iterators)             
            val_files   = io.load_training_data(self.root, validateRunIDs, self.voxel_size, self.tomo_algorithm, 
                                                self.target_name, self.target_session_id, self.target_user_id,
                                                progress_update=False) 
            self._check_max_label_value(val_files)

            # Create validation dataset
            val_ds = CacheDataset(data=val_files, transform=augment.get_transforms(), cache_rate=1.0)

            # Delete the validation files to free memory
            val_files = None
            gc.collect()

            # # I need to read (nx,ny,nz) and scale the crop size to make sure it isnt larger than nx.
            # if self.nx is None:
            #     (self.nx,self.ny,self.nz) = val_ds[0]['image'].shape[1:]

            # if crop_size > self.nx: self.input_dim = (self.nx, crop_size, crop_size)
            # else:                   self.input_dim = (crop_size, crop_size, crop_size)

            # Wrap the cached dataset to apply random transforms during iteration
            self.dynamic_validation_dataset = dataset.DynamicDataset( data=val_ds )

            dataset_size = len(self.dynamic_validation_dataset)
            n_procs = min(mp.cpu_count(), 8)

            # Create validation DataLoader
            self.val_loader  = DataLoader(
                self.dynamic_validation_dataset,
                batch_size=val_batch_size,
                num_workers=n_procs,
                pin_memory=torch.cuda.is_available(),
                shuffle=False,  # Ensure the data order remains consistent,            
            )
        else:
            validateRunIDs = self._extract_run_ids('val_data_iter', self._initialize_val_iterators)             
            val_files   = io.load_training_data(self.root, validateRunIDs, self.voxel_size, self.tomo_algorithm, 
                                                self.target_name, self.target_session_id, self.target_user_id,
                                                progress_update=False)  
            self._check_max_label_value(val_files)

        return self.train_loader, self.val_loader
    
    def get_reload_frequency(self, num_epochs: int):
        """
        Automatically calculate the reload frequency for the dataset during training.

        Returns:
            int: Reload frequency (number of epochs between dataset reloads).
        """
        if not self.reload_training_dataset:
            # No need to reload if all tomograms fit in memory
            print("All training samples fit in memory. No reloading required.")
            self.reload_frequency = -1

        else:
            # Calculate the number of segments based on total training runs and batch size
            num_segments = (len(self.myRunIDs['train']) + self.tomo_batch_size - 1) // self.tomo_batch_size

            # Calculate reload frequency to distribute reloading evenly over epochs
            self.reload_frequency = max(num_epochs // num_segments, 1)

            print(f"\nReloading {self.tomo_batch_size} tomograms every {self.reload_frequency} epochs\n")

            # Warn if the number of epochs is insufficient for full dataset coverage
            if num_epochs < num_segments:
                print(
                    f"Warning: Chosen number of epochs ({num_epochs}) may not be sufficient "
                    f"to train over all training samples. Consider increasing the number of epochs "
                    f"to at least {num_segments}\n."
                )

    def _check_max_label_value(self, train_files):
        max_label_value = max(file['label'].max() for file in train_files)
        if max_label_value > self.Nclasses:
            print(f"Warning: Maximum class label value {max_label_value} exceeds the number of classes {self.Nclasses}.")
            print("This may cause issues with the model's output layer.")
            print("Consider adjusting the number of classes or the label values in your data.\n")

    def get_dataloader_parameters(self):

        parameters = {
            'config': self.config,
            'target_name': self.target_name,
            'target_session_id': self.target_session_id,
            'target_user_id': self.target_user_id,
            'voxel_size': self.voxel_size,
            'tomo_algorithm': self.tomo_algorithm,
            'tomo_batch_size': self.tomo_batch_size,
            'reload_frequency': self.reload_frequency,
            'testRunIDs': self.myRunIDs['test'],
            'valRunIDs': self.myRunIDs['validate'],    
            'trainRunIDs': self.myRunIDs['train'],
        }

        return parameters                

class PredictLoaderManager:

    def __init__(self, 
                 config: str, 
                 voxel_size: float = 10, 
                 tomo_algorithm: str = 'wbp', 
                 tomo_batch_size: int = 15, # Number of Tomograms to Load Per Sub-Epoch    
                 Nclasses: int = 3):
        
        # Read Copick Project
        self.copick_config = config
        self.root = io.load_copick_config(config)

        # Copick Query For Input Tomogram
        self.voxel_size = voxel_size
        self.tomo_algorithm = tomo_algorithm

        self.Nclasses = Nclasses
        self.tomo_batch_size = tomo_batch_size 

        # Initialize the input dimensions   
        self.nx = None
        self.ny = None
        self.nz = None


    def create_predict_dataloader(
        self, 
        voxel_spacing: float, 
        tomo_algorithm: str,       
        runIDs: str = None):

        # Split trainRunIDs, validateRunIDs, testRunIDs
        if runIDs is None:
            runIDs = [run.name for run in self.root.runs]
            
        # Load the test data
        test_files = io.load_predict_data(self.root, runIDs, voxel_spacing, tomo_algorithm)  

        # Create the cached dataset with non-random transforms
        test_ds = CacheDataset(data=test_files, transform=augment.get_predict_transforms())

        # Read (nx,ny,nz) for input tomograms.
        if self.nx is None:
            (self.nx,self.ny,self.nz) = test_ds[0]['image'].shape[1:]

        # Create the DataLoader
        test_loader = DataLoader(test_ds, 
                                batch_size=4, 
                                shuffle=False, 
                                num_workers=4, 
                                pin_memory=torch.cuda.is_available())
        return test_loader
    
    def get_dataloader_parameters(self):

        parameters = {
            'config': self.copick_config,
            'voxel_size': self.voxel_size,
            'tomo_algorithm': self.tomo_algorithm
        }

        return parameters    
    
    