"""
Data Preprocessor
=================

This module provides a complete preprocessing pipeline for the Sentiment140 dataset.
"""

import logging
from pathlib import Path
from typing import Tuple, Optional
import pandas as pd
import numpy as np
from sklearn.pipeline import Pipeline

from .data_cleaner import TextCleaner
from .data_splitter import DataSplitter

logger = logging.getLogger(__name__)


class DataPreprocessor:
    """
    Complete preprocessing pipeline for Sentiment140 data.
    """
    
    def __init__(
        self,
        use_spacy_lemma: bool = False,
        test_size: float = 0.2,
        val_size: float = 0.1,
        random_state: int = 42,
        binary_only: bool = False
    ):
        """
        Initialize the data preprocessor.
        
        Args:
            use_spacy_lemma: Whether to use spaCy for lemmatization
            test_size: Proportion of data for testing
            val_size: Proportion of training data for validation
            random_state: Random seed for reproducibility
            binary_only: If True, only keep negative (0) and positive (4) labels
        """
        self.cleaner = TextCleaner(use_spacy_lemma=use_spacy_lemma)
        self.splitter = DataSplitter(
            test_size=test_size,
            val_size=val_size,
            random_state=random_state,
            binary_only=binary_only
        )
        
    def preprocess_and_split(
        self,
        input_path: Path,
        output_dir: Path,
        save_splits: bool = True
    ) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        """
        Complete preprocessing pipeline: load, split, and optionally save data.
        
        Args:
            input_path: Path to input CSV file
            output_dir: Directory to save splits
            save_splits: Whether to save the splits to disk
            
        Returns:
            Tuple of (train_df, val_df, test_df)
        """
        logger.info("Starting data preprocessing pipeline...")
        
        # Load and split data
        train_df, val_df, test_df = self.splitter.process(input_path, output_dir)
        
        logger.info("Data preprocessing completed successfully")
        
        return train_df, val_df, test_df
    
    def get_preprocessed_data(
        self,
        df: pd.DataFrame
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Get preprocessed features and labels from a DataFrame.
        
        Args:
            df: DataFrame with 'text' and 'target' columns
            
        Returns:
            Tuple of (X, y) where X is cleaned texts and y is labels
        """
        X = df["text"].tolist()
        y = df["target"].to_numpy()
        
        # Clean texts
        X_cleaned = self.cleaner.transform(X)
        
        return X_cleaned, y


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    # Example usage
    preprocessor = DataPreprocessor(
        use_spacy_lemma=False,
        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 = preprocessor.preprocess_and_split(input_path, output_dir)
    
    print(f"\nPreprocessed data ready:")
    print(f"Train: {len(train_df):,} samples")
    print(f"Val: {len(val_df):,} samples")
    print(f"Test: {len(test_df):,} samples")
