data_processing.py

import re
import os
import ast
import json

import numpy as np
import pandas as pd

from tqdm import tqdm
from spacy import displacy

from sklearn.model_selection import train_test_split, StratifiedKFold
from imblearn.over_sampling import RandomOverSampler
from imblearn.under_sampling import RandomUnderSampler

from prediction_properties import PredictionProperties

class DataProcessing:
    """A class to preprocess data"""

    def concat_dfs(dfs: list[pd.DataFrame], axis: int = 0, ignore_index: bool = True) -> pd.DataFrame:
        """Concatenate multiple DataFrames"""
        df = pd.concat(dfs, axis=axis, ignore_index=ignore_index)
        return df
    
    def shuffle_df(df: pd.DataFrame, random_state):
        """Shuffle the data"""
        df = df.sample(frac=1, random_state=random_state).reset_index(drop=True)
        return df
    
    def split_data(
            features_df, 
            labels_df,
            random_state,
            val_size=None, 
            test_size=None,
            stratify_by=None, 
            stratify_kfold: int = None,
            **kwargs):
        """
        Split features and labels into train/test or train/val/test sets.
        
        Parameters
        ----------
        features_df : pd.DataFrame or np.ndarray
            Features to split
        labels_df : pd.DataFrame
            Label columns (e.g., ['Sentence Label', 'Author Type'])
        test_size : float, default=0.2
            Fraction for test set
        val_size : float, default=None
            Fraction for validation set (if None, only train/test split)
        random_state : int, default=42
            Random seed for reproducibility
        stratify_by : str, default=None
            Column name to stratify on (e.g., 'Sentence Label')
            If None, no stratification is used
        
        Returns
        -------
        tuple
            If val_size is None: (X_train, X_test, y_train, y_test)
            If val_size is set: (X_train, X_val, X_test, y_train, y_val, y_test)
        
        Examples
        --------
        # 2-way split with stratification
        X_train, X_test, y_train, y_test = split_data(
            features, labels, stratify_by='Sentence Label'
        )
        
        # 3-way split (60% train, 20% val, 20% test)
        X_train, X_val, X_test, y_train, y_val, y_test = split_data(
            features, labels, val_size=0.2, stratify_by='Sentence Label'
        )
        """
        stratify = labels_df[stratify_by] if stratify_by else None

        # 1) CHECK K-FOLD FIRST
        if stratify_kfold:
            skf = StratifiedKFold(n_splits=stratify_kfold, shuffle=True, random_state=random_state)

            all_folds = []

            for train_idx, val_idx in skf.split(features_df, labels_df):
                X_train = features_df.iloc[train_idx]
                X_val = features_df.iloc[val_idx]
                y_train = labels_df.iloc[train_idx]
                y_val = labels_df.iloc[val_idx]
                
                all_folds.append((X_train, X_val, y_train, y_val))
                
            return all_folds
        
        # No validation set - simple 2-way split
        if val_size is None:
            return train_test_split(
                features_df, 
                labels_df, 
                test_size=test_size, 
                random_state=random_state,
                stratify=stratify
            )
        
        # With validation set - two splits
        # Split 1: Remove test set
        X_temp, X_test, y_temp, y_test = train_test_split(
            features_df, 
            labels_df,
            test_size=test_size,
            random_state=random_state,
            stratify=stratify
        )
        
        # Split 2: Split remaining into train and val
        val_size_adjusted = val_size / (1 - test_size)
        stratify_temp = y_temp[stratify_by] if stratify_by else None
        
        X_train, X_val, y_train, y_val = train_test_split(
            X_temp, 
            y_temp,
            test_size=val_size_adjusted,
            random_state=random_state,
            stratify=stratify_temp
        )
        
        return X_train, X_val, X_test, y_train, y_val, y_test

    def join_predictions_with_labels(df: pd.DataFrame, true_labels: pd.Series, y_predictions: pd.Series, model) -> pd.DataFrame:
        """Join the predictions with the true labels DF
        
        Parameters:
        -----------
        df: `pd.DataFrame`
            A DataFrame containing the data

        true_labels: `pd.DataFrame`
            A DataFrame containing the true labels

        y_predictions: `pd.DataFrame`
            A DataFrame containing the predictions from the model

        model: `object`
            An instance of the model
        
        Returns:
        --------
        `pd.DataFrame`
            A DataFrame containing the sentence, labels, and predictions
        """
        assert len(true_labels) == len(y_predictions), "The length of the true labels and predictions must be the same"

        N = len(df) - len(y_predictions)
        test_df = df.loc[N:].copy()
        model_name = model.__name__()
        model_col_name = f"{model_name} Prediction"
        test_df.loc[:, model_col_name] = y_predictions.values

        # Move "Prediction Label" column to second to last position
        cols = list(test_df.columns)
        cols.remove("Prediction Label")
        cols.insert(-1, "Prediction Label")
        test_df = test_df[cols]
        
        return test_df
    
    def join_predictions_with_sentences(df: pd.DataFrame, y_predictions: pd.Series, model) -> pd.DataFrame:
        """Join the model predictions with the sentences. In this case, no true labels are provided
        
        Parameters:
        -----------
        df: `pd.DataFrame`
            A DataFrame containing the data

        y_predictions: `pd.DataFrame`
            A DataFrame containing the predictions from the model

        model: `object`
            An instance of the model
        
        Returns:
        --------
        `pd.DataFrame`
            A DataFrame containing the sentence, labels, and predictions
        """
        assert len(df) == len(y_predictions), "The length of the true labels and predictions must be the same"

        joint_df = df.copy()
        model_name = model.__name__()
        model_col_name = f"{model_name} Prediction"
        joint_df.loc[:, model_col_name] = y_predictions.values
        
        return joint_df
    
    def reformat_df_with_template_number(df: pd.DataFrame, prediction_templates: list, col_name: str) -> pd.DataFrame:
        """Reformat the DataFrame with the template number
        
        Parameters:
        -----------
        df: `pd.DataFrame`
            A DataFrame containing the data
        
        col_name: `str`
            The column name to use for the template number
        
        Returns:
        --------
        `pd.DataFrame`
            A DataFrame with the template number in the specified column
        """
        
        template_numbers = []
        template_texts = []
        reformat_predictions = []
        indices_to_keep = []

        for idx, prediction in enumerate(df[col_name].values):
            first_word = prediction.split()[0]
            if first_word == "T0:":
                template_numbers.append(0)
                reformat_predictions.append(prediction[4:])
                indices_to_keep.append(idx)
            elif first_word == "T1:":
                template_numbers.append(1)
                template_texts.append(prediction_templates[0])
                reformat_predictions.append(prediction[4:])
                indices_to_keep.append(idx)
            elif first_word == "T2:":
                template_numbers.append(2)
                template_texts.append(prediction_templates[1])
                reformat_predictions.append(prediction[4:])
                indices_to_keep.append(idx)
            elif first_word == "T3:":
                template_numbers.append(3)
                template_texts.append(prediction_templates[2])
                reformat_predictions.append(prediction[4:])
                indices_to_keep.append(idx)
            elif first_word == "T4:":
                template_numbers.append(4)
                template_texts.append(prediction_templates[3])
                reformat_predictions.append(prediction[4:])
                indices_to_keep.append(idx)
            elif first_word == "T5:":
                template_numbers.append(5)
                template_texts.append(prediction_templates[4])
                reformat_predictions.append(prediction[4:])
                indices_to_keep.append(idx)
            elif first_word == "T6:":
                template_numbers.append(6)
                template_texts.append(prediction_templates[5])
                reformat_predictions.append(prediction[4:])
                indices_to_keep.append(idx)
            else:
                continue
        
        new_df = df.iloc[indices_to_keep].copy()
        new_df[col_name] = reformat_predictions
        new_df['Template Number'] = template_numbers
        new_df['Template Text'] = template_texts
        return new_df
    
    def df_to_list(df: pd.DataFrame, col: str = None, type_of_df: str = "Standard") -> list:
        """Convert a DataFrame to a list
        
        Parameters:
        -----------
        df: `pd.DataFrame`
            A DataFrame containing the data
        
        col: `str`
            The column to convert to a list

        type_of_df: `str`
            The type of DataFrame either Standard or Pivot Table

        Returns:
        --------
        list
            A list containing the data from the DataFrame
        """
        if type_of_df == "Standard":
            return df[col].values.tolist()
        elif type_of_df == "Pivot Table":
            return df.index.values.tolist()
        else:
            return "Options: Standard or Pivot Table"

    def drop_df_columns(df: pd.DataFrame, columns: list):
        """Drop columns
        
        Parameters
        ----------
        df: pd.DataFrame
            The DataFrame to drop columns from
        
        columns: list 
            The list of columns to drop from the DataFrame
        
        Returns
        -------
        df: pd.DataFrame
            The DataFrame with the columns dropped
        """
        df = df.drop(columns=columns)
        return df

    def encode_tags_entities_df(df: pd.DataFrame, sentence_and_label_df: pd.DataFrame) -> pd.DataFrame:
        """Encode the tags or entities in the DataFrame. Use 1 for the presence of the tag or entity and 0 if NaN
        
        Parameters
        ----------
        df: pd.DataFrame
            The DataFrame to encode

        sentence_and_label_df: pd.DataFrame
            The DataFrame containing the sentence and prediction labels

        Returns
        -------
        encoded_df: pd.DataFrame
            The DataFrame with the tags or entities encoded
        """
        bool_df = df.notnull() # Convert the DataFrame to boolean where if presence, place True and NaN place False
        encoded_df = bool_df.astype('int') # Convert the boolean DataFrame to integer where True is 1 and False is 0
        updated_encoded_df = DataProcessing.include_sentence_and_label(encoded_df, sentence_and_label_df)

        return updated_encoded_df

    def array_to_df(data: np.array) -> pd.DataFrame:
        """Convert a numpy array to a DataFrame
        
        Parameters:
        -----------
        data: `np.array`
            An array containing the data to convert to a DataFrame

        Returns:
        --------
        `pd.Series`
            A Series containing the data
        """
        return pd.Series(data)
    
    def convert_tags_entities_to_dataframe(keys_of_mappings: set, mappings: list[list]) -> pd.DataFrame:
        """
        Convert extracted features (Part-of-Speech (POS) tags and Named Entities Recognition (NER)) into a pandas DataFrame.
        
        Parameters:
        -----------
        entities : list
            A list of entities extracted from documents.
        
        all_ner_entities : list
            A list of all unique NER (Named Entity Recognition) tags.

        Returns:
        --------
        pd.DataFrame
            A DataFrame containing entities organized by NER tags.
        """
        df_ner = pd.DataFrame(columns=list(keys_of_mappings))
        for i, document_mapping in tqdm(enumerate(mappings)):
            for text, label in document_mapping:
                df_ner.at[i, label] = text
        return df_ner
    
    def include_sentence_and_label(df_to_update: pd.DataFrame, sentence_and_label_df: pd.DataFrame) -> pd.DataFrame:
        """Include the sentence and prediction labels in the DataFrame"""

        if len(df_to_update) == len(sentence_and_label_df):
            df_to_update.insert(0, 'Base Sentence', sentence_and_label_df['Base Sentence'].values)
            df_to_update.insert(1, 'Sentence Label', sentence_and_label_df['Sentence Label'].values)
            return df_to_update
        else:
            print(f"Error: The lengths of df_to_update ({len(df_to_update)}) and sentence_and_label_df ({len(sentence_and_label_df)}) do not match.")
            return df_to_update
        
    def convert_to_df(data, mapping=None):
        """Convert data to a DataFrame or Series.

        Parameters:
        -----------
        data: `np.array`, `list`, `dict`, or `set`
            An array, list, dictionary, or set containing the data to convert to a DataFrame or Series.
        
        mapping: `any`, `optional`
            convert_tags_entities_to_dataframe() --- A dictionary containing mappings for tags/entities


        Returns:
        --------
        `pd.DataFrame` or `pd.Series`
            A DataFrame or Series containing the data.
        """
        if isinstance(data, np.ndarray):
            return DataProcessing.array_to_df(data)
        elif isinstance(data, set) and isinstance(mapping, list):
            return DataProcessing.convert_tags_entities_to_dataframe(data, mapping)
        elif isinstance(data, list) and mapping == 'Open Measures':
            sources = []
            for hit in data:
                sources.append(hit['_source'])
            return pd.DataFrame(sources)
        # elif isinstance(data, tuple[int, json]):
        #     return DataProcessing.json_to_pd(data)
        else:
            raise ValueError("Invalid input: data must be a numpy array, dictionary, list, or set with a mapping.")
    
    def load_prediction_properties():
        return PredictionProperties.get_prediction_properties()
        
    def patterns(nlp):
        # ruler = nlp.add_pipe("entity_ruler", before="ner")
        # patterns = [
        #     {"label": "DATE", "pattern": [{"TEXT": {"REGEX": "20\\d{2}/\\d{2}/\\d{2}"}}]},
        #     {"label": "DATE", "pattern": [{"TEXT": {"REGEX": "20\\d{2}"}}, {"TEXT": {"REGEX": "Q[1-4]"}}]},
        #     {"label": "GPE", "pattern": [{"TEXT": {"REGEX": "LOC"}}, {"TEXT": {"REGEX": "LOC_\d"}}]},
        #     {"label": "DATE", "pattern": [{"TEXT": {"REGEX": "\d{1,2} [a-zA-Z]+ \d{4}"}}]} # NOT WORKING for ex 15 November 2021
        # ]
        # ruler.add_patterns(patterns)
        pass

    def update_ner(label, text):
        """Updates the NER label based on provided rules."""
        if label == "LOC" or label.startswith("LOC_"):
            return "GPE"
        elif label == "FAC":
            return "PERSON"
        elif label == "WORK_OF_ART":
            if re.match(r"Q\d of \d{4}|Quarter \d|\d Q\d", text):
                return "DATE"
            elif text in ["Weather Underground", "The Weather Channel]"]:
                return "ORG"
            else:
                return label  # Keep original label if no match
        elif label == "CARDINAL":
            if re.match(r"\d{1,2} [a-zA-Z]+ \d{4}", text):
                return "DATE"
            elif re.match(r"-\d+°[CF]|-\d+°[CF] to -\d+°[CF]|\d+°[CF] to \d+°[CF]|\d+°[CF]", text):
                return "TEMPERATURE"
            else: return label
        else:
            return label
    
    def select_pattern(template_number: int) -> str:
        """Select the pattern to use based on the template number
        
        Parameters:
        -----------
        template_number: `int`
           The template number that corresponds to the pattern to use

        Returns:
        --------
        `str`
            A string containing the pattern to use
        """

        if template_number == 1:
            pattern = r"On \[(.*?)\], \[(.*?)\] (.*?) that the \[(.*?)\] at \[(.*?)\] \[(.*?)\] \[(.*?)\] by \[(.*?)\] in \[(.*?)\]."
            return pattern
        elif template_number == 2:
            pattern = r"In \[(.*?)\], \[(.*?)\] from \[(.*?)\], (.*?) that the \[(.*?)\] \[(.*?)\] \[(.*?)\] from \[(.*?)\] to \[(.*?)\] in \[(.*?)\]."
            return pattern
        elif template_number == 3:
            pattern = r"\[(.*?)\] (.*?) on \[(.*?)\] that the \[(.*?)\] in \[(.*?)\] \[(.*?)\] \[(.*?)\] by \[(.*?)\] in \[(.*?)\]."
            return pattern
        elif template_number == 4:
            pattern = r"According to a \[(.*?)\] from \[(.*?)\], on \[(.*?)\], the \[(.*?)\] \[(.*?)\] \[(.*?)\] beyond \[(.*?)\] in the timeframe of \[(.*?)\]."
            return pattern
        elif template_number == 5:
            pattern = r"In \[(.*?)\], the \[(.*?)\] in \[(.*?)\] \[(.*?)\] \[(.*?)\] by \[(.*?)\], as (.*?) by \[(.*?)\] on \[(.*?)\]."
            return pattern
        else:
            raise ValueError("The template number is not recognized.")
    
    def sentence_to_df(df: pd.DataFrame) -> pd.DataFrame:
        """Convert an example sentence to a DataFrame

        NOTE: This function is specific to the template used in the example sentence
        
        Parameters:
        -----------
        sentence: `str`
            A sentence containing the variables to extract

        Returns:
        --------
        `pd.DataFrame`
            A DataFrame containing the extracted variables
        """
        template_numbers = df['Template Number'].values
        base_sentences = df['Base Sentence'].values
        
        extracted_data = []

        for template_number, sentence in zip(template_numbers, base_sentences):
            print(sentence)
            pattern = DataProcessing.select_pattern(template_number)
            print(pattern)
            
            match = re.match(pattern, sentence)
            
            if match:
                p_t, p_p, p_a, p_o, p_v, p_s, p_m, p_f = match.groups()
                data = {
                    'p_p': p_p,
                    'p_o': p_o,
                    'p_t': p_t,
                    'p_f': p_f,
                    'p_a': p_a,
                    'p_s': p_s,
                    'p_m': p_m,
                    'p_v': p_v,
                    'p_l': None  # Assuming y_l is not used in this template
                }
                extracted_data.append(data)
            else:
                print(sentence)
                raise ValueError("The sentence does not match the expected template.")
        
        return pd.DataFrame(extracted_data)
    
    def visualize_spacy_doc(doc, save_dir: str = "../paper/exSentence.png"):

        from PIL import Image
        import io
        import cairosvg

        options = {"compact": True, 
                   "bg": "#09a3d5", 
                   "font": "Source Sans Pro"}
        svg = displacy.render(doc, style="dep", jupyter=False, options=options)

            # Convert SVG to PNG
        png_data = cairosvg.svg2png(bytestring=svg.encode('utf-8'))
        img = Image.open(io.BytesIO(png_data))
        img.save(save_dir)

    def get_next_file_number(directory: str, prefix: str, extensions: str = ('.json', '.log', '.csv', '.png')):
        """
        Determine the next available file number based on existing files in a directory.
        
        Parameters
        ----------
        directory : str
            Path to the directory where files are stored.
        prefix : str
            The prefix used in filenames (e.g., 'siteA' for files like 'siteA-v3.json').
        extensions : tuple of str, optional
            File extensions to consider when scanning for existing files. Default is ('.json', '.log', '.csv', '.png').
        
        Returns
        -------
        int
            The next available file number (e.g., returns 4 if files with numbers 1, 2, and 3 exist).
        """
        if not os.path.exists(directory):
            return 1
        
        numbers = []
        for name in os.listdir(directory):
            full_path = os.path.join(directory, name)
            
            # Only check files
            if not os.path.isfile(full_path):
                continue
            
            if name.startswith(prefix) and name.endswith(extensions):
                try:
                    # Assumes format: prefix-vN.ext (e.g., siteA-v3.json)
                    after_prefix = name[len(prefix):]  # e.g., "-v3.json"
                    without_ext = after_prefix.rsplit('.', 1)[0]  # e.g., "-v3"
                    
                    # Extract just the number (handle both "-v3" and "-3" formats)
                    if without_ext.startswith('-v'):
                        number_part = without_ext[2:]  # Remove "-v"
                    elif without_ext.startswith('-'):
                        number_part = without_ext[1:]  # Remove "-"
                    else:
                        continue
                    
                    number = int(number_part)
                    numbers.append(number)
                    
                except (ValueError, IndexError):
                    continue
        
        next_num = max(numbers, default=0) + 1
        return next_num

    def get_next_directory_number(directory: str, prefix: str, date_filter: str = None):
        """
        Determine the next available directory number based on existing directories.
        
        Parameters
        ----------
        directory : str
            Path to the parent directory where subdirectories are stored.
        prefix : str
            The prefix used in directory names (e.g., 'undersampled_96d-v4_').
        date_filter : str, optional
            Filter directories by date (e.g., '2026-02-17'). Only counts directories with this date.
        
        Returns
        -------
        int
            The next available directory number.
        
        Notes
        -----
        Expected directory format: {prefix}v{N}_{date}
        Example: undersampled_96d-v4_v2_2026-02-17
        """
        if not os.path.exists(directory):
            return 1
        
        numbers = []
        print(f"\n[DEBUG] Looking for directories with prefix: '{prefix}' and date: '{date_filter}'")
        
        for name in os.listdir(directory):
            full_path = os.path.join(directory, name)
            
            # Only check directories
            if not os.path.isdir(full_path):
                continue
            
            print(f"[DEBUG] Checking directory: '{name}'")
            
            # Check if starts with prefix
            if name.startswith(prefix):
                # If date filter specified, check if directory ends with that date
                if date_filter and not name.endswith(date_filter):
                    print(f"[DEBUG] ✗ Doesn't match date filter '{date_filter}'")
                    continue
                
                print(f"[DEBUG] ✓ Matches prefix and date")
                try:
                    # Extract part after prefix: "v1_2026-02-17" or "v2_2026-02-17"
                    after_prefix = name[len(prefix):]
                    print(f"[DEBUG] After prefix: '{after_prefix}'")
                    
                    # Split by underscore: ["v1", "2026-02-17"]
                    parts = after_prefix.split('_')
                    print(f"[DEBUG] Parts: {parts}")
                    
                    # First part should be version: "v1", "v2", etc.
                    if parts and parts[0].startswith('v') and len(parts[0]) > 1:
                        number_part = parts[0][1:]  # Remove 'v'
                        number = int(number_part)
                        numbers.append(number)
                        print(f"[DEBUG] Found version number: {number}")
                            
                except (ValueError, IndexError) as e:
                    print(f"[DEBUG] ✗ Error extracting number: {e}")
                    continue
            else:
                print(f"[DEBUG] ✗ Doesn't match prefix")
        
        next_num = max(numbers, default=0) + 1
        print(f"[DEBUG] Numbers found: {numbers}")
        print(f"[DEBUG] Next number: {next_num}\n")
        return next_num
    
    def save_to_file(data, 
                    path: str, 
                    prefix: str, 
                    save_file_type: str, 
                    include_version: bool = True,
                    append: bool = False,
                    **kwargs: dict) -> None:
        """ 
        Save data to any file with an incremented filename based on existing files.

        Parameters
        ----------
        data : dict, list, or pd.DataFrame
            The data to be saved.
        path : str
            Directory path where the file will be saved.
        prefix : str
            Prefix for the filename (e.g., 'results' → 'results.csv' or 'results-v1.csv').
        save_file_type : str
            File type: 'json', 'csv', or 'png'.
        include_version : bool, optional
            If True, appends version number to filename (e.g., results-v1.csv).
            If False, saves directly without version suffix.
            Default is True.
        append : bool, optional
            If True, appends to an existing CSV file instead of overwriting.
            Only applies to CSV files. If file does not exist, creates it.
            Default is False.
        **kwargs : dict
            Additional arguments for specific file types:
            - For PNG: dpi (default=300), bbox_inches (default='tight')

        Notes
        -----
        append=True is useful for batch processing pipelines where results
        are written incrementally (e.g., every BATCH_SIZE sentences) to the
        same file. The header is only written if the file does not exist yet.
        """
        os.makedirs(path, exist_ok=True)

        if include_version and not append:
            next_number = DataProcessing.get_next_file_number(path, prefix)
            print(f"Using file number: {next_number}")

        if save_file_type == 'json':
            if include_version:
                file_name = f"{prefix}-v{next_number}.json"
            else:
                file_name = f"{prefix}.json"
            file_path = os.path.join(path, file_name)
            print(f"Saving JSON file to: {file_path}")
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=4)

        elif save_file_type in ['csv', '.csv']:
            if append:
                # Append mode — always use the same filename, no versioning
                file_name = f"{prefix}.csv"
                file_path = os.path.join(path, file_name)
                file_exists = os.path.exists(file_path)
                print(f"Appending to CSV file: {file_path}")
                data.to_csv(file_path, mode='a', header=not file_exists, index=False)
            elif include_version:
                file_name = f"{prefix}-v{next_number}.csv"
                file_path = os.path.join(path, file_name)
                print(f"Saving CSV file to: {file_path}")
                data.to_csv(file_path, index=False)
            else:
                file_name = f"{prefix}.csv"
                file_path = os.path.join(path, file_name)
                print(f"Saving CSV file to: {file_path}")
                data.to_csv(file_path, index=False)

        elif save_file_type in ['png', '.png', 'PNG']:
            import matplotlib.pyplot as plt
            if include_version:
                file_name = f"{prefix}-v{next_number}.png"
            else:
                file_name = f"{prefix}.png"
            file_path = os.path.join(path, file_name)
            print(f"Saving PNG file to: {file_path}")
            dpi = kwargs.get('dpi', 300)
            bbox_inches = kwargs.get('bbox_inches', 'tight')
            plt.savefig(file_path, dpi=dpi, bbox_inches=bbox_inches)
            plt.close()

        else:
            raise ValueError(f"Unsupported file type: {save_file_type}. Choose from [json, csv, png]")

    def load_from_file(path: str, 
                    file_type: str = 'csv', 
                    sep = ",", 
                    encoding = 'utf-8',
                    **kwargs
                    ):
        """Load data from directory
        
        Parameters
        ----------
        path : str
            Directory path where the file will be loaded from.
        file_type : str
            File types such as json, csv, etc
        sep : str
            Delimiter for CSV files
        encoding : str
            File encoding
        **kwargs
            Additional keyword arguments passed to pd.read_csv()
            (e.g., header=None, names=['col1', 'col2'], dtype=...)
        
        Returns
        -------
        pd.DataFrame
            Loaded dataframe
        """
        
        if file_type == 'csv': 
            df = pd.read_csv(path, sep=sep, encoding=encoding, **kwargs)
            return df
        elif file_type == 'xlsx': 
            df = pd.read_excel(path, **kwargs)
            return df
        elif file_type == 'json':
            data = pd.read_json(path)
            df = pd.DataFrame(data)
            return df
        else:
            return 'Did not properly load'

    def get_latest_file(directory: str, prefix: str, file_type: str = 'csv') -> str:
        """
        Parameters
        ----------
        directory : str
            Path to the directory to search in.
        prefix : str
            Filename prefix to match (e.g., 'classifications').
        file_type : str, default 'csv'
            File extension to match (without dot).

        Notes
        -----
        Relies on the same versioning format as get_next_file_number()
        (e.g., classifications-v3.csv). Returns the filename with the
        highest version number.

        Returns
        -------
        str
            Filename of the latest versioned file.
        """
        if not os.path.exists(directory):
            raise FileNotFoundError(f"Directory not found: {directory}")

        matched = {}
        for name in os.listdir(directory):
            full_path = os.path.join(directory, name)
            if not os.path.isfile(full_path):
                continue
            if name.startswith(prefix) and name.endswith(f'.{file_type}'):
                try:
                    after_prefix = name[len(prefix):]        # e.g., "-v3.csv"
                    without_ext = after_prefix.rsplit('.', 1)[0]  # e.g., "-v3"
                    if without_ext.startswith('-v'):
                        number = int(without_ext[2:])
                    elif without_ext.startswith('-'):
                        number = int(without_ext[1:])
                    else:
                        continue
                    matched[number] = name
                except (ValueError, IndexError):
                    continue

        if not matched:
            raise FileNotFoundError(
                f"No versioned '{prefix}' files found in: {directory}"
            )

        # Return the filename with the highest version number
        latest_filename = matched[max(matched)]
        return latest_filename

    def remove_duplicates(df):
        filt_no_duplicates = (df.duplicated() == False)
        no_duplicates_df = df[filt_no_duplicates]

        return no_duplicates_df
    
    def json_to_df(data) -> pd.DataFrame:
        """
        Notes
        -----
        Modular to where we can use directly or call convert_to_df().
        Need to update calling in convert_to_df().
        """
        row, value = data
        print(row)
        df = pd.DataFrame(value, index=row)
        return df
    
    def load_base_data_path(notebook_dir: str) -> str:
        """Path to data/"""
        return os.path.join(notebook_dir, "../data")
        
    def _build_batch_path(notebook_dir: str, data_type: str, batch_idx: int) -> list:
        """
        Build the file paths for a specific batch across all annotator directories.
        
        Parameters
        ----------
        notebook_dir : str
            Base notebook directory
        data_type : str
            Either 'prediction' or 'observation'
        batch_idx : int
            Batch index number
        
        Returns
        -------
        list
            List of full paths to the batch CSV files that exist
        """
        if data_type not in ['prediction', 'observation']:
            raise ValueError("data_type must be either 'prediction' or 'observation'")
        
        base_data_path = DataProcessing.load_base_data_path(notebook_dir)
        batch_folder_name = f"batch_{batch_idx}-{data_type}"
        file_name = f"batch_{batch_idx}-from_df.csv"
        
        # Create a list of all potential directories to search
        data_folders = []
        data_folders.append(f"{data_type}_logs")
        data_folders.append(f"{data_type}_logs-abigail")
        data_folders.append(f"{data_type}_logs-adriano")
        
        found_paths = []
        
        # Check each folder to see if the file exists there
        for data_folder in data_folders:
            potential_path = os.path.join(base_data_path, data_folder, batch_folder_name, file_name)
            if os.path.exists(potential_path):
                found_paths.append(potential_path)
                
        # If the loop finishes and no paths were found, raise an error
        if len(found_paths) == 0:
            raise FileNotFoundError(f"Batch {batch_idx} ({file_name}) not found in any log directories.")
            
        return found_paths

    def load_single_synthetic_data(notebook_dir: str, 
                                   batch_idx: int, 
                                   sep: str, 
                                   data_type: str = 'prediction',
                                   return_as: str = 'dataframe'):
        """
        Load a single batch of synthetic data (including all annotators).
        
        Parameters
        ----------
        notebook_dir : str
            Base notebook directory
        batch_idx : int
            Batch index number to load. Default is 7.
        sep : str
            Separator for CSV file
        data_type : str
            Either 'prediction' or 'observation'. Default is 'prediction'.
        return_as : str
            Either 'dataframe' or 'path'. Default is 'dataframe'.
        
        Returns
        -------
        pd.DataFrame or list
            DataFrame containing the batch data or list of path strings
        """
        # This will now get a list of ALL correct paths from the directories for this batch index
        file_paths = DataProcessing._build_batch_path(notebook_dir, data_type, batch_idx)
        
        if return_as.lower() in ['path', 'string']:
            return file_paths
        else:  # default to dataframe
            dataframes_to_combine = []
            
            for file_path in file_paths:
                print(f"Loading: {file_path}")
                df = DataProcessing.load_from_file(file_path, 'csv', sep)
                dataframes_to_combine.append(df)
            
            # Combine all annotators for this specific batch index
            combined_batch_df = DataProcessing.concat_dfs(dataframes_to_combine)
            return combined_batch_df

    def load_multiple_batches(notebook_dir: str, sep: str = ',', data_type: str = 'prediction', 
                            batch_indices: list = None, start_idx: int = 1, 
                            end_idx: int = None, return_as: str = 'dataframe') -> pd.DataFrame:
        """
        Load multiple batches of synthetic data and concatenate them.
        
        Parameters
        ----------
        notebook_dir : str
            Base notebook directory
        sep : str
            Separator for CSV file. Default is ','.
        data_type : str
            Either 'prediction' or 'observation'. Default is 'prediction'.
        batch_indices : list, optional
            Specific list of batch indices to load. If provided, start_idx and end_idx are ignored.
        start_idx : int, optional
            Starting batch index (inclusive). Default is 1.
        end_idx : int, optional
            Ending batch index (inclusive). If None, automatically detects the last available batch.
        return_as : str
            Either 'dataframe' or 'path'. Default is 'dataframe'.
            If 'path', returns list of file paths instead of loading data.
        
        Returns
        -------
        pd.DataFrame or list
            Concatenated DataFrame containing all batch data, or list of file paths
        """
        if data_type not in ['prediction', 'observation']:
            raise ValueError("data_type must be either 'prediction' or 'observation'")
        
        # Determine which batches to load
        if batch_indices is not None:
            indices = batch_indices
        else:
            # Auto-detect the number of available batches if end_idx is None
            if end_idx is None:
                base_data_path = DataProcessing.load_base_data_path(notebook_dir)
                
                log_directories = []
                log_directories.append(os.path.join(base_data_path, f"{data_type}_logs"))
                log_directories.append(os.path.join(base_data_path, f"{data_type}_logs-abigail"))
                log_directories.append(os.path.join(base_data_path, f"{data_type}_logs-adriano"))
                
                valid_directories = []
                for directory in log_directories:
                    if os.path.exists(directory):
                        valid_directories.append(directory)
                        
                if len(valid_directories) == 0:
                    raise FileNotFoundError("No log directories found.")
                
                batch_dirs = []
                for valid_directory in valid_directories:
                    for directory_name in os.listdir(valid_directory):
                        full_path = os.path.join(valid_directory, directory_name)
                        if os.path.isdir(full_path):
                            if directory_name.startswith('batch_'):
                                if directory_name.endswith(f'-{data_type}'):
                                    batch_dirs.append(directory_name)
                
                if len(batch_dirs) == 0:
                    raise ValueError("No batch directories found in the log directories.")
                
                # Extract batch numbers and find the maximum
                batch_numbers = []
                for directory_name in batch_dirs:
                    try:
                        num_str = directory_name.split('_')[1].split('-')[0]
                        num = int(num_str)
                        batch_numbers.append(num)
                    except (IndexError, ValueError):
                        continue
                
                if len(batch_numbers) > 0:
                    end_idx = max(batch_numbers)
                else:
                    end_idx = start_idx
            
            indices = range(start_idx, end_idx + 1)
        
        # If return_as is 'path', return list of paths
        if return_as.lower() in ['path', 'string']:
            paths = []
            for idx in indices:
                try:
                    path = DataProcessing.load_single_synthetic_data(
                        notebook_dir=notebook_dir, 
                        batch_idx=idx,
                        sep=sep,
                        data_type=data_type,
                        return_as='path'
                    )
                    paths.append(path)
                    print(f"✓ Found batch {idx}: {path}")
                except FileNotFoundError:
                    print(f"⚠ Warning: Batch {idx} not found, skipping...")
                    continue
                except Exception as e:
                    print(f"⚠ Error locating batch {idx}: {e}")
                    continue
            
            if len(paths) == 0:
                raise ValueError("No batch paths were found")
            
            print(f"\nFound {len(paths)} batch file(s)")
            return paths
        
        # Default behavior: load as dataframes
        dfs = []
        for idx in indices:
            try:
                df = DataProcessing.load_single_synthetic_data(
                    notebook_dir=notebook_dir, 
                    batch_idx=idx,
                    sep=sep,
                    data_type=data_type,
                    return_as='dataframe'
                )
                dfs.append(df)
                print(f"✓ Loaded batch {idx}")
            except FileNotFoundError:
                print(f"⚠ Warning: Batch {idx} not found, skipping...")
                continue
            except Exception as e:
                print(f"⚠ Error loading batch {idx}: {e}")
                continue
        
        if len(dfs) == 0:
            raise ValueError("No batches were successfully loaded")
        
        # Concatenate all dataframes
        combined_df = DataProcessing.concat_dfs(dfs)
        print(f"\nSuccessfully loaded and combined {len(dfs)} batches")
        print(f"Total rows: {len(combined_df)}")
        
        return combined_df

    def match_text_label_to_int(df: pd.DataFrame, text_label_col_name: str, 
                            target_label: str = 'PREDICTION',
                            int_label_col_name: str = 'Binary Label') -> pd.DataFrame:
        """
        Convert text labels to integer labels (binary classification).
        
        Parameters:
        -----------
        df : pd.DataFrame
            DataFrame containing the text labels
        text_label_col_name : str
            Name of the column containing text labels (e.g., 'maya_label')
        target_label : str
            The label to encode as 1. Default is 'PREDICTION'.
            Common values: 'PREDICTION', 'NON-PREDICTION', 'OBSERVATION', 'NON-OBSERVATION'
        int_label_col_name : str
            Name for the new integer label column. Default is 'Binary Label'.
        
        Returns:
        --------
        pd.DataFrame
            DataFrame with added integer label column where:
            - 1 = matches the target_label
            - 0 = does not match the target_label
        """
        # Create a copy to avoid modifying the original
        result_df = df.copy()
        
        # Create filter for matching labels
        filt_match = (result_df[text_label_col_name] == target_label)
        
        # Create two dataframes: matches and non-matches
        match_df = result_df[filt_match].copy()
        match_df[int_label_col_name] = 1
        
        non_match_df = result_df[~filt_match].copy()
        non_match_df[int_label_col_name] = 0
        
        # Concatenate and return
        numerical_labels_df = DataProcessing.concat_dfs([match_df, non_match_df])
        
        return numerical_labels_df
    
    def apply_resampling_full_dimensions(
            df: pd.DataFrame,
            embedding_col: str,
            label_col: str,
            random_state,
            method: str,
            sampling_strategy: str = 'auto',
            save_path: str = None,
            save_prefix: str = None,
            save_file_type: str = None) -> pd.DataFrame:
        """Apply resampling to full 96-dimensional embeddings."""
        X_full = np.stack(df[embedding_col].values)
        y = df[label_col].values

        print("\n" + "="*40)
        print(f"RESAMPLE: {embedding_col} | {label_col}")
        print("="*40)
        print(f"X_full Shape: {X_full.shape}")
        print(f"\nX_full Preview:\n{X_full}\n")
        print(f"y Shape: {y.shape}")
        print(f"\ny Preview:\n{y}\n")


        if method == 'oversample':
            resampler = RandomOverSampler(sampling_strategy=sampling_strategy, random_state=random_state)
        elif method == 'undersample':
            resampler = RandomUnderSampler(sampling_strategy=sampling_strategy, random_state=random_state)
        else:
            raise ValueError(f"Unknown method: {method}. Use 'oversample' or 'undersample'.")
        
        X_resampled, y_resampled = resampler.fit_resample(X_full, y)
        
        # Get indices from resampling
        indices = resampler.sample_indices_ if hasattr(resampler, 'sample_indices_') else None
        
        # Create resampled dataframe by duplicating original rows
        if indices is not None:
            resampled_df = df.iloc[indices].copy().reset_index(drop=True)
            # Update the embedding and label columns with resampled values
            resampled_df[embedding_col] = list(X_resampled)
            resampled_df[label_col] = y_resampled
        else:
            # Fallback: just include embedding and label columns
            resampled_df = pd.DataFrame({
                embedding_col: list(X_resampled),
                label_col: y_resampled
            })
        
        if save_path:
            DataProcessing.save_to_file(
                data=resampled_df,
                path=save_path,
                prefix=save_prefix,
                save_file_type=save_file_type
            )
        
        return resampled_df
    
    # def extract_features_for_visualization(df: pd.DataFrame, 
    #                                    embedding_col: str, 
    #                                    label_col: str) -> pd.DataFrame:
    #     """Extract first 2 dimensions from embeddings for visualization."""
    #     print("\n" + "="*40)
    #     print(f"EXTRACT FEATURES FOR VISUALIZING: {embedding_col} | {label_col}")
    #     print("="*40)
    #     print(f"df Shape: {df.shape}")
    #     print(f"\ndf Preview:\n{df.head(7)}\n")

    #     embeddings_array = np.stack(df[embedding_col].values)
    #     features_df = pd.DataFrame({
    #         'Feature_1': embeddings_array[:, 0],
    #         'Feature_2': embeddings_array[:, 1],
    #         'Label': df[label_col]
    #     })
    #     return features_df

    def parse_llm_json_response(text: str):
        """
        Robustly parse a JSON dictionary from a raw LLM response string.

        Parameters
        ----------
        text : str
            Raw string output from an LLM response.

        Returns
        -------
        dict or None
            Parsed dictionary if successful, or None if parsing fails.

        Notes
        -----
        Handles common LLM response formatting issues including:
        - Markdown code blocks (```json ... ``` or ``` ... ```)
        - Conversational filler text before or after the JSON
        - Single-quoted Python dictionary strings via ast.literal_eval fallback
        """
        text = str(text).strip()

        match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', text, re.DOTALL)
        if match:
            json_str = match.group(1)
        else:
            match = re.search(r'(\{.*\})', text, re.DOTALL)
            json_str = match.group(1) if match else text

        try:
            return json.loads(json_str)
        except json.JSONDecodeError:
            try:
                return ast.literal_eval(json_str)
            except:
                return None
    
    def expand_pipe_separated_rows(df: pd.DataFrame, col_names: list) -> pd.DataFrame:
        """
        Expand pipe-separated values in property columns into separate rows.

        Each pipe-separated value gets its own row, with all other columns duplicated from the original row. A Token_Index column is added to
        track the position of each token within its original sentence row, allowing traceability back to the original Input_Index.

        Parameters
        ----------
        df : pd.DataFrame
            DataFrame containing pipe-separated property columns.
        col_names : list of str
            Column names to expand (e.g., ['No Property', 'Source', 'Target', 'Date', 'Outcome']).

        Returns
        -------
        pd.DataFrame
            Expanded DataFrame where each pipe-separated value has its own row
            with a Token_Index column added.

        Examples
        --------
        Input:
            Input_Index | Sentence  | Source   | Outcome
            0           | Apple...  | JPMorgan | stock price|remain stable

        Output:
            Input_Index | Token_Index | Sentence  | Source   | Outcome
            0           | 0           | Apple...  | JPMorgan | stock price
            0           | 1           | Apple...  | JPMorgan | remain stable
        """
        expanded_df = df.copy()

        # Split each property column on pipe into a list
        for col_name in col_names:
            if col_name in expanded_df.columns:
                expanded_df[col_name] = expanded_df[col_name].astype(str).str.split('|')

        # Add Token_Index based on the length of the first property column list
        first_col = next((c for c in col_names if c in expanded_df.columns), None)
        if first_col:
            expanded_df['Token_Index'] = expanded_df[first_col].apply(
                lambda x: list(range(len(x))) if isinstance(x, list) else [0]
            )

        # Explode each property column
        for col_name in col_names:
            if col_name in expanded_df.columns:
                expanded_df = expanded_df.explode(col_name).reset_index(drop=True)

        # Explode Token_Index to match
        expanded_df = expanded_df.explode('Token_Index').reset_index(drop=True)

        # Strip whitespace
        for col_name in col_names:
            if col_name in expanded_df.columns:
                expanded_df[col_name] = expanded_df[col_name].str.strip()

        # Reorder so Input_Index and Token_Index are first
        priority_cols = ['Input_Index', 'Token_Index']
        remaining_cols = [col for col in expanded_df.columns if col not in priority_cols]
        expanded_df = expanded_df[priority_cols + remaining_cols]

        return expanded_df
    # ==============================================================
    # HELPER: Standardize columns
    # ==============================================================
    def standardize_columns(df, text_col, label_col, label_map=None):
        """
        Standardize column names and ordering for any dataset.

        Parameters
        ----------
        df : pd.DataFrame
            Raw loaded dataset
        text_col : str
            Name of the sentence/text column to map to 'Base Sentence'
        label_col : str
            Name of the label column to map to 'Ground Truth'
        label_map : dict or None
            Optional mapping to convert raw label values to integers
            before renaming to 'Ground Truth'. e.g. {'prediction': 1, 'not-prediction': 0}

        Returns
        -------
        pd.DataFrame
            Dataset with standardized column names and ordering
        """
        if label_map is not None:
            if label_col not in df.columns:
                raise ValueError(f"Expected label column '{label_col}' not found for label_map application.")
            df[label_col] = df[label_col].str.lower().map(label_map)

        rename_map = {}
        if text_col != 'Base Sentence' and text_col in df.columns:
            rename_map[text_col] = 'Base Sentence'
        if label_col != 'Ground Truth' and label_col in df.columns:
            rename_map[label_col] = 'Ground Truth'
        if rename_map:
            df = df.rename(columns=rename_map)

        if 'Base Sentence' not in df.columns:
            raise ValueError(f"Expected text column '{text_col}' not found in dataset.")
        if 'Ground Truth' not in df.columns:
            raise ValueError(f"Expected label column '{label_col}' not found in dataset.")

        priority_cols = ['Base Sentence', 'Ground Truth']
        remaining_cols = []
        for col in df.columns:
            if col not in priority_cols:
                remaining_cols.append(col)

        df = df[priority_cols + remaining_cols]
        return df

    def load_synthetic_predictions_dataset(script_dir, sep=','):
        print("\n" + "="*60)
        print("LOAD PREDICTIONS DATASET")
        print("="*60)
        predictions_df = DataProcessing.load_multiple_batches(
            script_dir,
            sep=sep,
            data_type='prediction'
        )
        print(f"Shape: {predictions_df.shape}")
        print(f"Columns: {list(predictions_df.columns)}")
        skip = ['Base Sentence']
        for col in predictions_df.columns:
            if col not in skip:
                print(f"\n=== {col} value counts ===")
                print(predictions_df[col].value_counts())
        predictions_df = DataProcessing.standardize_columns(
            df=predictions_df,
            text_col='Base Sentence',
            label_col='Sentence Label'
        )
        predictions_df['Dataset Name'] = 'synthetic_predictions'
        print(f"\nPreview:\n{predictions_df.head(3)}\n")
        return predictions_df

    def load_synthetic_non_predictions_dataset(script_dir, sep=','):
        print("\n" + "="*60)
        print("LOAD NON-PREDICTIONS DATASET")
        print("="*60)
        non_predictions_df = DataProcessing.load_multiple_batches(
            script_dir,
            sep=sep,
            data_type='observation'
        )
        print(f"Shape: {non_predictions_df.shape}")
        print(f"Columns: {list(non_predictions_df.columns)}")
        skip = ['Base Sentence']
        for col in non_predictions_df.columns:
            if col not in skip:
                print(f"\n=== {col} value counts ===")
                print(non_predictions_df[col].value_counts())
        non_predictions_df = DataProcessing.standardize_columns(
            df=non_predictions_df,
            text_col='Base Sentence',
            label_col='Sentence Label'
        )
        non_predictions_df['Dataset Name'] = 'synthetic_non_predictions'
        print(f"\nPreview:\n{non_predictions_df.head(3)}\n")
        return non_predictions_df

    def load_synthetic_dataset(script_dir, sep=',', predictions_only: bool = False, visualize: bool = True, **kwargs):
        print("\n" + "="*60)
        print("LOAD SYNTHETIC DATASET")
        print("="*60)

        predictions_df     = DataProcessing.load_synthetic_predictions_dataset(script_dir, sep=sep)
        non_predictions_df = DataProcessing.load_synthetic_non_predictions_dataset(script_dir, sep=sep)

        if predictions_only:
            print(f"Filtered shape (predictions only): {predictions_df.shape}")
            return predictions_df

        synthetic_df = DataProcessing.concat_dfs([predictions_df, non_predictions_df])
        print(f"Combined shape: {synthetic_df.shape}")

        if visualize:
            from data_visualizing import DataVisualizing
            DataVisualizing.plot_class_distribution(
                df=synthetic_df,
                label_col='Ground Truth',
                title='Synthetic Dataset Class Distribution'
            )

        print("\nGround Truth distribution:")
        print(synthetic_df['Ground Truth'].value_counts())

        synthetic_df['Dataset Name'] = 'synthetic'
        print(f"\nPreview:\n{synthetic_df.head(7)}\n")

        return synthetic_df

    def load_financial_phrasebank_dataset(script_dir, sep=',', encoding='utf-8', predictions_only: bool = False, visualize: bool = True, **kwargs):
        print("\n" + "="*60)
        print("LOAD FINANCIAL PHRASEBANK DATASET")
        print("="*60)
        
        base_data_path = DataProcessing.load_base_data_path(script_dir)
        fpb_path = os.path.join(
            base_data_path,
            # 'financial_phrase_bank/annotators/fpb-maya-binary-imbalanced-96d-v1.csv'
            'financial_phrase_bank/annotators/fpb-maya-with_sentiment-final.xlsx'
        )
        
        print(f"Loading from: {fpb_path}")
        fpb_df = DataProcessing.load_from_file(fpb_path, 'xlsx', sep=None, encoding=None)
        print(f"Loaded shape: {fpb_df.shape}")
        
        original_len = len(fpb_df)
        fpb_df.dropna(inplace=True)
        dropped_count = original_len - len(fpb_df)
        if dropped_count > 0:
            print(f"✓ Dropped {dropped_count} rows without labels")
        
        print("\nConverting text labels to binary...")
        fpb_df = DataProcessing.match_text_label_to_int(
            fpb_df,
            text_label_col_name='maya_label',
            target_label='PREDICTION'
        )
        fpb_df = DataProcessing.standardize_columns(
            df=fpb_df,
            text_col='statement',
            label_col='Binary Label'
        )
        
        if predictions_only:
            fpb_df = fpb_df[fpb_df['Ground Truth'] == 1]
            print(f"Filtered shape (predictions only): {fpb_df.shape}")
        elif visualize:
            from data_visualizing import DataVisualizing
            DataVisualizing.plot_class_distribution(
                df=fpb_df,
                label_col='Ground Truth',
                title='FPB Class Distribution'
            )
        
        print(f"Columns: {list(fpb_df.columns)}")
        print(f"\nGround Truth distribution:")
        print(fpb_df['Ground Truth'].value_counts())
        
        fpb_df['Dataset Name'] = 'financial_phrasebank'
        print(f"\nPreview:\n{fpb_df.head(3)}\n")
        
        return fpb_df

    def load_chronicle2050_dataset(script_dir, sep=',', encoding='utf-8', predictions_only: bool = False, visualize: bool = True, **kwargs):
        print("\n" + "="*60)
        print("LOAD CHRONICLE2050 DATASET")
        print("="*60)
        
        base_data_path = DataProcessing.load_base_data_path(script_dir)
        chronicle2050_path = os.path.join(
            base_data_path,
            'chronicle2050',
            'annotators',
            'chronicle2050-shawnick-binary-v7.csv'
        )
        print(f"Loading from: {chronicle2050_path}")
        
        chronicle2050_df = DataProcessing.load_from_file(
            chronicle2050_path,
            'csv',
            sep=sep,
            encoding=encoding
        )
        print(f"Loaded shape: {chronicle2050_df.shape}")
        
        original_len = len(chronicle2050_df)
        chronicle2050_df.dropna(subset=['shawnick_labels'], inplace=True)
        dropped_count = original_len - len(chronicle2050_df)
        if dropped_count > 0:
            print(f"✓ Dropped {dropped_count} rows without labels")
        
        chronicle2050_df = DataProcessing.standardize_columns(
            df=chronicle2050_df,
            text_col='sentence',
            label_col='shawnick_labels',
            label_map={'prediction': 1, 'not-prediction': 0}
        )
        
        if predictions_only:
            chronicle2050_df = chronicle2050_df[chronicle2050_df['Ground Truth'] == 1]
            print(f"Filtered shape (predictions only): {chronicle2050_df.shape}")
        elif visualize:
            from data_visualizing import DataVisualizing
            DataVisualizing.plot_class_distribution(
                df=chronicle2050_df,
                label_col='Ground Truth',
                title='Chronicle2050 Class Distribution'
            )
        
        print(f"Columns: {list(chronicle2050_df.columns)}")
        print("\nGround Truth distribution:")
        print(chronicle2050_df['Ground Truth'].value_counts())
        
        chronicle2050_df['Dataset Name'] = 'chronicle2050'
        print(f"\nPreview:\n{chronicle2050_df.head(7)}\n")
        
        return chronicle2050_df

    def load_timebank_dataset(script_dir, sep=',', predictions_only: bool = False, visualize: bool = True, **kwargs):
        print("\n" + "="*60)
        print("LOAD TIMEBANK DATASET")
        print("="*60)
        
        base_data_path = DataProcessing.load_base_data_path(script_dir)
        tb_path = os.path.join(base_data_path, "timebank_1_2", "annotators")
        print(f"Loading from: {tb_path}")
        
        dfs = []
        for filename in os.listdir(tb_path):
            if not filename.endswith(".csv"):
                continue
            filepath = os.path.join(tb_path, filename)
            print(f"Loading: {filename}")
            df = DataProcessing.load_from_file(filepath, file_type="csv", sep=sep)
            dfs.append(df)
        
        if not dfs:
            print("⚠️ No TimeBank CSVs found.")
            return pd.DataFrame()
        
        tb_df = DataProcessing.concat_dfs(dfs)
        print(f"Loaded shape (all rows): {tb_df.shape}")
        
        if 'Label' not in tb_df.columns:
            raise ValueError("Expected 'Label' column in TimeBank dataset")
        
        tb_df = DataProcessing.standardize_columns(
            df=tb_df,
            text_col='Sentence',
            label_col='Label'
        )
        
        if predictions_only:
            tb_df = tb_df[tb_df['Ground Truth'] == 1]
            print(f"Filtered shape (predictions only): {tb_df.shape}")
        elif visualize:
            from data_visualizing import DataVisualizing
            DataVisualizing.plot_class_distribution(
                df=tb_df,
                label_col='Ground Truth',
                title='TimeBank Class Distribution'
            )
        
        print("\nGround Truth distribution:")
        print(tb_df['Ground Truth'].value_counts())
        
        tb_df['Dataset Name'] = 'timebank'
        print(f"\nPreview:\n{tb_df.head(7)}\n")
        print(f"\nPreview:\n{tb_df.tail(7)}\n")
        
        return tb_df
    
    def load_yt_dataset(script_dir, sep=',', predictions_only: bool = False, visualize: bool = True, **kwargs):
        print("\n" + "="*60)
        print("LOAD YT DATASET")
        print("="*60)
        
        base_data_path = DataProcessing.load_base_data_path(script_dir)
        yt_path = os.path.join(base_data_path, "yt", "annotators", "sports")
        print(f"Loading from: {yt_path}")
        
        dfs = []
        for filename in os.listdir(yt_path):
            if not filename.endswith(".csv"):
                continue
            filepath = os.path.join(yt_path, filename)
            print(f"Loading: {filename}")
            df = DataProcessing.load_from_file(filepath, file_type="csv", sep=sep)
            dfs.append(df)
        
        if not dfs:
            print("⚠️ No YT CSVs found.")
            return pd.DataFrame()
        
        yt_df = DataProcessing.concat_dfs(dfs)
        print(f"Loaded shape (all rows): {yt_df.shape}")
        
        if 'Human Annotation' not in yt_df.columns:
            raise ValueError("Expected 'Human Annotation' column in YT dataset")
        
        yt_df = DataProcessing.standardize_columns(
            df=yt_df,
            text_col='Base Sentence',
            label_col='Human Annotation'
        )
        
        if predictions_only:
            yt_df = yt_df[yt_df['Ground Truth'] == 1]
            print(f"Filtered shape (predictions only): {yt_df.shape}")
        elif visualize:
            from data_visualizing import DataVisualizing
            DataVisualizing.plot_class_distribution(
                df=yt_df,
                label_col='Ground Truth',
                title='YT Class Distribution'
            )
        
        print("\nGround Truth distribution:")
        print(yt_df['Ground Truth'].value_counts())
        
        yt_df['Dataset Name'] = 'yt'
        print(f"\nPreview:\n{yt_df.head(7)}\n")
        
        return yt_df

    def load_news_api_dataset(script_dir, sep=',', predictions_only: bool = False, visualize: bool = True, **kwargs):
        print("\n" + "="*60)
        print("LOAD NEWS API DATASET")
        print("="*60)
        
        base_data_path = DataProcessing.load_base_data_path(script_dir)
        news_api_path = os.path.join(base_data_path, "news_api", "annotators", "all_domains")

        ### Get specific file
        # news_api_path = os.path.join(base_data_path, "news_api", "annotators", "news_articles_election_vote_polling_legislation_expected_likely_projected_forecast_2026-01-01_to_2026-04-26_predictions-v7_policy_1_human_annotation.csv")
        # print(f"Loading from: {news_api_path}")
        # df = DataProcessing.load_from_file(news_api_path, file_type="csv", sep=sep)
        
        dfs = []
        for filename in os.listdir(news_api_path):
            if not filename.endswith(".csv"):
                continue
            filepath = os.path.join(news_api_path, filename)
            print(f"Loading: {filename}")
            df = DataProcessing.load_from_file(filepath, file_type="csv", sep=sep)
            print(f"   Loaded shape (all rows): {df.shape}")
            dfs.append(df)
        
        if not dfs:
            print("⚠️ No NewsAPI CSVs found.")
            return pd.DataFrame()
        
        news_api_df = DataProcessing.concat_dfs(dfs)
        print(f"Loaded shape (all rows): {news_api_df.shape}")
        
        if 'Human Annotation' not in news_api_df.columns:
            raise ValueError("Expected 'Human Annotation' column in NewsAPI dataset")
        
        news_api_df = DataProcessing.standardize_columns(
            df=news_api_df,
            text_col='Base Sentence',
            label_col='Sentence Label'
        )
        
        if predictions_only:
            news_api_df = news_api_df[news_api_df['Ground Truth'] == 1]
            print(f"Filtered shape (predictions only): {news_api_df.shape}")
        elif visualize:
            from data_visualizing import DataVisualizing
            DataVisualizing.plot_class_distribution(
                df=news_api_df,
                label_col='Ground Truth',
                title='NewsAPI Class Distribution'
            )
        
        print("\nGround Truth distribution:")
        print(news_api_df['Ground Truth'].value_counts())
        
        news_api_df['Dataset Name'] = 'news_api'
        print(f"\nPreview:\n{news_api_df.head(7)}\n")
        
        return news_api_df

    def load_mf_climate_dataset(script_dir, sep=',', encoding='latin', predictions_only: bool = False, visualize: bool = True, **kwargs):
        print("\n" + "="*60)
        print("LOAD MF CLIMATE DATASET")
        print("="*60)
        
        base_data_path = DataProcessing.load_base_data_path(script_dir)
        # mf_climate_path = os.path.join(
        #     base_data_path,
        #     'multimodal_forecast/climate_forecasts/climate_2014_2023_final.csv'
        # )
        
        mf_climate_path = os.path.join(
            base_data_path,
            'multimodal_forecast/climate_forecasts/climate_forecasts-v5.csv'
        ) # dropped bad sentences -> notebooks -> explore_climate_2014_2023-pre_label.ipynb

        print(f"Loading from: {mf_climate_path}")
        mf_climate_df = DataProcessing.load_from_file(mf_climate_path, 'csv', sep=sep, encoding=encoding)
        mf_climate_df['Ground Truth'] = 1
        print(f"Loaded shape: {mf_climate_df.shape}")
        
        original_len = len(mf_climate_df)
        mf_climate_df.dropna(inplace=True)
        dropped_count = original_len - len(mf_climate_df)
        if dropped_count > 0:
            print(f"✓ Dropped {dropped_count} rows without labels")
        

        mf_climate_df = DataProcessing.standardize_columns(
            df=mf_climate_df,
            text_col='text',
            label_col='Ground Truth'
        )
        
        if predictions_only:
            mf_climate_df = mf_climate_df[mf_climate_df['Ground Truth'] == 1]
            print(f"Filtered shape (predictions only): {mf_climate_df.shape}")
        elif visualize:
            from data_visualizing import DataVisualizing
            DataVisualizing.plot_class_distribution(
                df=mf_climate_df,
                label_col='Ground Truth',
                title='MF Climate Class Distribution'
            )
        
        print(f"Columns: {list(mf_climate_df.columns)}")
        print(f"\nGround Truth distribution:")
        print(mf_climate_df['Ground Truth'].value_counts())
        
        mf_climate_df['Dataset Name'] = 'mf_climate'
        print(f"\nPreview:\n{mf_climate_df.head(3)}\n")
        
        return mf_climate_df
        
    def load_clients_rivals_rouges_dataset(
        script_dir,
        sep="\t",
        encoding="utf-8",
        predictions_only: bool = True,
        visualize: bool = False,
        **kwargs
    ):
        print("\n" + "="*60)
        print("LOAD CLIENTS, RIVALS, ROUGES DATASET")
        print("="*60)

        base_data_path = DataProcessing.load_base_data_path(script_dir)
        dataset_path = os.path.join(
            base_data_path,
            "clients_rivals_rogues/clients_rivals_rogues.tsv"
        )

        print(f"Loading from: {dataset_path}")
        df = DataProcessing.load_from_file(
            dataset_path,
            file_type="csv",
            sep=sep,
            encoding=encoding,
            **kwargs
        )

        print(f"Loaded shape: {df.shape}")
        print(f"Columns: {list(df.columns)}")

        # ---- Safety check BEFORE standardization ----
        required_cols = {
            "Theory",
            "Theory prediction",
            "Empirical support",
            "Human Annotation"
        }
        missing = required_cols - set(df.columns)
        if missing:
            raise ValueError(f"Missing required columns: {missing}")

        # ---- Standardize ----
        df = DataProcessing.standardize_columns(
            df=df,
            text_col="Theory prediction",
            label_col="Human Annotation"
        )

        # ---- Optional filtering / visualization ----
        if predictions_only:
            df = df[df["Ground Truth"] == 1]
            print(f"Filtered shape (predictions only): {df.shape}")

        elif visualize:
            from data_visualizing import DataVisualizing
            DataVisualizing.plot_class_distribution(
                df=df,
                label_col="Ground Truth",
                title="Clients, Rivals, Rouges Class Distribution"
            )

        # ---- Rename remaining metadata ----
        df = df.rename(columns={
            "Theory": "Source Theory",
            "Empirical support": "Empirical Support"
        })

        print("\nGround Truth distribution:")
        print(df["Ground Truth"].value_counts())

        df["Dataset Name"] = "clients_rivals_rouges"

        print(f"\nPreview:\n{df.head(3)}\n")
        return df

    def load_forecast_bench_dataset(script_dir, sep=',', encoding='utf-8', predictions_only: bool = False, visualize: bool = True):
        print("\n" + "="*60)
        print("LOAD FORECAST BENCH (QUESTIONS) DATASET")
        print("="*60)
        
        base_data_path = DataProcessing.load_base_data_path(script_dir)
        forecast_bench_path = os.path.join(
            base_data_path,
            'forecast_bench',
            'forecast_bench_questions.csv'
        )
        print(f"Loading from: {forecast_bench_path}")
        
        forecast_bench_df = DataProcessing.load_from_file(
            forecast_bench_path,
            'csv'
        )
        print(f"Loaded shape: {forecast_bench_df.shape}")

                # ---- Standardize ----
        df = DataProcessing.standardize_columns(
            df=forecast_bench_df,
            text_col="question",
            label_col="Ground Truth"
        )
        df['Dataset Name'] = 'forecast_bench'

        return df
    
    # not yet
    def load_smart_hospitals_dataset(
        script_dir,
        sep="\t",
        encoding="utf-8",
        predictions_only: bool = True,
        visualize: bool = False,
        **kwargs
    ):
        print("\n" + "="*60)
        print("LOAD SMART HOSPITALS DATASET")
        print("="*60)

        base_data_path = DataProcessing.load_base_data_path(script_dir)
        dataset_path = os.path.join(
            base_data_path,
            "smart_hospitals/smart_hospitals.tsv"
        )

        print(f"Loading from: {dataset_path}")
        df = DataProcessing.load_from_file(
            dataset_path,
            file_type="csv",
            sep=sep,
            encoding=encoding,
            **kwargs
        )

        print(f"Loaded shape: {df.shape}")
        print(f"Columns: {list(df.columns)}")

        # ---- Safety check BEFORE standardization ----
        required_cols = {
            "Base Sentence",
            "Human Annotation",
            "Domain",
            "Dimensionality",
            "Time Horizon"
        }
        missing = required_cols - set(df.columns)
        if missing:
            raise ValueError(f"Missing required columns: {missing}")

        # ---- Standardize ----
        df = DataProcessing.standardize_columns(
            df=df,
            text_col="Base Sentence",
            label_col="Human Annotation"
        )

        if predictions_only:
            df = df[df["Ground Truth"] == 1]
            print(f"Filtered shape (predictions only): {df.shape}")

        elif visualize:
            from data_visualizing import DataVisualizing
            DataVisualizing.plot_class_distribution(
                df=df,
                label_col="Ground Truth",
                title="Smart Hospitals Class Distribution"
            )

        df["Dataset Name"] = "smart_hospitals"

        print("\nGround Truth distribution:")
        print(df["Ground Truth"].value_counts())

        print(f"\nPreview:\n{df.head(3)}\n")

        return df