model/splitters/base_splitters.py

import random
import pandas as pd
from abc import ABC, abstractmethod
from Bio import SeqIO, PDB
from sklearn.preprocessing import MinMaxScaler
import numpy as np
import pickle
import copy
import shutil

import os
import sys
root_folder = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))

from model.utils.other_utils import aa_dict_3to1
from model.utils.data_utils import find_mutation_positions_multithreaded, MutationFormat

class BaseSplitter(ABC):
    """Abstract base class for splitters."""

    """
    Attributes:
        wt_seq (str): Wild-type sequence of the protein.
        use_cache (bool): Flag to determine if caching should be used.
        random_state (int): Random state for reproducibility.
        data (pd.DataFrame): DataFrame containing the protein data.
        file_attrs (dict): Dictionary containing file attributes and paths.

    Example Usage:

    splitter = BaseSplitter(training_dataset_fname, 
                            wt_file, 
                            csv_has_header=True,  # Whether input CSV has a header row (default: True)
                            use_cache=True        # Whether to cache split results (default: True)
                            )
    splitter.split_data()
    """

    def __init__(self, name, data, wt_file, csv_has_header=False, use_cache=False,
                 random_state=42, type='biomolecules',
                 **kwargs):
        """
        Args:
        - name (str): Name of the biomolecule.
        - data (str or pd.DataFrame): Accepts a table of data, containing two columns: one for sequences (for multi-chain proteins, colon-separated sequences) and the second one for the labels.
        - wt_file (str or list): File path(s) to wild-type sequence of protein of interest in FASTA format.
        - csv_has_header (bool): Flag to modify if your data has headers (True) or does not (False).
        - use_cache (bool): Flag to use cache.
        - random_state (int): Random state for reproducibility.
        - type (str): Type of biomolecule.
        - **kwargs: Additional keyword arguments.
        """

        self.wt_seq_lens = []
        self.wt_seqs = []
        if isinstance(wt_file, str):
            self.wt_seq_lens.append(len(str(SeqIO.read(wt_file, "fasta").seq)))
            self.wt_seqs.append(str(SeqIO.read(wt_file, "fasta").seq))
        elif isinstance(wt_file, list):
            for file in wt_file:
                self.wt_seq_lens.append(len(str(SeqIO.read(file, "fasta").seq)))
                self.wt_seqs.append(str(SeqIO.read(file, "fasta").seq))
        self.wt_seq = ''.join(self.wt_seqs)
        self.use_cache = use_cache
        self.random_state = random_state

         # If the data is a CSV file
        if isinstance(data, str) and data.endswith('.csv'):
            self.data = pd.read_csv(data, header=0 if csv_has_header else None)
            # Rename columns for consistency
            self.data.rename(columns={self.data.columns[0]: 0, self.data.columns[1]: 1}, inplace=True)
            dataset_name = os.path.splitext(os.path.basename(data))[0]
            dataset_file = os.path.join(root_folder, type, name, dataset_name + '.csv')
        elif isinstance(data, pd.DataFrame):  # If the data is already a DataFrame
            self.data = data.copy()
            # Ensure column names are standardized
            self.data.rename(columns={self.data.columns[0]: 0, self.data.columns[1]: 1}, inplace=True)
            dataset_name = 'dataframe_input'
            dataset_file = os.path.join(root_folder, type, name, dataset_name + '.csv')
        else:
            raise ValueError("Invalid data format: data must be a file path to a CSV or a DataFrame.")

        # Define file attributes and split directory
        self.file_attrs = {
            'dataset_file': dataset_file,
            'dataset_name': dataset_name,
            'dataset_dir': os.path.join(root_folder, type, name),
            'split_dir': os.path.join(root_folder, type, name, 'split_cache', dataset_name)
        }

        # Create cache directory if needed
        if self.use_cache:
            os.makedirs(self.file_attrs['split_dir'], exist_ok=True)

        # copy dataset file to new location
        if not os.path.exists(self.file_attrs['dataset_file']):
            os.makedirs(os.path.dirname(self.file_attrs['dataset_file']), exist_ok=True)
            
            if isinstance(data, str):  # data is a file path
                shutil.copy(data, self.file_attrs['dataset_file'])
            else:  # data is a DataFrame
                data.to_csv(self.file_attrs['dataset_file'], index=False)
        
# Note: For new classes on top of ProteinSplitter, all unique args for split_data() should be used for generate split_type attribute for proper cache storage

class ProteinSplitter(BaseSplitter):
    """Class for splitting protein datasets."""    

    """
    Attributes:
        y_scaling (str): String indicating whether y values are scaled or not.
        val_split (float): Fraction of data for validation set.
        file_attrs (dict): Dictionary containing file attributes and paths.
            base_splitter_path (str): Path to the base splitter pickle file.
        splits (dict): Dictionary to store the data splits.

    Example Usage:

    splitter = ProteinSplitter(training_dataset_fname, 
                                wt_file, 
                                csv_has_header=True,  # Whether input CSV has a header row
                                use_cache=True,       # Whether to cache results to disk
                                y_scaling=True,       # Whether to scale y values to [0,1]
                                val_split=None        # Fraction of data for validation (None=no validation split)
                                )
    splitter.split_data()
    """

    def __init__(self, protein_name, data, wt_file, csv_has_header=False, use_cache=False, 
                 random_state=42,
                 y_scaling=False, 
                 val_split=None,
                 **kwargs):
        """
        Args:
        - protein_name (str): Name of the protein.
        - data (str or pd.DataFrame): Accepts a table of data, containing two columns: one for sequences and the second one for the labels.
        - wt_file (str or list): File path(s) to wild-type sequence of protein of interest in FASTA format.
        - csv_has_header (bool): Flag to modify if your data has headers (True) or does not (False).
        - use_cache (bool): Flag to use cache.
        - y_scaling (bool): Flag to scale y values.
        - val_split (float): Fraction of data to partition into the validation set.
        - random_state (int): Random state for reproducibility.
        - **kwargs: Additional keyword arguments.
        """

        super().__init__(protein_name, data, wt_file, csv_has_header=csv_has_header, use_cache=use_cache, 
                         random_state=random_state, type='proteins',
                         **kwargs)
        
        # Define base protein splitter path
        if y_scaling == False:
            self.y_scaling = "y_unscaled"
        else:
            self.y_scaling = "y_scaled"

        self.val_split = val_split
        self.kfold_splits = False

        self.file_attrs['base_splitter_path'] = os.path.join(self.file_attrs['split_dir'], "base_splitter" + '_' + self.y_scaling + ".pkl")

        # Load existing base protein splitter or create a new one
        if self.use_cache and os.path.exists(self.file_attrs['base_splitter_path']):
            self.data = pd.read_pickle(self.file_attrs['base_splitter_path'])
            
        # Create base protein splitter if not found and save is use_cache is True
        else:
            
            if len(self.wt_seq_lens) > 1:
                # check MutationFormat of column 0
                if MutationFormat(self.data[0].iloc[0].split(':')[0], self.wt_seq).format == 'Mutation String':
                    self.data[0] = self.data[0].apply(lambda x: self._shift_mutation_position(x, self.wt_seq_lens, 'Mutation String'))
                elif MutationFormat(self.data[0].iloc[0].split(':')[0], self.wt_seq).format == 'Mutation List':
                    self.data[0] = self.data[0].apply(lambda x: self._shift_mutation_position(x, self.wt_seq_lens, 'Mutation List'))
                elif MutationFormat(self.data[0].iloc[0].split(':')[0], self.wt_seq).format == 'Full Sequence':
                    self.data[0] = self.data[0].apply(lambda x: ''.join(x.split(':')))
                else:
                    raise ValueError('Mutation format not recognized')

            self.data[0] = self.data[0].apply(lambda x: MutationFormat(x, self.wt_seq).to_full_sequence())
            mut_positions = find_mutation_positions_multithreaded(self.wt_seq, self.data[0].tolist())
            self.data['mut_positions'] = mut_positions
            self.data['muts'] = self.data[0].apply(lambda x: MutationFormat(x, self.wt_seq).to_mutation_string())
            self.data['mut_load'] = self.data['mut_positions'].apply(lambda x: len(x))

            if y_scaling == True:
                scaler = MinMaxScaler()
                scaled_data = scaler.fit_transform(np.array(self.data[1]).reshape(-1, 1))
                self.data[1] = scaled_data

            # save as pickle file
            if self.use_cache:
                self.data.to_pickle(self.file_attrs['base_splitter_path'])

    def _shift_mutation_position(self, inputs, lengths, type):
        # inputs is a list of mutation strings or mutation lists
        # goal is to return a combined mutation string/mutation list
        inputs = inputs.split(':')
        lengths = [sum(lengths[:i+1]) for i in range(len(lengths))]
        new_inputs  = []
        # Process first input differently based on type
        if type == 'Mutation String':
            new_inputs.extend(inputs[0].split('/'))
        
        # Process remaining inputs
        for i in range(1, len(inputs)):
            mutations = inputs[i].split('/') if type == 'Mutation String' else inputs[i]
            for mut in mutations:
                if mut == 'WT':
                    break
                else:
                    mut_pos = mut[1:-1]
                    mut_pos = str(int(mut_pos) + lengths[i-1]) 
                    mut = mut[0] + mut_pos + mut[-1]
                    new_inputs.append(mut)
        
        if len(new_inputs) > 1 and "WT" in new_inputs:
            new_inputs.remove("WT")

        return '/'.join(new_inputs)

    def split_data(self, iter=None):
        """
        Splits data into training and test sets. 
        Data split into the test set is given a group label of 1,
        while data split into the training set is given a group label of 0.
        
        Args:
        - iter (int, optional): Iteration number for naming the split file.
        """
                   
        self.split_type = 'base'

        raise NotImplementedError("This method should be implemented by the subclass.")

    def _save_splits(self, iter=None):
        """
        Save splits for specific split type.
        
        Args:
        - iter (int, optional): Iteration number for naming the split file.
        """

        test_size = self.data['group'].sum() / len(self.data)

        # 0 for train, 1 for test, 2 for val

        if self.val_split is not None:

            if self.kfold_splits == True:
                pass
            else:
                # Separate train set into train and val sets
                rows_with_marker_0 = self.data[self.data['group'] == 0]
                num_to_sample = int(len(rows_with_marker_0) * self.val_split)
                sampled_indices = rows_with_marker_0.sample(n=num_to_sample).index
                self.data.loc[sampled_indices, 'group'] = 2

            X_train = self.data[self.data['group'] == 0][0].values
            X_val = self.data[self.data['group'] == 2][0].values
            X_test = self.data[self.data['group'] == 1][0].values
            y_train = self.data[self.data['group'] == 0][1].values
            y_val = self.data[self.data['group'] == 2][1].values
            y_test = self.data[self.data['group'] == 1][1].values

            train_size = len(X_train) / len(self.data)
            val_size = len(X_val) / len(self.data)

            # return splits
            # initialize dictionary to store splits
            self.splits = {'X_train': X_train, 'X_val': X_val, 'X_test': X_test, 'y_train': y_train, 'y_val': y_val, 'y_test': y_test}


            # check if iter is none:
            if iter is None:
                split_name = f'split_by_{self.split_type}_{int(train_size*100)}-{int(val_size*100)}-{int(test_size*100)}-{self.y_scaling}'
            else:
                split_name = f'split_by_{self.split_type}_{int(train_size*100)}-{int(val_size*100)}-{int(test_size*100)}-{self.y_scaling}_iter{iter}'
            
        else:
            X_train = self.data[self.data['group'] == 0][0].values
            X_test = self.data[self.data['group'] == 1][0].values
            y_train = self.data[self.data['group'] == 0][1].values
            y_test = self.data[self.data['group'] == 1][1].values

            # initialize dictionary to store splits
            self.splits = {'X_train': X_train, 'X_test': X_test, 'y_train': y_train, 'y_test': y_test}
            
            # check if iter is none:
            if iter is None:
                split_name = f'split_by_{self.split_type}_{int((1-test_size)*100)}-{int(test_size*100)}-{self.y_scaling}'
            else:
                split_name = f'split_by_{self.split_type}_{int((1-test_size)*100)}-{int(test_size*100)}-{self.y_scaling}_iter{iter}'

        # Save split or load split if it already exists to prevent overwriting
        self.splits['split_name'] = split_name

        file_name = os.path.join(self.file_attrs['split_dir'], split_name + ".pkl")

        if self.use_cache:
            if not os.path.exists(file_name):
                with open(file_name, 'wb') as file:
                    print("Split saved.")
                    pickle.dump(self.splits, file)
            
            else:
                with open(file_name, 'rb') as file:
                    print("Split already exists. Generated splits not saved. Loading pre-existing split.")
                    self.splits = pickle.load(file)

    def _assign_folds(self, k_folds):
        """
        Assigns fold labels to the data for K-Fold cross-validation if there is a validation split.
        
        Args:
        - k_folds (int): Number of folds.
        """

        if self.random_state is not None:
            np.random.seed(self.random_state)

        # initialize fold columns
        self.data['fold'] = None

        # get indices of rows designated for training/validation
        train_indices = self.data[self.data['group'] == 0].index

        # Shuffle the indices
        shuffled_indices = np.random.permutation(train_indices)

        # Determine the number of points per fold
        fold_sizes = [len(shuffled_indices) // k_folds] * k_folds
        for i in range(len(shuffled_indices) % k_folds):
            fold_sizes[i] += 1

        # Assign fold labels
        groups = np.zeros(len(self.data), dtype=int)
        current_idx = 0
        for fold_idx, fold_size in enumerate(fold_sizes):
            for i in range(current_idx, current_idx + fold_size):
                self.data.loc[shuffled_indices[i], 'fold'] = fold_idx
            current_idx += fold_size

    def _save_folds(self,k_folds):

        # assert at least 2 folds or code breaks
        assert k_folds >= 2, "At least 2 folds are required"

        # mark object as having kfold splits
        self.kfold_splits = True
        self._assign_folds(k_folds=k_folds)

        folds = []
        for fold_num in range(k_folds):

            # Create a deep copy of self and add it to folds
            fold_copy = copy.deepcopy(self)
            
            # Assign fold num to validation set if match fold num, otherwise keep in train or test
            fold_copy.data['group'] = np.where(
                fold_copy.data['fold'] == fold_num, 2, fold_copy.data['group']
            )

            fold_copy._save_splits(iter=fold_num)
            folds.append(fold_copy)
            
        self.folds = folds
    
    def load_splits(self, file_path):
        """
        Loads the split from a pickle file.
        
        Args:
        - file_path (str): Path to the pickle file containing the splits.
        """
        with open(file_path, 'rb') as file:
            self.splits = pickle.load(file)

class KFoldProteinSplitter(ProteinSplitter):
    """
    Class for K-Fold splitting of protein datasets.

    Attributes:
        data (pd.DataFrame): The protein dataset.
        random_state (int): Random seed for reproducibility.
        splits (dict): Dictionary to store the data splits.
        split_type (str): Type of split being performed.
        file_attrs (dict): Dictionary containing file attributes.
        use_cache (bool): Whether to use cached splits.
        y_scaling (str): Method for scaling y values.
        val_split (float): Fraction of data to use for validation.

    Example Usage:

    splitter = KFoldProteinSplitter(training_dataset_fname, 
                                    wt_file, 
                                    csv_has_header=True,  # Whether input CSV has a header row
                                    use_cache=True,       # Whether to cache results to disk
                                    y_scaling=True,       # Whether to scale y values to [0,1]
                                    val_split=None        # Fraction of data for validation (None=no validation split)
                                    )
    splitter.split_data(5) # Performs a 5-Fold split for 5-Fold cross-validation
    splits = splitter.generate_splits() # Returns a list of 5 splitter objects, one for each fold
    """

    def _assign_folds(self, n_splits):
        """
        Assigns fold labels to the data for K-Fold cross-validation.
        
        Args:
        - n_splits (int): Number of folds.
        """
    
        if self.random_state is not None:
            np.random.seed(self.random_state)
        
        # Shuffle the DataFrame indices
        shuffled_indices = np.random.permutation(self.data.index)
        
        # Determine the number of points per fold
        fold_sizes = [len(shuffled_indices) // n_splits] * n_splits
        for i in range(len(shuffled_indices) % n_splits):
            fold_sizes[i] += 1
        
        # Assign fold labels
        groups = np.zeros(len(self.data), dtype=int)
        current_idx = 0
        for fold_idx, fold_size in enumerate(fold_sizes):
            for i in range(current_idx, current_idx + fold_size):
                groups[shuffled_indices[i]] = fold_idx
            current_idx += fold_size
        
        # Add 'group' column to the DataFrame
        self.data['fold'] = groups

    def _split_data(self, fold_number):
        """
        Splits data into training and test sets based on the specified fold number.
        
        Args:
        - fold_number (int): Fold number to include in the test set.
        """

        self.data['group'] = np.where(
            self.data['fold'] == fold_number, 1, 0
        )

        # Save splits
        self.split_type = f'kfold-{fold_number}'
        self._save_splits()

    def generate_splits(self, n_splits):
      
      self._assign_folds(n_splits=n_splits)

      splits = []
      for fold_num in range(n_splits):
          self._split_data(fold_number=fold_num)
          # Create a deep copy of self and add it to splits
          split_copy = copy.deepcopy(self)
          splits.append(split_copy)
      
      return splits

class RoundProteinSplitter(ProteinSplitter):
    """
    Class for splitting protein datasets based on round number.

    Attributes:
        data (pd.DataFrame): The protein dataset.
        random_state (int): Random seed for reproducibility.
        splits (dict): Dictionary to store the data splits.
        split_type (str): Type of split being performed.
        file_attrs (dict): Dictionary containing file attributes.
        use_cache (bool): Whether to use cached splits.
        y_scaling (str): Method for scaling y values.
        val_split (float): Fraction of data to use for validation.

    Example Usage:

    splitter = RoundProteinSplitter(training_dataset_fname, 
                                    wt_file, 
                                    csv_has_header=True,  # Whether input CSV has a header row
                                    use_cache=True,       # Whether to cache results to disk
                                    y_scaling=True,       # Whether to scale y values to [0,1]
                                    val_split=None        # Fraction of data for validation (None=no validation split)
                                    )
    splitter.split_data(1, 3) # Splits data such that all data from round 1 and 2 are in the training set, and all data from round 3 and above are in the test set
    """

    def split_data(self, max_train_round, min_test_round, iter=None, k_folds=None):
        """
        Splits data into training and test sets based on round number.
        
        Args:
        - max_train_round (int): Maximum round number to include in the training set.
        - min_test_round (int): Minimum round number to include in the test set.
        - iter (int, optional): Iteration number for naming the split file.
        - k_folds (int, optional): Number of folds to generate.
        """
        assert 'round' in self.data.columns, "DataFrame must contain a 'round' column"
        assert max_train_round < min_test_round, "Maximum training round must be less than minimum test round"

        self.data['group'] = np.where(
            self.data['round'] <= max_train_round, 0,
            np.where(
                self.data['round'] >= min_test_round, 1, np.nan
            )
        )

        # Save splits
        self.split_type = f'round-{max_train_round}-{min_test_round}'

        if k_folds is not None:
            if self.val_split is not None:
                self._save_folds(k_folds=k_folds)
            elif self.val_split is None:
                print("No validation split, cannot generate kfolds")
                return

        else:
            
            if iter is None:
                self._save_splits()
            else:
                self._save_splits(iter=iter)

class RandomProteinSplitter(ProteinSplitter):
    """
    Class for splitting protein datasets randomly.

    Attributes:
        data (pd.DataFrame): The protein dataset.
        random_state (int): Random seed for reproducibility.
        splits (dict): Dictionary to store the data splits.
        split_type (str): Type of split being performed.
        file_attrs (dict): Dictionary containing file attributes.
        use_cache (bool): Whether to use cached splits.
        y_scaling (str): Method for scaling y values.
        val_split (float): Fraction of data to use for validation.

    Example Usage:

    splitter = RandomProteinSplitter(training_dataset_fname, 
                                    wt_file, 
                                    csv_has_header=True,  # Whether input CSV has a header row
                                    use_cache=True,       # Whether to cache results to disk
                                    y_scaling=True,       # Whether to scale y values to [0,1]
                                    val_split=None        # Fraction of data for validation (None=no validation split)
                                    )
    splitter.split_data(test_size=0.2) # Splits data into 80% training and 20% test sets
    """

    def split_data(self, test_size=0.2, iter=None, k_folds=None):
        """
        Splits data into training and test sets randomly.
        
        Args:
        - test_size (float): Fraction of data to partition into the test set.
        - iter (int, optional): Iteration number for naming the split file.
        - k_folds (int, optional): Number of folds to generate.
        """      
        random.seed(self.random_state)  # Add this line
        test_len = int(test_size * len(self.data))
        test_indices = random.sample(range(len(self.data)), test_len)
        self.data['group'] = 0
        self.data.loc[test_indices, 'group'] = 1
        
        # Save splits
        self.split_type = 'random'

        if k_folds is not None:
            if self.val_split is not None:
                self._save_folds(k_folds=k_folds)
            elif self.val_split is None:
                print("No validation split, cannot generate kfolds")
                return

        else:
            
            if iter is None:
                self._save_splits()
            else:
                self._save_splits(iter=iter)

class PositionProteinSplitter(ProteinSplitter):
    """
    Class for splitting protein datasets based on mutation positions.

    Attributes:
        data (pd.DataFrame): The protein dataset.
        random_state (int): Random seed for reproducibility.
        splits (dict): Dictionary to store the data splits.
        split_type (str): Type of split being performed.
        file_attrs (dict): Dictionary containing file attributes.
        use_cache (bool): Whether to use cached splits.
        y_scaling (str): Method for scaling y values.
        val_split (float): Fraction of data to use for validation.
    
    Example Usage:

    splitter = PositionProteinSplitter(training_dataset_fname, 
                                    wt_file, 
                                    csv_has_header=True,  # Whether input CSV has a header row
                                    use_cache=True,       # Whether to cache results to disk
                                    y_scaling=True,       # Whether to scale y values to [0,1]
                                    val_split=None        # Fraction of data for validation (None=no validation split)
                                    )
    # Splits data into training and test sets based on mutation positions:
    # 1. Randomly samples test_size_sample fraction of variants to get mutation positions to exclude
    # 2. Any variant containing those positions goes into test set
    # 3. Repeats sampling up to iter times until test set size is between test_size_min and test_size_max
    # 4. If test set size requirements not met after iter attempts, uses best attempt

    splitter.split_data(test_size_sample=0.2, iter=3, test_size_min=0.2, test_size_max=0.3)
    """

    def split_data(self, test_size_sample, sample_iter=3, test_size_min=0.2, test_size_max=0.3, iter=None, k_folds=None):
        """
        Splits the dataset into training and test sets based on mutation positions occupied in each variant.
        
        Args:
        - test_size_sample (float): Fraction to sample to retrieve mutation positions to exclude out of the training set.
        - sample_iter (int): Number of sampling iterations.
        - test_size_min (float): Minimum test size set desired.
        - test_size_max (float): Maximum test size set allowed.
        - iter (int, optional): Iteration number for naming the split file.
        - k_folds (int, optional): Number of folds to generate.
        """

        def are_any_elements_present(row):
            return 1 if any(element in test_positions for element in row['mut_positions']) else 0

        i = 0
        test_size = test_size_sample
        random.seed(self.random_state)
        while not (test_size_min < test_size < test_size_max) and i < sample_iter:
            i += 1
            positions = [sublist for sublist in self.data['mut_positions'].values]
            test_len = int(test_size_sample * len(positions))
            test_positions_ls = random.sample(positions, test_len)
            test_positions = [item for sublist in test_positions_ls for item in sublist]
            self.data['group'] = self.data.apply(are_any_elements_present, axis=1)
            test_size = self.data['group'].sum() / len(self.data)

        if (test_size_min < test_size < test_size_max):
            print(f'Test set size ({round(test_size,2)}) passes the recommended requirements (i.e. between {test_size_min} and {test_size_max}).')
        elif test_size < test_size_min: 
            print(f'Test set size ({round(test_size,2)}) is lower than the recommended minimum size ({test_size_min}). If necessary, rerun with the same or higher test set sample size.')
        elif test_size > test_size_max:
            print(f'Test set size ({round(test_size,2)}) is higher than the recommended maximum size ({test_size_max}). If necessary, rerun with the same or lower test set sample size.')

        # Save splits
        self.split_type = 'position'

        if k_folds is not None:
            if self.val_split is not None:
                self._save_folds(k_folds=k_folds)
            elif self.val_split is None:
                print("No validation split, cannot generate kfolds")
                return

        else:
            
            if iter is None:
                self._save_splits()
            else:
                self._save_splits(iter=iter)
    
    
class RegionProteinSplitter(ProteinSplitter):
    """
    Class for splitting protein datasets based on mutation positions.

    Attributes:
        data (pd.DataFrame): The protein dataset.
        random_state (int): Random seed for reproducibility.
        splits (dict): Dictionary to store the data splits.
        split_type (str): Type of split being performed.
        file_attrs (dict): Dictionary containing file attributes.
        use_cache (bool): Whether to use cached splits.
        y_scaling (str): Method for scaling y values.
        val_split (float): Fraction of data to use for validation.

    Example Usage:

    splitter = RegionProteinSplitter(training_dataset_fname, 
                                    wt_file, 
                                    csv_has_header=True,  # Whether input CSV has a header row
                                    use_cache=True,       # Whether to cache results to disk
                                    y_scaling=True,       # Whether to scale y values to [0,1]
                                    val_split=None        # Fraction of data for validation (None=no validation split)
                                    )
    splitter.split_data(region=[1, 60]) # Splits data such that all variants containing mutations in the first 60 positions are in the test set
    """

    def split_data(self, region, iter=None, k_folds=None):
        """
        Exclude a region or domain of a protein into the test set, the remaining regions are placed into the test set.
        
        Args:
        - region (list): Provided as a 2-number list defining the boundaries of the region to exclude (e.g. [1, 60]).
        - iter (int, optional): Iteration number for naming the split file.
        - k_folds (int, optional): Number of folds to generate.
        """

        region_i = region[0]
        region_f = region[1]

        def are_any_mutations_present(row):
            return 1 if any(element in range(region_i, region_f + 1, 1) for element in row['mut_positions']) else 0

        self.data['group'] = self.data.apply(are_any_mutations_present, axis=1)
        
        # Save splits
        self.split_type = f'region_{region[0]}-{region[1]}'

        if k_folds is not None:
            if self.val_split is not None:
                self._save_folds(k_folds=k_folds)
            elif self.val_split is None:
                print("No validation split, cannot generate kfolds")
                return

        else:
            if iter is None:
                self._save_splits()
            else:
                self._save_splits(iter=iter)
    
class PropertyProteinSplitter(ProteinSplitter):
    """
    Class for splitting protein datasets by value.

    Attributes:
        data (pd.DataFrame): The protein dataset.
        random_state (int): Random seed for reproducibility.
        splits (dict): Dictionary to store the data splits.
        split_type (str): Type of split being performed.
        file_attrs (dict): Dictionary containing file attributes.
        use_cache (bool): Whether to use cached splits.
        y_scaling (str): Method for scaling y values.
        val_split (float): Fraction of data to use for validation.

    Example Usage:

    splitter = PropertyProteinSplitter(training_dataset_fname, 
                                    wt_file, 
                                    csv_has_header=True,  # Whether input CSV has a header row
                                    use_cache=True,       # Whether to cache results to disk
                                    y_scaling=True,       # Whether to scale y values to [0,1]
                                    val_split=None        # Fraction of data for validation (None=no validation split)
                                    )
    splitter.split_data(
        property=0.5,      # Value to split on (e.g. 0.5 for median split)
        above_or_below='above'  # 'above': variants with y > property in test set
                               # 'below': variants with y < property in test set
    )
    """
    
    def split_data(self, property, above_or_below, iter=None, k_folds=None):
        """
        Splits data by the property represented by the y values.
        
        Args:
        - property (float): Value of property to split on.
        - above_or_below (str): 'above' or 'below', values to leave out into the test set based on the given property value.
        - iter (int, optional): Iteration number for naming the split file.
        - k_folds (int, optional): Number of folds to generate.
        """

        if above_or_below == 'above':
            self.data['group'] = np.where(self.data[1] > property, 1, 0)
        
        elif above_or_below == 'below':
            self.data['group'] = np.where(self.data[1] < property, 1, 0)

        # Save splits
        self.split_type = f'y_{above_or_below}_{property}'

        if k_folds is not None:
            if self.val_split is not None:
                self._save_folds(k_folds=k_folds)
            elif self.val_split is None:
                print("No validation split, cannot generate kfolds")
                return
    
        else:
            if iter is None:
                self._save_splits()
            else:
                self._save_splits(iter=iter)

class MutLoadProteinSplitter(ProteinSplitter):
    """
    Class for splitting protein datasets by mutational load.

    Attributes:
        data (pd.DataFrame): The protein dataset.
        random_state (int): Random seed for reproducibility.
        splits (dict): Dictionary to store the data splits.
        split_type (str): Type of split being performed.
        file_attrs (dict): Dictionary containing file attributes.
        use_cache (bool): Whether to use cached splits.
        y_scaling (str): Method for scaling y values.
        val_split (float): Fraction of data to use for validation.

    Example Usage:

    splitter = MutLoadProteinSplitter(training_dataset_fname, 
                                    wt_file, 
                                    csv_has_header=True,  # Whether input CSV has a header row
                                    use_cache=True,       # Whether to cache results to disk
                                    y_scaling=True,       # Whether to scale y values to [0,1]
                                    val_split=None        # Fraction of data for validation (None=no validation split)
                                    )
    splitter.split_data(
        max_train_muts=2,      # Maximum number of mutations to include in training set
        min_test_muts=5        # Minimum number of mutations to include in test set
    )
    """
    
    def split_data(self, max_train_muts, min_test_muts, iter=None, k_folds=None):
        """
        Splits data into training and test sets based on mutational load.
        
        Args:
        - max_train_muts (int): Maximum mutational load to include in the training set.
        - min_test_muts (int): Minimum mutational load to include in the test set.
        - iter (int, optional): Iteration number for naming the split file.
        - k_folds (int, optional): Number of folds to generate.
        """
        assert 'mut_load' in self.data.columns, "DataFrame must contain a 'mut_load' column"
        assert max_train_muts < min_test_muts, "Maximum training mutational load must be less than minimum test mutational load"

        self.data['group'] = np.where(
            self.data['mut_load'] <= max_train_muts, 0,
            np.where(
                self.data['mut_load'] >= min_test_muts, 1, np.nan
            )
        )

        # Save splits
        self.split_type = f'muts-{max_train_muts}-{min_test_muts}'

        if k_folds is not None:
            if self.val_split is not None:
                self._save_folds(k_folds=k_folds)
            elif self.val_split is None:
                print("No validation split, cannot generate kfolds")
                return

        else:
            if iter is None:
                self._save_splits()
            else:
                self._save_splits(iter=iter)

class ResidueDistanceSplitter(ProteinSplitter):
    """
    Class for splitting protein datasets based on residue distances in 3D structure.

    Attributes:
        data (pd.DataFrame): The protein dataset.
        random_state (int): Random seed for reproducibility.
        splits (dict): Dictionary to store the data splits.
        split_type (str): Type of split being performed.
        file_attrs (dict): Dictionary containing file attributes.
        use_cache (bool): Whether to use cached splits.
        y_scaling (str): Method for scaling y values.
        val_split (float): Fraction of data to use for validation.
        pdb_file (str): Path to PDB/CIF structure file.
        chain_ids (list): List of chain IDs to analyze.
        dist_dict (dict): Dictionary mapping mutation pairs to distances.

    Example Usage:

    splitter = ResidueDistanceSplitter(training_dataset_fname, 
                                    wt_file, 
                                    csv_has_header=True,  # Whether input CSV has a header row
                                    use_cache=True,       # Whether to cache results to disk
                                    y_scaling=True,       # Whether to scale y values to [0,1]
                                    val_split=None,       # Fraction of data for validation (None=no validation split)
                                    pdb_file='1abc.pdb', # Path to structure file
                                    chain_ids=['A','B']  # Chain IDs to analyze
                                    )
    splitter.split_data(
        percentile_threshold=50,  # Distance percentile threshold for training set
        min_test_muts=5,         # Minimum mutations for test set
        max_train_muts=2         # Maximum mutations for training set
    )
    """

    def __init__(self, protein_name, data, wt_file, csv_has_header=False, use_cache=False, 
                 y_scaling=False, 
                 val_split=None,
                 random_state=42,
                 pdb_file=None,
                 chain_ids=None,
                 **kwargs):
        """
        Args:
            data (str or pd.DataFrame): Input data containing sequences and labels.
            wt_file (str or list): Path(s) to wild-type sequence file(s).
            csv_has_header (bool): Whether input CSV has header.
            use_cache (bool): Whether to cache results.
            y_scaling (bool): Whether to scale y values.
            val_split (float): Fraction of data for validation.
            random_state (int): Random seed.
            pdb_file (str): Path to PDB/CIF structure file.
            chain_ids (list): List of chain IDs to analyze.
            **kwargs: Additional keyword arguments.
        """
        super().__init__(protein_name, data, wt_file, csv_has_header=csv_has_header, use_cache=use_cache, 
                         random_state=random_state,
                         y_scaling=y_scaling, 
                         val_split=val_split,
                         **kwargs)
        
        self.pdb_file = pdb_file
        self.chain_ids = chain_ids

    def _calculate_ca_distances(self):
        """
        Calculate pairwise distances between all alpha carbons in a protein structure.
        
        Calculates distances between CA atoms and stores in self.dist_dict mapping
        mutation pairs to their 3D distance in Angstroms.
        """
        # Initialize PDB parser
        if self.pdb_file.endswith(".pdb"):
            parser = PDB.PDBParser(QUIET=True)
        elif self.pdb_file.endswith(".cif"):
            parser = PDB.MMCIFParser(QUIET=True)
        else:
            raise ValueError("Invalid file type. Please provide a PDB or CIF file.")
        
        structure = parser.get_structure('protein', self.pdb_file)
        
        # Get all alpha carbons
        ca_atoms = []
        residue_info = []
        
        for model in structure:
            for chain in model:
                if chain.id in self.chain_ids:
                    for residue in chain:
                        if 'CA' in residue:
                            ca_atoms.append(residue['CA'])
                            residue_info.append((
                                chain.id,
                                residue.get_id()[1],  # residue number
                                residue.get_resname()  # residue name
                            ))
        
        # Calculate distance matrix
        n_residues = len(ca_atoms)
        distance_matrix = np.zeros((n_residues, n_residues))
        
        for i in range(n_residues):
            for j in range(i+1, n_residues):
                distance = ca_atoms[i] - ca_atoms[j]  # Returns distance in Angstroms
                distance_matrix[i,j] = distance
                distance_matrix[j,i] = distance
        
        dist_dict = {}
        
        for i in range(len(residue_info)):
            for j in range(i+1, len(residue_info)):  # Only upper triangle to avoid duplicates
                chain_i, resnum_i, resname_i = residue_info[i]
                chain_j, resnum_j, resname_j = residue_info[j]
                resname_i = aa_dict_3to1[resname_i]
                resname_j = aa_dict_3to1[resname_j]
                
                # Calculate adjusted residue numbers based on chain index
                resnum_i_adj = resnum_i + self.wt_seq_lens[self.chain_ids.index(chain_i)-1] if self.chain_ids.index(chain_i) != 0 else resnum_i
                resnum_j_adj = resnum_j + self.wt_seq_lens[self.chain_ids.index(chain_j)-1] if self.chain_ids.index(chain_j) != 0 else resnum_j
                
                # Create key string with adjusted residue numbers
                key = f'{resname_i}{resnum_i_adj}_{resname_j}{resnum_j_adj}'
                
                # Store distance in dictionary
                dist_dict[key] = distance_matrix[i,j]

        if self.randomized_control:
            # modify dist_dict to be randomized
            # Set random seed before shuffling
            random.seed(self.random_state)
            # Get all values and shuffle them
            values = list(dist_dict.values())
            random.shuffle(values)
            # Reassign shuffled values to the same keys
            dist_dict = dict(zip(dist_dict.keys(), values))
            
        self.dist_dict = dist_dict

        self.data['dist'] = self.data['muts'].apply(lambda x: self._get_dist(x.split('/')))

        self._get_dist_percentile()

    def _get_dist(self, muts):
        """
        Get sum of pairwise distances between mutations.
        
        Args:
            muts (list): List of mutation strings in format 'A123B'.
            
        Returns:
            float: Sum of pairwise distances between mutations.
        """
        distances = []
        for i in range(len(muts)):
            for j in range(i+1, len(muts)):
                mut_pair = f"{muts[i][:-1]}_{muts[j][:-1]}"
                if mut_pair in self.dist_dict:
                    distances.append(self.dist_dict[mut_pair])
        return sum(distances)
    
    def _get_dist_percentile(self):
        """
        Calculate distance percentile for each variant within its mutational load group.
        Updates self.data with 'dist_percentile' column.
        """
        for mut_load in self.data['mut_load'].unique():
            subset = self.data[self.data['mut_load'] == mut_load].copy()
            if mut_load == 0 or mut_load == 1:
                subset['dist_percentile'] = 0
            else:
                subset['dist_percentile'] = subset['dist'].rank(pct=True) *100
            self.data.loc[self.data['mut_load'] == mut_load, 'dist_percentile'] = subset['dist_percentile']

    def split_data(self, percentile_threshold, min_test_muts, max_train_muts, randomized_control=False, iter=None, k_folds=None):
        """
        Split data based on mutation distances and counts.
        
        Args:
            percentile_threshold (float): Maximum distance percentile for training set.
            min_test_muts (int): Minimum mutations for test set.
            max_train_muts (int): Maximum mutations for training set.
            randomized_control (bool): Whether to randomize the distance dictionary.
            iter (int, optional): Iteration number for naming the split file.
            k_folds (int, optional): Number of folds to generate.
        """
        assert 'mut_load' in self.data.columns, "DataFrame must contain a 'mut_load' column"
        assert max_train_muts < min_test_muts, "Maximum training mutational load must be less than minimum test mutational load"

        self.randomized_control = randomized_control
        self._calculate_ca_distances()

        self.data['group'] = np.where(
            (self.data['mut_load'] <= max_train_muts) & (self.data['dist_percentile'] <= percentile_threshold), 0,
            np.where(
                self.data['mut_load'] >= min_test_muts, 1, np.nan
            )
        )

        # Save splits
        self.split_type = f'dist-p{percentile_threshold}-{max_train_muts}-{min_test_muts}{"-randomized-" + str(self.random_state) if self.randomized_control else ""}'

        if k_folds is not None:
            if self.val_split is not None:
                self._save_folds(k_folds=k_folds)
            elif self.val_split is None:
                print("No validation split, cannot generate kfolds")
                return
        else:
            if iter is None:
                self._save_splits()
            else:
                self._save_splits(iter=iter)