"""
Data Splitter Module
====================

This module provides functionality for splitting data into train, validation, and test sets.
"""

import logging
from pathlib import Path
from typing import Tuple, Optional
import pandas as pd
from sklearn.model_selection import train_test_split

logger = logging.getLogger(__name__)


class DataSplitter:
    """
    Splits data into training, validation, and test sets with stratification.
    """
    
    LABEL_MAP = {0: "negative", 2: "neutral", 4: "positive"}
    
    def __init__(
        self,
        test_size: float = 0.2,
        val_size: float = 0.1,
        random_state: int = 42,
        binary_only: bool = False
    ):
        """
        Initialize the data splitter.
        
        Args:
            test_size: Proportion of data for testing (0.0-1.0)
            val_size: Proportion of training data for validation (0.0-1.0)
            random_state: Random seed for reproducibility
            binary_only: If True, only keep negative (0) and positive (4) labels
        """
        self.test_size = test_size
        self.val_size = val_size
        self.random_state = random_state
        self.binary_only = binary_only
        
    def load_data(self, data_path: Path) -> pd.DataFrame:
        """
        Load Sentiment140 dataset from CSV.
        
        Args:
            data_path: Path to the CSV file
            
        Returns:
            DataFrame with the data
        """
        logger.info(f"Loading data from {data_path}")
        
        colnames = ["target", "ids", "date", "flag", "user", "text"]
        df = pd.read_csv(data_path, encoding="latin-1", names=colnames)
        
        # Remove rows with missing text
        df = df[~df["text"].isna()].reset_index(drop=True)
        df["target"] = df["target"].astype(int)
        df["label_name"] = df["target"].map(self.LABEL_MAP)
        
        # Filter to binary classification if requested
        if self.binary_only:
            logger.info("Filtering to binary classification (0=negative, 4=positive)")
            df = df[df["target"].isin([0, 4])].reset_index(drop=True)
        
        logger.info(f"Loaded {len(df)} samples")
        logger.info(f"Class distribution:\n{df['target'].value_counts()}")
        
        return df
    
    def split_data(
        self,
        df: pd.DataFrame
    ) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        """
        Split data into train, validation, and test sets.
        
        Args:
            df: DataFrame to split
            
        Returns:
            Tuple of (train_df, val_df, test_df)
        """
        logger.info(f"Splitting data: test_size={self.test_size}, val_size={self.val_size}")
        
        # First split: train+val vs test
        train_val_df, test_df = train_test_split(
            df,
            test_size=self.test_size,
            random_state=self.random_state,
            stratify=df["target"]
        )
        
        # Second split: train vs val
        # Adjust val_size to be relative to train+val
        val_relative_size = self.val_size / (1 - self.test_size)
        train_df, val_df = train_test_split(
            train_val_df,
            test_size=val_relative_size,
            random_state=self.random_state,
            stratify=train_val_df["target"]
        )
        
        logger.info(f"Split sizes - Train: {len(train_df)}, Val: {len(val_df)}, Test: {len(test_df)}")
        
        # Log class distributions
        for name, dataset in [("Train", train_df), ("Val", val_df), ("Test", test_df)]:
            dist = dataset["target"].value_counts(normalize=True).sort_index()
            logger.info(f"{name} distribution: {dist.to_dict()}")
        
        return train_df, val_df, test_df
    
    def save_splits(
        self,
        train_df: pd.DataFrame,
        val_df: pd.DataFrame,
        test_df: pd.DataFrame,
        output_dir: Path
    ):
        """
        Save train, validation, and test splits to CSV files.
        
        Args:
            train_df: Training data
            val_df: Validation data
            test_df: Test data
            output_dir: Directory to save the splits
        """
        output_dir.mkdir(parents=True, exist_ok=True)
        
        train_path = output_dir / "train_split.csv"
        val_path = output_dir / "val_split.csv"
        test_path = output_dir / "test_split.csv"
        
        train_df.to_csv(train_path, index=False)
        val_df.to_csv(val_path, index=False)
        test_df.to_csv(test_path, index=False)
        
        logger.info(f"Splits saved to {output_dir}")
        logger.info(f"  - Train: {train_path}")
        logger.info(f"  - Val: {val_path}")
        logger.info(f"  - Test: {test_path}")
    
    def process(self, input_path: Path, output_dir: Path) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        """
        Complete pipeline: load, split, and save data.
        
        Args:
            input_path: Path to input CSV file
            output_dir: Directory to save splits
            
        Returns:
            Tuple of (train_df, val_df, test_df)
        """
        # Load data
        df = self.load_data(input_path)
        
        # Split data
        train_df, val_df, test_df = self.split_data(df)
        
        # Save splits
        self.save_splits(train_df, val_df, test_df, output_dir)
        
        return train_df, val_df, test_df


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    # Example usage
    splitter = DataSplitter(
        test_size=0.2,
        val_size=0.1,
        random_state=42,
        binary_only=False
    )
    
    input_path = Path("data/sentiment140/train.csv")
    output_dir = Path("data/sentiment140")
    
    train_df, val_df, test_df = splitter.process(input_path, output_dir)
    
    print(f"\nFinal split sizes:")
    print(f"Train: {len(train_df):,}")
    print(f"Val: {len(val_df):,}")
    print(f"Test: {len(test_df):,}")
