src/preprocessing/transformer_user_data_preparation.py

#!/usr/bin/env python3
"""
Transformer User Data Preparation

Enhanced user dataset creation without interaction history caps for transformer-based models.
Supports variable-length sequences with proper attention masking.
"""

import pandas as pd
import numpy as np
import tensorflow as tf
from typing import Dict, List, Tuple, Optional
from datetime import datetime
import pickle
import os
from functools import lru_cache

from src.preprocessing.data_loader import DataProcessor


class TransformerUserDatasetCreator:
    """Creates user training dataset with uncapped interaction histories for transformer models."""
    
    def __init__(self, 
                 max_history_length: Optional[int] = None,  # None = no cap
                 min_history_length: int = 1,
                 artifacts_prefix: str = "transformer_"):
        """
        Initialize transformer user dataset creator.
        
        Args:
            max_history_length: Maximum history length (None for no cap)
            min_history_length: Minimum interactions required for a user
            artifacts_prefix: Prefix for transformer artifact files
        """
        self.max_history_length = max_history_length
        self.min_history_length = min_history_length
        self.artifacts_prefix = artifacts_prefix
        self.data_processor = DataProcessor()
    
    def categorize_age(self, age: float) -> int:
        """Categorize age into 6 demographic groups."""
        if age < 18:
            return 0  # Teen
        elif age < 26:
            return 1  # Young Adult
        elif age < 36:
            return 2  # Adult
        elif age < 51:
            return 3  # Middle Age
        elif age < 66:
            return 4  # Mature
        else:
            return 5  # Senior
    
    def categorize_income(self, income_series: pd.Series) -> np.ndarray:
        """Categorize income into 5 percentile-based groups."""
        percentiles = [0, 20, 40, 60, 80, 100]
        income_thresholds = np.percentile(income_series, percentiles)
        categories = np.digitize(income_series, income_thresholds[1:-1])
        categories = np.clip(categories, 0, 4)
        return categories.astype(np.int32)
    
    def categorize_profession(self, profession: str) -> int:
        """Categorize profession into numeric categories."""
        profession_map = {
            "Technology": 0, "Healthcare": 1, "Education": 2, "Finance": 3,
            "Retail": 4, "Manufacturing": 5, "Services": 6, "Other": 7
        }
        return profession_map.get(profession, 7)
    
    def categorize_location(self, location: str) -> int:
        """Categorize location into numeric categories."""
        location_map = {"Urban": 0, "Suburban": 1, "Rural": 2}
        return location_map.get(location, 0)
    
    def categorize_education_level(self, education: str) -> int:
        """Categorize education level into numeric categories."""
        education_map = {
            "High School": 0, "Some College": 1, "Bachelor's": 2,
            "Master's": 3, "PhD+": 4
        }
        return education_map.get(education, 0)
    
    def categorize_marital_status(self, marital_status: str) -> int:
        """Categorize marital status into numeric categories."""
        marital_map = {"Single": 0, "Married": 1, "Divorced": 2, "Widowed": 3}
        return marital_map.get(marital_status, 0)
    
    @lru_cache(maxsize=1)
    def load_item_embeddings(self, embeddings_path: str = None) -> Dict[int, np.ndarray]:
        """Load transformer item embeddings with caching."""
        if embeddings_path is None:
            embeddings_path = f"src/artifacts/transformers/{self.artifacts_prefix}item_embeddings.npy"
        
        try:
            embeddings = np.load(embeddings_path, allow_pickle=True).item()
            print(f"Loaded {len(embeddings)} transformer item embeddings")
            return embeddings
        except FileNotFoundError:
            print(f"Warning: {embeddings_path} not found. Creating dummy embeddings...")
            # Create dummy embeddings for demo purposes
            processor = DataProcessor()
            items_df, users_df, interactions_df = processor.load_data()
            
            num_items = len(items_df['product_id'].unique())
            item_ids = items_df['product_id'].unique()
            embedding_matrix = np.random.rand(num_items, 128).astype(np.float32)
            
            dummy_embeddings = dict(zip(item_ids, embedding_matrix))
            print(f"Created dummy embeddings for {len(dummy_embeddings)} items")
            return dummy_embeddings
    
    def create_variable_length_user_histories(self, 
                                            interactions_df: pd.DataFrame,
                                            items_df: pd.DataFrame) -> Tuple[Dict[int, List[int]], Dict[int, int]]:
        """Create user interaction histories without length restrictions."""
        
        # Convert timestamp and sort
        interactions_df = interactions_df.copy()
        interactions_df['event_time'] = pd.to_datetime(interactions_df['event_time'], utc=True)
        interactions_sorted = interactions_df.sort_values(['user_id', 'event_time'])
        
        # Build user histories without caps
        user_histories = {}
        user_history_lengths = {}
        
        for user_id, user_interactions in interactions_sorted.groupby('user_id'):
            item_ids = []
            for _, row in user_interactions.iterrows():
                item_vocab_id = self.data_processor.item_vocab.get(row['product_id'], 0)
                item_ids.append(item_vocab_id)
            
            # Only include users with minimum interactions
            if len(item_ids) >= self.min_history_length:
                # Apply max cap only if specified
                if self.max_history_length is not None and len(item_ids) > self.max_history_length:
                    item_ids = item_ids[-self.max_history_length:]
                
                user_histories[user_id] = item_ids
                user_history_lengths[user_id] = len(item_ids)
        
        print(f"Created variable-length histories for {len(user_histories)} users")
        print(f"History length stats:")
        lengths = list(user_history_lengths.values())
        print(f"  Min: {min(lengths)}, Max: {max(lengths)}")
        print(f"  Mean: {np.mean(lengths):.1f}, Median: {np.median(lengths):.1f}")
        
        return user_histories, user_history_lengths
    
    def aggregate_variable_length_embeddings(self, 
                                           user_histories: Dict[int, List[int]],
                                           item_embeddings: Dict[int, np.ndarray],
                                           embedding_dim: int = 128) -> Dict[int, Tuple[np.ndarray, int]]:
        """Aggregate embeddings for variable-length histories."""
        
        user_aggregated_embeddings = {}
        vocab_to_item_id = {vocab_idx: item_id for item_id, vocab_idx in self.data_processor.item_vocab.items()}
        
        for user_id, item_history in user_histories.items():
            if not item_history:
                continue
            
            # Get embeddings for interaction history
            history_embeddings = []
            for vocab_idx in item_history:
                actual_item_id = vocab_to_item_id.get(vocab_idx)
                
                if actual_item_id and actual_item_id in item_embeddings:
                    history_embeddings.append(item_embeddings[actual_item_id])
                else:
                    # Use zero embedding for unknown items
                    history_embeddings.append(np.zeros(embedding_dim))
            
            if history_embeddings:
                history_embeddings = np.array(history_embeddings)
                sequence_length = len(history_embeddings)
                
                # Store embedding sequence and its actual length
                user_aggregated_embeddings[user_id] = (history_embeddings, sequence_length)
        
        return user_aggregated_embeddings
    
    def create_padded_sequences(self, 
                              user_embeddings: Dict[int, Tuple[np.ndarray, int]],
                              max_sequence_length: Optional[int] = None) -> Dict[int, Dict]:
        """Create padded sequences with attention masks for transformer input."""
        
        if max_sequence_length is None:
            # Determine max length from data
            max_sequence_length = max(length for _, length in user_embeddings.values())
            print(f"Auto-determined max sequence length: {max_sequence_length}")
        
        padded_user_data = {}
        
        for user_id, (embeddings, actual_length) in user_embeddings.items():
            # Pad sequences to max length
            if actual_length < max_sequence_length:
                padding_length = max_sequence_length - actual_length
                embedding_dim = embeddings.shape[1]
                
                # Pad with zeros at the end
                padding = np.zeros((padding_length, embedding_dim))
                padded_embeddings = np.vstack([embeddings, padding])
            else:
                # Truncate if longer than max (shouldn't happen with proper max calculation)
                padded_embeddings = embeddings[:max_sequence_length]
                actual_length = max_sequence_length
            
            # Create attention mask (1 for real tokens, 0 for padding)
            attention_mask = np.zeros(max_sequence_length)
            attention_mask[:actual_length] = 1
            
            padded_user_data[user_id] = {
                'embeddings': padded_embeddings.astype(np.float32),
                'attention_mask': attention_mask.astype(np.float32),
                'sequence_length': actual_length
            }
        
        print(f"Created padded sequences for {len(padded_user_data)} users")
        print(f"Sequence length: {max_sequence_length}")
        
        return padded_user_data, max_sequence_length
    
    def prepare_user_features(self, 
                            users_df: pd.DataFrame,
                            padded_user_data: Dict[int, Dict]) -> Dict[str, np.ndarray]:
        """Prepare user features combining demographics and variable-length interaction sequences."""
        
        # Filter users that have both demographics and interaction data
        valid_users = set(users_df['user_id']) & set(padded_user_data.keys())
        valid_users = sorted(list(valid_users))
        
        # Prepare demographic features
        user_demographics = users_df[users_df['user_id'].isin(valid_users)].copy()
        user_demographics = user_demographics.sort_values('user_id')
        
        # Convert demographics to categorical
        user_demographics['gender_numeric'] = (user_demographics['gender'] == 'male').astype(int)
        user_demographics['age_category'] = user_demographics['age'].apply(self.categorize_age)
        user_demographics['income_category'] = self.categorize_income(user_demographics['income'])
        user_demographics['profession_category'] = user_demographics['profession'].apply(self.categorize_profession)
        user_demographics['location_category'] = user_demographics['location'].apply(self.categorize_location)
        user_demographics['education_category'] = user_demographics['education_level'].apply(self.categorize_education_level)
        user_demographics['marital_category'] = user_demographics['marital_status'].apply(self.categorize_marital_status)
        
        # Create user features with variable-length sequences
        user_features = {
            'user_ids': user_demographics['user_id'].values,
            'age': user_demographics['age_category'].values.astype(np.int32),
            'gender': user_demographics['gender_numeric'].values.astype(np.int32),
            'income': user_demographics['income_category'].values.astype(np.int32),
            'profession': user_demographics['profession_category'].values.astype(np.int32),
            'location': user_demographics['location_category'].values.astype(np.int32),
            'education_level': user_demographics['education_category'].values.astype(np.int32),
            'marital_status': user_demographics['marital_category'].values.astype(np.int32),
        }
        
        # Add interaction sequences and masks
        embeddings_list = []
        attention_masks_list = []
        sequence_lengths_list = []
        
        for user_id in user_demographics['user_id']:
            user_data = padded_user_data[user_id]
            embeddings_list.append(user_data['embeddings'])
            attention_masks_list.append(user_data['attention_mask'])
            sequence_lengths_list.append(user_data['sequence_length'])
        
        user_features['item_history_embeddings'] = np.array(embeddings_list)
        user_features['attention_masks'] = np.array(attention_masks_list)
        user_features['sequence_lengths'] = np.array(sequence_lengths_list, dtype=np.int32)
        
        print(f"Prepared transformer user features for {len(valid_users)} users")
        print(f"Feature shapes:")
        for key, value in user_features.items():
            if isinstance(value, np.ndarray):
                print(f"  {key}: {value.shape}")
        
        return user_features
    
    def create_temporal_split(self, 
                            interactions_df: pd.DataFrame,
                            split_date: str = "2019-11-15") -> Tuple[pd.DataFrame, pd.DataFrame]:
        """Split interactions temporally for training and validation."""
        interactions_df = interactions_df.copy()
        interactions_df['event_time'] = pd.to_datetime(interactions_df['event_time'], utc=True)
        split_timestamp = pd.to_datetime(split_date, utc=True)
        
        train_interactions = interactions_df[interactions_df['event_time'] < split_timestamp]
        val_interactions = interactions_df[interactions_df['event_time'] >= split_timestamp]
        
        print(f"Temporal split:")
        print(f"  Training interactions: {len(train_interactions)} (before {split_date})")
        print(f"  Validation interactions: {len(val_interactions)} (after {split_date})")
        
        return train_interactions, val_interactions
    
    def create_training_dataset(self, 
                              interactions_df: pd.DataFrame,
                              items_df: pd.DataFrame,
                              users_df: pd.DataFrame,
                              item_embeddings: Dict[int, np.ndarray],
                              negative_samples_per_positive: int = 4,
                              vocab_path: str = "src/artifacts/transformers/transformer_vocabularies.pkl") -> Tuple[Dict[str, np.ndarray], int]:
        """Create complete training dataset with variable-length sequences."""
        
        if os.path.exists(vocab_path):
            print(f"📂 Loading vocabularies from {vocab_path}")
            self.data_processor.load_vocabularies(vocab_path)
        else:
            raise FileNotFoundError(
                f"❌ Vocabularies not found at {vocab_path}. "
                "Run Phase 1 (item pretraining) first to generate them."
            )
        # Create variable-length user histories
        print("Creating variable-length user interaction histories...")
        user_histories, user_history_lengths = self.create_variable_length_user_histories(
            interactions_df, items_df
        )
        
        # Aggregate embeddings for variable-length sequences
        print("Aggregating variable-length embeddings...")
        user_embeddings = self.aggregate_variable_length_embeddings(
            user_histories, item_embeddings
        )
        
        # Create padded sequences with attention masks
        print("Creating padded sequences with attention masks...")
        padded_user_data, max_sequence_length = self.create_padded_sequences(user_embeddings)
        
        # Create positive/negative pairs
        print("Creating positive/negative pairs...")
        training_pairs = self.data_processor.create_positive_negative_pairs(
            interactions_df, items_df, negative_samples_per_positive
        )
        
        # Prepare user features
        user_features = self.prepare_user_features(users_df, padded_user_data)
        
        # Prepare item features
        item_features = self.data_processor.prepare_item_features(items_df)
        
        # Create aligned dataset
        print("Creating aligned training dataset...")
        valid_pairs = []
        for _, row in training_pairs.iterrows():
            user_id = row['user_id']
            item_id = row['product_id']
            rating = row['rating']
            
            if (user_id in self.data_processor.user_vocab and 
                item_id in self.data_processor.item_vocab and
                user_id in padded_user_data):
                valid_pairs.append({
                    'user_id': user_id,
                    'product_id': item_id,
                    'rating': rating
                })
        
        valid_pairs_df = pd.DataFrame(valid_pairs)
        
        # Create feature arrays for training
        training_features = {}
        
        # Map users to feature indices
        user_id_to_index = {uid: idx for idx, uid in enumerate(user_features['user_ids'])}
        user_indices = []
        valid_user_pairs = []
        
        for _, row in valid_pairs_df.iterrows():
            user_id = row['user_id']
            if user_id in user_id_to_index:
                user_indices.append(user_id_to_index[user_id])
                valid_user_pairs.append(row)
        
        if len(valid_user_pairs) == 0:
            print("Warning: No valid user-item pairs found!")
            return {}, max_sequence_length
        
        valid_pairs_df = pd.DataFrame(valid_user_pairs)
        
        # User features for each pair
        training_features['age'] = user_features['age'][user_indices]
        training_features['gender'] = user_features['gender'][user_indices]
        training_features['income'] = user_features['income'][user_indices]
        training_features['profession'] = user_features['profession'][user_indices]
        training_features['location'] = user_features['location'][user_indices]
        training_features['education_level'] = user_features['education_level'][user_indices]
        training_features['marital_status'] = user_features['marital_status'][user_indices]
        training_features['item_history_embeddings'] = user_features['item_history_embeddings'][user_indices]
        training_features['attention_masks'] = user_features['attention_masks'][user_indices]
        training_features['sequence_lengths'] = user_features['sequence_lengths'][user_indices]
        
        # Item features for each pair
        item_indices = [self.data_processor.item_vocab[iid] for iid in valid_pairs_df['product_id']]
        training_features['product_id'] = item_features['product_id'][item_indices]
        training_features['category_id'] = item_features['category_id'][item_indices]
        training_features['category_code_id'] = item_features['category_code_id'][item_indices]
        training_features['brand_id'] = item_features['brand_id'][item_indices]
        training_features['price'] = item_features['price'][item_indices]
        
        # Ratings
        training_features['rating'] = valid_pairs_df['rating'].values.astype(np.float32)
        
        print(f"Created transformer training dataset with {len(valid_pairs_df)} samples")
        print(f"Max sequence length: {max_sequence_length}")
        
        return training_features, max_sequence_length
    
    def save_dataset(self, 
                    training_features: Dict[str, np.ndarray],
                    max_sequence_length: int,
                    save_path: str = "src/artifacts/transformers/"):
        """Save the transformer training dataset."""
        
        os.makedirs(save_path, exist_ok=True)
        
        # Save features
        features_path = f"{save_path}/{self.artifacts_prefix}training_features.pkl"
        with open(features_path, 'wb') as f:
            pickle.dump(training_features, f)
        
        # Save dataset statistics
        stats = {
            'num_samples': len(training_features['rating']),
            'num_positive': np.sum(training_features['rating'] > 0.5),
            'num_negative': np.sum(training_features['rating'] <= 0.5),
            'max_sequence_length': max_sequence_length,
            'embedding_dim': training_features['item_history_embeddings'].shape[2],
            'has_attention_masks': True,
            'has_variable_lengths': True
        }
        
        stats_path = f"{save_path}/{self.artifacts_prefix}dataset_stats.txt"
        with open(stats_path, 'w') as f:
            for key, value in stats.items():
                f.write(f"{key}: {value}\n")

        # Save vocabularies with transformer prefix
        vocab_path = f"{save_path}/{self.artifacts_prefix}vocabularies.pkl"
        vocab_data = {
            'item_vocab': self.data_processor.item_vocab,
            'category_vocab': self.data_processor.category_vocab,
            'category_code_vocab': self.data_processor.category_code_vocab,
            'brand_vocab': self.data_processor.brand_vocab,
            'user_vocab': self.data_processor.user_vocab
        }
        with open(vocab_path, 'wb') as f:
            pickle.dump(vocab_data, f)

        print(f"✅ Transformer training dataset saved:")
        print(f"   - Features: {features_path}")
        print(f"   - Stats: {stats_path}")
        print(f"   - Vocabularies: {vocab_path}")
        print(f"   - Dataset statistics: {stats}")
    
    def load_dataset(self, load_path: str = None) -> Dict[str, np.ndarray]:
        """Load saved transformer training dataset."""
        if load_path is None:
            load_path = f"src/artifacts/transformers/{self.artifacts_prefix}training_features.pkl"
        
        with open(load_path, 'rb') as f:
            training_features = pickle.load(f)
        
        print(f"Loaded transformer training dataset with {len(training_features['rating'])} samples")
        return training_features


def main():
    """Main function for transformer user dataset creation."""
    
    print("🚀 Starting Transformer User Dataset Creation")
    print("=" * 60)
    
    # Initialize dataset creator (no history cap)
    dataset_creator = TransformerUserDatasetCreator(
        max_history_length=None,  # No cap for transformer version
        min_history_length=1
    )
    
    # Load data
    print("Loading data...")
    data_processor = DataProcessor()
    items_df, users_df, interactions_df = data_processor.load_data()
    
    # Load transformer item embeddings
    print("Loading transformer item embeddings...")
    item_embeddings = dataset_creator.load_item_embeddings()
    
    # Use full dataset
    print("Using full dataset...")
    sample_users = users_df
    user_ids = set(sample_users['user_id'])
    
    # Filter interactions to users
    sample_interactions = interactions_df[interactions_df['user_id'].isin(user_ids)]
    
    # Filter items to those in interactions
    item_ids = set(sample_interactions['product_id'])
    sample_items = items_df[items_df['product_id'].isin(item_ids)]
    
    print(f"Full dataset: {len(sample_items)} items, {len(sample_users)} users, {len(sample_interactions)} interactions")
    
    # Create temporal split
    print("Creating temporal split...")
    train_interactions, val_interactions = dataset_creator.create_temporal_split(sample_interactions)
    
    # Create training dataset
    print("Creating transformer training dataset...")
    training_features, max_sequence_length = dataset_creator.create_training_dataset(
        train_interactions, sample_items, sample_users, item_embeddings,
        negative_samples_per_positive=2
    )
    
    # Save training dataset
    print("Saving transformer training dataset...")
    dataset_creator.save_dataset(training_features, max_sequence_length)
    
    # Create validation dataset
    print("Creating transformer validation dataset...")
    val_sample_size = min(5000, max(len(val_interactions) // 10, len(val_interactions)))
    val_sample = val_interactions.sample(val_sample_size) if val_sample_size > 0 and val_sample_size < len(val_interactions) else val_interactions
    
    val_training_features, _ = dataset_creator.create_training_dataset(
        val_sample, sample_items, sample_users, item_embeddings,
        negative_samples_per_positive=1
    )
    
    # Save validation dataset
    val_path = f"src/artifacts/transformers/{dataset_creator.artifacts_prefix}validation_features.pkl"
    with open(val_path, 'wb') as f:
        pickle.dump(val_training_features, f)
    
    print("✅ Transformer User Dataset Creation Completed!")
    print(f"   - Max sequence length: {max_sequence_length}")
    print(f"   - Training samples: {len(training_features['rating'])}")
    print(f"   - Validation samples: {len(val_training_features['rating'])}")


if __name__ == "__main__":
    main()