data_preprocessing.py

"""
Data Preprocessing Module for Sentiment Analysis
Supports IMDB, Sentiment140, and Yelp Review datasets
"""

import pandas as pd
import numpy as np
import re
import os
from typing import Tuple, List, Optional
from sklearn.model_selection import train_test_split
try:
    from tensorflow.keras.preprocessing.text import Tokenizer
    from tensorflow.keras.preprocessing.sequence import pad_sequences
except ImportError:
    # For newer TensorFlow/Keras versions
    from keras.preprocessing.text import Tokenizer
    from keras.preprocessing.sequence import pad_sequences
import pickle


class DataPreprocessor:
    """Handles data loading and preprocessing for sentiment analysis"""
    
    def __init__(self, max_vocab_size: int = 10000, max_sequence_length: int = 200):
        self.max_vocab_size = max_vocab_size
        self.max_sequence_length = max_sequence_length
        self.tokenizer = None
        self.vocab_size = 0
        
    def clean_text(self, text: str) -> str:
        """Clean and preprocess text data"""
        if pd.isna(text):
            return ""
        
        # Convert to string
        text = str(text)
        
        # Remove URLs
        text = re.sub(r'http\S+|www\S+|https\S+', '', text, flags=re.MULTILINE)
        
        # Remove user mentions and hashtags (for Twitter data)
        text = re.sub(r'@\w+|#\w+', '', text)
        
        # Remove special characters but keep spaces
        text = re.sub(r'[^a-zA-Z\s]', '', text)
        
        # Convert to lowercase
        text = text.lower()
        
        # Remove extra whitespace
        text = ' '.join(text.split())
        
        return text
    
    def load_imdb_data(self, data_path: str) -> pd.DataFrame:
        """Load IMDB dataset from CSV file"""
        # Expected format: CSV with 'review' and 'sentiment' columns
        # sentiment: 'positive' or 'negative'
        df = pd.read_csv(data_path)
        
        # Standardize column names
        if 'sentiment' in df.columns:
            df['label'] = df['sentiment'].map({'positive': 1, 'negative': 0})
        elif 'label' in df.columns:
            pass  # Already has label column
        else:
            raise ValueError("Dataset must have 'sentiment' or 'label' column")
        
        if 'review' in df.columns:
            df['text'] = df['review']
        elif 'text' not in df.columns:
            raise ValueError("Dataset must have 'review' or 'text' column")
        
        return df[['text', 'label']]
    
    def load_sentiment140_data(self, data_path: str) -> pd.DataFrame:
        """Load Sentiment140 dataset from CSV file"""
        # Expected format: CSV with columns including text and sentiment
        # sentiment: 0 (negative), 2 (neutral), 4 (positive)
        df = pd.read_csv(data_path, encoding='latin-1')

        # If header is missing (common in Sentiment140), pandas will assign integer column names.
        # The canonical order is: target, ids, date, flag, user, text
        if 'target' not in df.columns and 'sentiment' not in df.columns:
            try:
                df = pd.read_csv(
                    data_path,
                    encoding='latin-1',
                    header=None,
                    names=['target', 'ids', 'date', 'flag', 'user', 'text']
                )
            except Exception:
                pass
        
        # Sentiment140 typically has columns: target, ids, date, flag, user, text
        # Map sentiment: 0->0 (negative), 4->1 (positive), filter out 2 (neutral)
        if 'target' in df.columns:
            df = df[df['target'].isin([0, 4])]  # Remove neutral
            df['label'] = df['target'].map({0: 0, 4: 1})
        elif 'sentiment' in df.columns:
            df = df[df['sentiment'].isin([0, 4])]
            df['label'] = df['sentiment'].map({0: 0, 4: 1})
        else:
            raise ValueError("Dataset must have 'target' or 'sentiment' column")
        
        if 'text' not in df.columns:
            raise ValueError("Dataset must have 'text' column")
        
        return df[['text', 'label']]
    
    def load_yelp_data(self, data_path: str) -> pd.DataFrame:
        """Load Yelp Review dataset from CSV file"""
        # Expected format: CSV with 'text' and 'stars' or 'rating' columns
        df = pd.read_csv(data_path)
        
        # Map ratings to binary labels
        if 'stars' in df.columns:
            # 1-2 stars: negative (0), 4-5 stars: positive (1), filter 3 stars
            df = df[df['stars'].isin([1, 2, 4, 5])]
            df['label'] = df['stars'].map({1: 0, 2: 0, 4: 1, 5: 1})
        elif 'rating' in df.columns:
            df = df[df['rating'].isin([1, 2, 4, 5])]
            df['label'] = df['rating'].map({1: 0, 2: 0, 4: 1, 5: 1})
        elif 'sentiment' in df.columns:
            df['label'] = df['sentiment'].map({'positive': 1, 'negative': 0})
        else:
            raise ValueError("Dataset must have 'stars', 'rating', or 'sentiment' column")
        
        if 'text' not in df.columns and 'review' in df.columns:
            df['text'] = df['review']
        elif 'text' not in df.columns:
            raise ValueError("Dataset must have 'text' or 'review' column")
        
        return df[['text', 'label']]
    
    def load_data(self, data_path: str, dataset_type: str = 'auto') -> pd.DataFrame:
        """Load data from specified path with automatic type detection"""
        if dataset_type == 'auto':
            # Try to detect dataset type from filename or path
            if 'imdb' in data_path.lower():
                dataset_type = 'imdb'
            elif 'sentiment140' in data_path.lower() or 'sentiment' in data_path.lower():
                dataset_type = 'sentiment140'
            elif 'yelp' in data_path.lower():
                dataset_type = 'yelp'
            else:
                # Default to generic CSV format
                dataset_type = 'generic'
        
        if dataset_type == 'imdb':
            df = self.load_imdb_data(data_path)
        elif dataset_type == 'sentiment140':
            df = self.load_sentiment140_data(data_path)
        elif dataset_type == 'yelp':
            df = self.load_yelp_data(data_path)
        else:
            # Generic CSV loader
            df = pd.read_csv(data_path)
            if 'text' not in df.columns and 'review' in df.columns:
                df['text'] = df['review']
            if 'label' not in df.columns and 'sentiment' in df.columns:
                df['label'] = df['sentiment'].map({'positive': 1, 'negative': 0})
        
        return df
    
    def preprocess_data(self, df: pd.DataFrame, fit_tokenizer: bool = True) -> Tuple[np.ndarray, np.ndarray]:
        """Preprocess text data and convert to sequences"""
        # Clean text
        df['text'] = df['text'].apply(self.clean_text)
        
        # Remove empty texts
        df = df[df['text'].str.len() > 0]
        
        texts = df['text'].values
        labels = df['label'].values
        
        # Tokenize and convert to sequences
        if fit_tokenizer:
            self.tokenizer = Tokenizer(num_words=self.max_vocab_size, oov_token='<OOV>')
            self.tokenizer.fit_on_texts(texts)
            self.vocab_size = min(len(self.tokenizer.word_index) + 1, self.max_vocab_size)
        else:
            if self.tokenizer is None:
                raise ValueError("Tokenizer not fitted. Set fit_tokenizer=True for first call.")
        
        sequences = self.tokenizer.texts_to_sequences(texts)
        padded_sequences = pad_sequences(sequences, maxlen=self.max_sequence_length, padding='post', truncating='post')
        
        return padded_sequences, labels
    
    def save_tokenizer(self, filepath: str):
        """Save tokenizer to file"""
        if self.tokenizer is None:
            raise ValueError("Tokenizer not fitted yet")
        with open(filepath, 'wb') as f:
            pickle.dump(self.tokenizer, f)
    
    def load_tokenizer(self, filepath: str):
        """Load tokenizer from file"""
        with open(filepath, 'rb') as f:
            self.tokenizer = pickle.load(f)
        self.vocab_size = min(len(self.tokenizer.word_index) + 1, self.max_vocab_size)
    
    def prepare_train_test_split(self, data_path: str, dataset_type: str = 'auto', 
                                 test_size: float = 0.2, random_state: int = 42) -> Tuple:
        """Load data, preprocess, and split into train/test sets"""
        # Load data
        df = self.load_data(data_path, dataset_type)
        
        # Preprocess
        X, y = self.preprocess_data(df, fit_tokenizer=True)
        
        # Split into train and test
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, random_state=random_state, stratify=y
        )
        
        return X_train, X_test, y_train, y_test