src/scripts/combine_dataset.py

#!/usr/bin/env python3
"""

Script to combine multiple CSV datasets into a single dataset file.



This script takes multiple input CSV files and concatenates them into a single

output CSV file. It handles deduplication, ID reassignment, and validates

data consistency across input files.

"""

import argparse
import csv
import logging
from pathlib import Path

import pandas as pd

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def validate_csv_structure(

    file_path: Path, expected_columns: set[str] | None = None

) -> set[str]:
    """

    Validate the structure of a CSV file and return its columns.



    Args:

        file_path: Path to the CSV file to validate

        expected_columns: Optional set of expected column names for validation



    Returns:

        set[str]: Set of column names found in the file



    Raises:

        ValueError: If file validation fails

    """
    if not file_path.exists():
        raise ValueError(f"Input file does not exist: {file_path}")

    if not file_path.suffix.lower() == ".csv":
        raise ValueError(f"Input file is not a CSV file: {file_path}")

    try:
        # Read just the header to validate structure
        df_sample = pd.read_csv(file_path, nrows=0)
        columns = set(df_sample.columns)

        if expected_columns and columns != expected_columns:
            raise ValueError(
                f"Column mismatch in {file_path}. "
                f"Expected: {expected_columns}, Found: {columns}"
            )

        logger.info(f"Validated {file_path}: {len(columns)} columns found")
        return columns

    except Exception as e:
        raise ValueError(f"Failed to read CSV file {file_path}: {e}") from e


def load_and_combine_datasets(

    input_paths: list[Path],

    remove_duplicates: bool = True,

    duplicate_columns: list[str] | None = None,

) -> pd.DataFrame:
    """

    Load and combine multiple CSV datasets into a single DataFrame.



    Args:

        input_paths: List of paths to input CSV files

        remove_duplicates: Whether to remove duplicate rows

        duplicate_columns: Columns to use for duplicate detection (if None, uses all columns)



    Returns:

        pd.DataFrame: Combined dataset with reassigned IDs



    Raises:

        ValueError: If datasets cannot be combined

    """
    if not input_paths:
        raise ValueError("No input paths provided")

    dataframes = []
    expected_columns = None

    for i, path in enumerate(input_paths):
        logger.info(f"Loading dataset {i + 1}/{len(input_paths)}: {path}")

        # Validate file structure
        columns = validate_csv_structure(path, expected_columns)
        if expected_columns is None:
            expected_columns = columns

        # Load the dataset
        try:
            df = pd.read_csv(path)
            logger.info(f"Loaded {len(df)} rows from {path}")
            dataframes.append(df)
        except Exception as e:
            raise ValueError(f"Failed to load dataset from {path}: {e}") from e

    # Combine all dataframes
    logger.info("Combining datasets...")
    combined_df = pd.concat(dataframes, ignore_index=True)
    logger.info(f"Combined dataset has {len(combined_df)} total rows")

    # Remove duplicates if requested
    if remove_duplicates:
        initial_count = len(combined_df)
        if duplicate_columns:
            combined_df = combined_df.drop_duplicates(
                subset=duplicate_columns, keep="first"
            )
        else:
            combined_df = combined_df.drop_duplicates(keep="first")

        duplicates_removed = initial_count - len(combined_df)
        if duplicates_removed > 0:
            logger.info(f"Removed {duplicates_removed} duplicate rows")

    # Reassign IDs if 'id' column exists
    if "id" in combined_df.columns:
        logger.info("Reassigning IDs sequentially...")
        combined_df["id"] = range(1, len(combined_df) + 1)

    return combined_df


def save_combined_dataset(df: pd.DataFrame, output_path: Path) -> None:
    """

    Save the combined dataset to a CSV file.



    Args:

        df: DataFrame to save

        output_path: Path where to save the combined dataset



    Raises:

        ValueError: If save operation fails

    """
    try:
        # Create output directory if it doesn't exist
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # Save to CSV
        df.to_csv(output_path, index=False, quoting=csv.QUOTE_MINIMAL)
        logger.info(f"Successfully saved combined dataset to {output_path}")
        logger.info(
            f"Final dataset contains {len(df)} rows and {len(df.columns)} columns"
        )

        # Log column names for verification
        logger.info(f"Columns: {', '.join(df.columns)}")

    except Exception as e:
        raise ValueError(
            f"Failed to save combined dataset to {output_path}: {e}"
        ) from e


def main() -> None:
    """

    Main function to combine multiple datasets.

    """
    parser = argparse.ArgumentParser(
        description="Combine multiple CSV datasets into a single dataset file",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""

Examples:

  # Combine two datasets

  python combine_dataset.py --input_paths data/train1.csv data/train2.csv --output_path data/combined_train.csv



  # Combine multiple datasets without removing duplicates

  python combine_dataset.py --input_paths data/*.csv --output_path data/all_data.csv --no-remove-duplicates



  # Combine datasets with custom duplicate detection columns

  python combine_dataset.py --input_paths data/train*.csv --output_path data/combined.csv --duplicate_columns problem_description correct_answer

        """,
    )

    parser.add_argument(
        "--input_paths",
        type=str,
        nargs="+",
        required=True,
        help="Paths to input CSV files to combine",
    )

    parser.add_argument(
        "--output_path",
        type=str,
        required=True,
        help="Path for the output combined CSV file",
    )

    parser.add_argument(
        "--no-remove-duplicates",
        action="store_true",
        help="Do not remove duplicate rows from the combined dataset",
    )

    parser.add_argument(
        "--duplicate_columns",
        type=str,
        nargs="*",
        help="Columns to use for duplicate detection (if not specified, uses all columns)",
    )

    parser.add_argument(
        "--log-level",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        default="INFO",
        help="Set the logging level (default: INFO)",
    )

    args = parser.parse_args()

    # Set logging level
    logging.getLogger().setLevel(getattr(logging, args.log_level))

    # Convert string paths to Path objects
    input_paths = [Path(path) for path in args.input_paths]
    output_path = Path(args.output_path)

    try:
        logger.info(f"Starting dataset combination with {len(input_paths)} input files")
        logger.info(f"Input files: {[str(p) for p in input_paths]}")
        logger.info(f"Output file: {output_path}")

        # Load and combine datasets
        combined_df = load_and_combine_datasets(
            input_paths=input_paths,
            remove_duplicates=not args.no_remove_duplicates,
            duplicate_columns=args.duplicate_columns,
        )

        # Save combined dataset
        save_combined_dataset(combined_df, output_path)

        logger.info("Dataset combination completed successfully")

    except Exception as e:
        logger.error(f"Dataset combination failed: {e}")
        raise


if __name__ == "__main__":
    main()