data_processing.py

"""
Comprehensive Multi-Dataset Processing for Insider Threat Detection
This script generates all features into combined_features.csv file.
It processes data from the r4.2 folder and efficiently creates all columns.
"""

import os
import time
import numpy as np
import pandas as pd
from collections import defaultdict
import warnings
import argparse
warnings.filterwarnings('ignore')

# Set global parameters for processing
SAMPLE_SIZE = 100000  # Samples per dataset
CHUNK_SIZE = 50000    # Chunk size for large files
RANDOM_SEED = 42      # Random seed for reproducibility

def load_and_preprocess_data(file_path, sample_size=None, chunk_size=CHUNK_SIZE):
    """
    Load and preprocess data from CSV file with efficient memory handling
    
    Parameters:
    -----------
    file_path : str
        Path to the CSV file
    sample_size : int, optional
        Number of rows to sample (for large files)
    chunk_size : int, optional
        Chunk size for reading large files
        
    Returns:
    --------
    df : DataFrame
        Preprocessed dataframe
    """
    print(f"Loading data from {file_path}")
    file_size = os.path.getsize(file_path)
    print(f"File size: {file_size / (1024 * 1024):.2f} MB")
    
    try:
        # For smaller files, load directly
        if file_size < 500e6:  # If less than 500MB
            if sample_size:
                df = pd.read_csv(file_path, nrows=sample_size)
                print(f"Loaded sample of {sample_size} rows")
            else:
                df = pd.read_csv(file_path)
                print(f"Loaded entire file: {df.shape[0]} rows")
        else:
            # For larger files, use chunking
            chunks = []
            total_rows = 0
            
            for chunk_num, chunk in enumerate(pd.read_csv(file_path, chunksize=chunk_size)):
                total_rows += len(chunk)
                if sample_size and total_rows > sample_size:
                    # Trim the last chunk to match sample_size exactly
                    overflow = total_rows - sample_size
                    chunk = chunk.iloc[:-overflow] if overflow < len(chunk) else chunk
                    chunks.append(chunk)
                    print(f"Processed {chunk_num+1} chunks, reached sample size of {sample_size} rows")
                    break
                chunks.append(chunk)
                print(f"Processed chunk {chunk_num+1} with {len(chunk)} rows, total rows so far: {total_rows}")
                
                # If not sampling and we've read the whole file
                if not sample_size and chunk_num % 10 == 0:
                    # Periodically consolidate chunks to save memory
                    if len(chunks) > 5:
                        print(f"Consolidating {len(chunks)} chunks to save memory...")
                        consolidated = pd.concat(chunks)
                        chunks = [consolidated]
                        
            df = pd.concat(chunks)
            if sample_size and len(df) > sample_size:
                df = df.iloc[:sample_size]
                
            if sample_size:
                print(f"Loaded {len(df)} rows using chunking (sample size: {sample_size})")
            else:
                print(f"Loaded entire file: {len(df)} rows using chunking")
    except Exception as e:
        print(f"Error loading {file_path}: {e}")
        # Try alternative approach with reduced columns
        try:
            print("Trying alternative loading approach...")
            # Read just the first row to get column names
            cols = pd.read_csv(file_path, nrows=1).columns.tolist()
            
            # Identify potentially useful columns (contains id, user, time, date, etc.)
            useful_cols = [col for col in cols if any(key in col.lower() for key in 
                          ['id', 'user', 'time', 'date', 'name', 'content', 'pc', 'activity'])]
            
            # If we couldn't identify useful columns, try a small subset
            if len(useful_cols) < 5 and len(cols) > 10:
                useful_cols = cols[:10]  # Take first 10 columns
                
            # Read only the useful columns
            if sample_size:
                df = pd.read_csv(file_path, usecols=useful_cols, nrows=sample_size)
            else:
                # Use chunking for full file with selected columns
                chunks = []
                for chunk in pd.read_csv(file_path, usecols=useful_cols, chunksize=chunk_size):
                    chunks.append(chunk)
                df = pd.concat(chunks)
                
        except Exception as e2:
            print(f"Alternative approach failed: {e2}")
            # Create an empty DataFrame with a minimal structure
            df = pd.DataFrame(columns=['id', 'user', 'date', 'activity'])
            print(f"WARNING: Using empty DataFrame for {file_path}")
    
    # Basic preprocessing
    # Convert string/object columns to category to save memory
    for col in df.select_dtypes(include=['object']).columns:
        if df[col].nunique() < df.shape[0] * 0.5:  # If cardinality is less than 50% of rows
            df[col] = df[col].astype('category')
    
    # Handle missing values
    for col in df.columns:
        if df[col].dtype.kind in 'fc':  # float or complex
            df[col] = df[col].fillna(df[col].mean() if not df[col].isna().all() else 0)
        else:
            # For categorical/object data, fill with mode or a placeholder
            if not df[col].isna().all():
                mode_val = df[col].mode().iloc[0] if not df[col].mode().empty else "unknown"
                df[col] = df[col].fillna(mode_val)
            else:
                df[col] = df[col].fillna("unknown")
    
    print(f"Data loaded successfully: {df.shape[0]} rows, {df.shape[1]} columns")
    return df

def identify_user_column(df):
    """
    Identify the user ID column in a dataframe
    
    Parameters:
    -----------
    df : DataFrame
        DataFrame to analyze
        
    Returns:
    --------
    user_col : str or None
        User ID column name, or None if not found
    """
    # Common user ID column names
    common_user_cols = ['user', 'user_id', 'userid', 'username', 'id', 'employee']
    
    # Check for exact column name matches
    for col_name in common_user_cols:
        if col_name in df.columns:
            return col_name
    
    # Check for columns containing user-related terms
    for col in df.columns:
        if any(user_term in col.lower() for user_term in ['user', 'employ', 'person']):
            return col
    
    # Default to the first column if no user column found
    print("WARNING: No user column identified. Using first column as proxy.")
    return df.columns[0]

def extract_datetime_features(df):
    """
    Extract datetime features from columns that appear to contain date/time information
    
    Parameters:
    -----------
    df : DataFrame
        DataFrame to process
        
    Returns:
    --------
    df : DataFrame
        DataFrame with added datetime features
    """
    # Identify potential date/time columns
    date_cols = [col for col in df.columns if any(date_term in col.lower() for date_term in 
                                              ['date', 'time', 'day', 'hour'])]
    
    for col in date_cols:
        try:
            # Convert to datetime
            df[f'{col}_dt'] = pd.to_datetime(df[col], errors='coerce')
            
            # Extract useful components
            df[f'{col}_hour'] = df[f'{col}_dt'].dt.hour
            df[f'{col}_day'] = df[f'{col}_dt'].dt.day
            df[f'{col}_weekday'] = df[f'{col}_dt'].dt.dayofweek
            df[f'{col}_month'] = df[f'{col}_dt'].dt.month
            
            # Working hours flag (8 AM to 6 PM)
            df[f'{col}_working_hours'] = ((df[f'{col}_dt'].dt.hour >= 8) & 
                                          (df[f'{col}_dt'].dt.hour < 18)).astype(int)
            
            # Weekend flag
            df[f'{col}_weekend'] = (df[f'{col}_dt'].dt.dayofweek >= 5).astype(int)
            
            # Drop the intermediate datetime column to save memory
            df = df.drop(f'{col}_dt', axis=1)
        except Exception as e:
            # Skip if column can't be converted to datetime
            print(f"Could not extract datetime features from {col}: {e}")
    
    return df

def combine_datasets(dataset_files, sample_size=SAMPLE_SIZE, chunk_size=CHUNK_SIZE):
    """
    Load and combine multiple datasets with comprehensive feature extraction
    
    Parameters:
    -----------
    dataset_files : dict
        Dictionary mapping dataset types to file paths
    sample_size : int or None
        Number of samples to use from each dataset, or None to use entire dataset
    chunk_size : int
        Chunk size for reading large files
        
    Returns:
    --------
    combined_df : DataFrame
        Combined data with all features
    """
    all_user_features = defaultdict(dict)
    dataset_types = list(dataset_files.keys())
    
    for dataset_type, file_path in dataset_files.items():
        try:
            print(f"\n--- Processing {dataset_type} dataset from {file_path} ---")
            
            # Load and preprocess the dataset
            df = load_and_preprocess_data(file_path, sample_size=sample_size, chunk_size=chunk_size)
            
            if df.empty:
                print(f"Skipping empty dataset: {dataset_type}")
                continue
            
            # Extract datetime features
            df = extract_datetime_features(df)
            
            # Identify the user column
            user_col = identify_user_column(df)
            print(f"Using {user_col} as user identifier")
            
            # For large datasets, process users in chunks to save memory
            user_groups = df.groupby(user_col)
            total_users = len(user_groups.groups)
            print(f"Processing {total_users} users in {dataset_type} dataset")
            
            # Process users in batches if there are many
            user_batch_size = 1000  # Process 1000 users at a time to save memory
            
            if total_users > user_batch_size and sample_size is None:
                user_count = 0
                for user_batch in np.array_split(list(user_groups.groups.keys()), 
                                               np.ceil(total_users / user_batch_size)):
                    for user in user_batch:
                        user_df = user_groups.get_group(user)
                        _process_user_features(user, user_df, dataset_type, all_user_features)
                        user_count += 1
                        
                        # Print progress periodically
                        if user_count % 100 == 0:
                            print(f"Processed {user_count}/{total_users} users...")
                    
                    # Garbage collection hint after each batch
                    import gc
                    gc.collect()
            else:
                # Process all users at once for smaller datasets
                for user, user_df in user_groups:
                    _process_user_features(user, user_df, dataset_type, all_user_features)
            
            print(f"Completed processing {total_users} users from {dataset_type} dataset")
            
        except Exception as e:
            print(f"Error processing {dataset_type} dataset: {e}")
    
    # Convert the dictionary to a DataFrame
    print("\nCreating combined DataFrame...")
    combined_df = pd.DataFrame.from_dict(all_user_features, orient='index')
    
    # Add the user_id column
    combined_df.reset_index(inplace=True)
    combined_df.rename(columns={'index': 'user'}, inplace=True)
    
    # Fill NaN values with 0
    combined_df.fillna(0, inplace=True)
    
    return combined_df

def _process_user_features(user, user_df, dataset_type, all_user_features):
    """Helper function to process features for a single user"""
    # Skip if user is None or NaN
    if user is None or (isinstance(user, float) and np.isnan(user)):
        return
    
    # Basic statistical features for each numeric column
    for col in user_df.select_dtypes(include=['number']).columns:
        try:
            col_data = user_df[col].dropna()
            if len(col_data) > 0:
                prefix = f"{dataset_type}_{col}"
                
                # Basic statistics
                all_user_features[user][f"{prefix}_mean"] = col_data.mean()
                all_user_features[user][f"{prefix}_std"] = col_data.std()
                all_user_features[user][f"{prefix}_min"] = col_data.min()
                all_user_features[user][f"{prefix}_max"] = col_data.max()
                all_user_features[user][f"{prefix}_count"] = len(col_data)
                
                # Additional stats for more comprehensive features
                if len(col_data) >= 3:  # Need at least 3 points for some stats
                    all_user_features[user][f"{prefix}_median"] = col_data.median()
                    all_user_features[user][f"{prefix}_q25"] = col_data.quantile(0.25)
                    all_user_features[user][f"{prefix}_q75"] = col_data.quantile(0.75)
                    
                    # Skewness and kurtosis
                    all_user_features[user][f"{prefix}_skew"] = col_data.skew()
                    all_user_features[user][f"{prefix}_kurtosis"] = col_data.kurtosis()
        except Exception as e:
            pass  # Silently continue to the next column
    
    # Count distinct values for categorical columns
    for col in user_df.select_dtypes(exclude=['number']).columns:
        try:
            prefix = f"{dataset_type}_{col}"
            all_user_features[user][f"distinct_{prefix}_count"] = user_df[col].nunique()
        except Exception:
            pass
    
    # Time-based features (if datetime features were extracted)
    time_cols = [col for col in user_df.columns if any(x in col for x in ['_hour', '_day', '_weekday', '_month'])]
    for col in time_cols:
        try:
            prefix = f"{dataset_type}_{col}"
            all_user_features[user][f"{prefix}_mean"] = user_df[col].mean()
            all_user_features[user][f"{prefix}_std"] = user_df[col].std()
        except Exception:
            pass
    
    # Working hours and weekend activity
    wh_cols = [col for col in user_df.columns if '_working_hours' in col]
    we_cols = [col for col in user_df.columns if '_weekend' in col]
    
    for col in wh_cols:
        try:
            base_col = col.replace('_working_hours', '')
            prefix = f"{dataset_type}_{base_col}"
            all_user_features[user][f"non_working_hours_{prefix}_mean"] = 1 - user_df[col].mean()
        except Exception:
            pass
    
    for col in we_cols:
        try:
            base_col = col.replace('_weekend', '')
            prefix = f"{dataset_type}_{base_col}"
            all_user_features[user][f"weekend_activity_{prefix}_mean"] = user_df[col].mean()
        except Exception:
            pass

def calculate_z_scores(df):
    """
    Calculate z-scores for all numeric columns
    
    Parameters:
    -----------
    df : DataFrame
        DataFrame with features
        
    Returns:
    --------
    df : DataFrame
        DataFrame with added z-score columns
    """
    print("Calculating z-scores...")
    
    # Focus on count columns for z-scores (to match the observed columns)
    count_cols = [col for col in df.columns if col.endswith('_count') and 'distinct' not in col]
    
    for col in count_cols:
        try:
            col_no_nulls = df[col].replace([np.inf, -np.inf], np.nan).dropna()
            if len(col_no_nulls) > 0:
                mean = col_no_nulls.mean()
                std = col_no_nulls.std()
                if std > 0:
                    df[f'z_score_{col}'] = (df[col] - mean) / std
                else:
                    df[f'z_score_{col}'] = 0
        except Exception as e:
            print(f"Error calculating z-score for {col}: {e}")
            df[f'z_score_{col}'] = 0
    
    return df

def calculate_ratio_to_median(df):
    """
    Calculate ratio to median for relevant columns
    
    Parameters:
    -----------
    df : DataFrame
        DataFrame with features
        
    Returns:
    --------
    df : DataFrame
        DataFrame with added ratio-to-median columns
    """
    print("Calculating ratios to median...")
    
    # Identify stat columns (mean, std, min, max) to calculate ratios for
    stat_cols = [col for col in df.columns if any(col.endswith(f'_{stat}') for stat in 
                                              ['mean', 'std', 'min', 'max', 'count'])]
    
    for col in stat_cols:
        try:
            col_no_nulls = df[col].replace([np.inf, -np.inf], np.nan).dropna()
            if len(col_no_nulls) > 0:
                median_val = col_no_nulls.median()
                if median_val != 0:
                    df[f'ratio_to_median_{col}'] = df[col] / median_val
        except Exception as e:
            print(f"Error calculating ratio-to-median for {col}: {e}")
    
    return df

def calculate_burstiness_features(df):
    """
    Calculate burstiness and consistency features
    
    Parameters:
    -----------
    df : DataFrame
        DataFrame with features
        
    Returns:
    --------
    df : DataFrame
        DataFrame with added burstiness features
    """
    print("Calculating burstiness features...")
    
    # Identify columns for burstiness calculation (focusing on std and max range)
    std_cols = [col for col in df.columns if col.endswith('_std')]
    max_range_cols = [col for col in df.columns if col.endswith('_max')]
    
    # Calculate burstiness (std-based)
    for col in std_cols:
        base_col = col.replace('_std', '')
        if f"{base_col}_mean" in df.columns:
            try:
                # Calculate coefficient of variation as a measure of burstiness
                mean_col = f"{base_col}_mean"
                df[f"{base_col}_std_burst"] = df[col] / df[mean_col].replace(0, 1)
            except Exception:
                pass
    
    # Calculate max range burstiness
    for col in max_range_cols:
        base_col = col.replace('_max', '')
        if f"{base_col}_min" in df.columns:
            try:
                min_col = f"{base_col}_min"
                mean_col = f"{base_col}_mean" if f"{base_col}_mean" in df.columns else None
                
                # Max range as (max - min) normalized by mean or just as a raw value
                if mean_col and (df[mean_col] != 0).all():
                    df[f"{base_col}_max_range_burst"] = (df[col] - df[min_col]) / df[mean_col]
                else:
                    df[f"{base_col}_max_range_burst"] = df[col] - df[min_col]
            except Exception:
                pass
    
    # Calculate consistency (inverse of burstiness)
    burst_cols = [col for col in df.columns if '_burst' in col]
    for col in burst_cols:
        try:
            # Consistency is high when burstiness is low
            base_col = col.replace('_burst', '')
            df[f"{base_col}_consistency"] = 1 / (1 + df[col])
        except Exception:
            pass
    
    return df

def calculate_extreme_outliers(df):
    """
    Identify extreme outliers across all features
    
    Parameters:
    -----------
    df : DataFrame
        DataFrame with features
        
    Returns:
    --------
    df : DataFrame
        DataFrame with added extreme outlier flags
    """
    print("Identifying extreme outliers...")
    
    # Focus on all numeric columns except existing flags and categorical columns
    numeric_cols = df.select_dtypes(include=['number']).columns
    exclude_terms = ['outlier', 'flag', 'label', 'user']
    
    # Filter columns
    cols_to_check = [col for col in numeric_cols if not any(term in col for term in exclude_terms)]
    
    for col in cols_to_check:
        try:
            # Use 3 standard deviations as threshold for extreme outliers
            col_data = df[col].replace([np.inf, -np.inf], np.nan).dropna()
            if len(col_data) > 10:  # Need enough data for reliable stats
                mean = col_data.mean()
                std = col_data.std()
                
                if std > 0:
                    z_scores = (df[col] - mean) / std
                    df[f'extreme_outlier_{col}'] = (abs(z_scores) > 3.0).astype(int)
        except Exception as e:
            print(f"Error calculating outliers for {col}: {e}")
    
    return df

def calculate_risk_scores(df):
    """
    Calculate composite risk scores based on various features
    
    Parameters:
    -----------
    df : DataFrame
        DataFrame with features
        
    Returns:
    --------
    df : DataFrame
        DataFrame with added risk score columns
    """
    print("Calculating risk scores...")
    
    # Identify key risk indicators
    outlier_cols = [col for col in df.columns if 'extreme_outlier' in col]
    burst_cols = [col for col in df.columns if 'burst' in col]
    non_work_cols = [col for col in df.columns if 'non_working_hours' in col]
    weekend_cols = [col for col in df.columns if 'weekend_activity' in col]
    
    # Calculate data exfiltration risk
    exfil_indicators = [col for col in outlier_cols if any(term in col for term in 
                                                       ['file', 'email', 'http', 'device'])]
    if exfil_indicators:
        df['data_exfiltration_risk'] = df[exfil_indicators].sum(axis=1) / len(exfil_indicators)
    else:
        df['data_exfiltration_risk'] = 0
    
    # Calculate account compromise risk
    compromise_indicators = [col for col in outlier_cols if any(term in col for term in 
                                                           ['logon', 'activity', 'http'])]
    if compromise_indicators:
        df['account_compromise_risk'] = df[compromise_indicators].sum(axis=1) / len(compromise_indicators)
    else:
        df['account_compromise_risk'] = 0
    
    # Calculate burstiness indicators
    if burst_cols:
        # Max burstiness across domains
        file_burst_cols = [col for col in burst_cols if 'file' in col]
        email_burst_cols = [col for col in burst_cols if 'email' in col]
        http_burst_cols = [col for col in burst_cols if 'http' in col]
        
        # Flag high access breadth (activity across multiple domains)
        if file_burst_cols and email_burst_cols and http_burst_cols:
            file_burst = df[file_burst_cols].max(axis=1)
            email_burst = df[email_burst_cols].max(axis=1)
            http_burst = df[http_burst_cols].max(axis=1)
            
            # High access breadth
            df['high_access_breadth'] = ((file_burst > file_burst.quantile(0.9)) & 
                                         (email_burst > email_burst.quantile(0.9)) &
                                         (http_burst > http_burst.quantile(0.9))).astype(int)
            
            # Cross-domain exfiltration
            df['cross_domain_exfiltration'] = ((df['data_exfiltration_risk'] > df['data_exfiltration_risk'].quantile(0.8)) & 
                                               (df['high_access_breadth'] == 1)).astype(int)
            
            # Cross-domain reconnaissance
            df['cross_domain_reconnaissance'] = ((df['high_access_breadth'] == 1) & 
                                                 (df[non_work_cols].mean(axis=1) > 0.3)).astype(int)
    
    # Calculate Mahalanobis distance (multivariate outlier detection)
    try:
        # Select a subset of key features for Mahalanobis distance
        mahalanobis_features = []
        
        # Add some count and burst features (limiting to avoid high dimensionality issues)
        count_cols = [col for col in df.columns if col.endswith('_count')][:10]
        burst_cols_subset = burst_cols[:10] if len(burst_cols) > 10 else burst_cols
        
        mahalanobis_features.extend(count_cols)
        mahalanobis_features.extend(burst_cols_subset)
        
        if mahalanobis_features:
            # Calculate Mahalanobis distance
            X = df[mahalanobis_features].values
            mean = np.mean(X, axis=0)
            cov = np.cov(X, rowvar=False)
            
            # Handle singular covariance matrix
            try:
                inv_cov = np.linalg.inv(cov)
                
                # Calculate distance for each point
                mahalanobis_dist = []
                for i in range(X.shape[0]):
                    x = X[i]
                    dist = np.sqrt(np.dot(np.dot((x - mean), inv_cov), (x - mean).T))
                    mahalanobis_dist.append(dist)
                
                df['mahalanobis_distance'] = mahalanobis_dist
            except np.linalg.LinAlgError:
                # If matrix is singular, use pseudo-inverse
                inv_cov = np.linalg.pinv(cov)
                
                # Calculate distance for each point
                mahalanobis_dist = []
                for i in range(X.shape[0]):
                    x = X[i]
                    dist = np.sqrt(np.dot(np.dot((x - mean), inv_cov), (x - mean).T))
                    mahalanobis_dist.append(dist)
                
                df['mahalanobis_distance'] = mahalanobis_dist
    except Exception as e:
        print(f"Error calculating Mahalanobis distance: {e}")
        df['mahalanobis_distance'] = 0
    
    # Calculate multi-domain baseline deviation
    if 'activity_diversity' in df.columns:
        # Higher diversity + higher non-working-hours activity indicates deviation
        if non_work_cols:
            df['multi_domain_baseline_deviation'] = df['activity_diversity'] * df[non_work_cols].mean(axis=1)
        else:
            df['multi_domain_baseline_deviation'] = df['activity_diversity']
    
    # Calculate final composite risk score
    risk_factors = ['data_exfiltration_risk', 'account_compromise_risk', 'high_access_breadth', 
                   'mahalanobis_distance', 'cross_domain_exfiltration', 'cross_domain_reconnaissance',
                   'multi_domain_baseline_deviation']
    
    # Standardize risk factors
    standardized_risks = []
    for col in risk_factors:
        if col in df.columns:
            try:
                # Min-max scaling to [0,1]
                col_data = df[col].replace([np.inf, -np.inf], np.nan).fillna(0)
                col_min = col_data.min()
                col_max = col_data.max()
                
                if col_max > col_min:
                    standardized = (col_data - col_min) / (col_max - col_min)
                else:
                    standardized = col_data - col_min
                
                standardized_risks.append(standardized)
            except Exception:
                pass
    
    # Calculate composite score
    if standardized_risks:
        df['composite_risk_score'] = sum(standardized_risks) / len(standardized_risks)
    else:
        df['composite_risk_score'] = 0
    
    return df

def calculate_variability_features(df):
    """
    Calculate temporal and behavioral variability features
    
    Parameters:
    -----------
    df : DataFrame
        DataFrame with features
        
    Returns:
    --------
    df : DataFrame
        DataFrame with added variability features
    """
    print("Calculating variability features...")
    
    # Get basic count columns
    count_cols = [col for col in df.columns if col.endswith('_count') and not col.startswith('distinct_')]
    
    # Calculate variability for count columns
    for col in count_cols:
        try:
            if f"{col.replace('_count', '_std')}" in df.columns and f"{col.replace('_count', '_mean')}" in df.columns:
                std_col = f"{col.replace('_count', '_std')}"
                mean_col = f"{col.replace('_count', '_mean')}"
                
                # Variability = std / mean (coefficient of variation)
                df[f"variability_{col}"] = df[std_col] / df[mean_col].replace(0, 1)
        except Exception:
            pass
    
    # Calculate variability for z-score columns
    z_score_cols = [col for col in df.columns if col.startswith('z_score_')]
    
    for col in z_score_cols:
        try:
            df[f"variability_{col}"] = df[col].abs()
        except Exception:
            pass
    
    # Calculate variability for ratio to median columns 
    ratio_cols = [col for col in df.columns if col.startswith('ratio_to_median_')]
    
    for col in ratio_cols:
        try:
            # Variability = |log(ratio)| to capture deviations in both directions
            log_ratio = np.log1p(df[col].replace(0, 1))
            df[f"variability_{col}"] = log_ratio.abs()
        except Exception:
            pass
    
    return df

def create_file_logon_and_email_ratios(df):
    """
    Create file-to-logon and email sent-to-received ratios
    
    Parameters:
    -----------
    df : DataFrame
        DataFrame with features
        
    Returns:
    --------
    df : DataFrame
        DataFrame with added ratio features
    """
    print("Creating special ratio features...")
    
    # File-to-logon ratio
    if 'file_id_count' in df.columns and 'logon_id_count' in df.columns:
        try:
            ratio = df['file_id_count'] / df['logon_id_count'].replace(0, 1)
            
            # Calculate z-score of the ratio
            mean = ratio.mean()
            std = ratio.std()
            if std > 0:
                df['file_logon_ratio_z'] = (ratio - mean) / std
            else:
                df['file_logon_ratio_z'] = 0
        except Exception:
            df['file_logon_ratio_z'] = 0
    
    # Email sent-to-received ratio
    if 'email_from_count' in df.columns and 'email_to_count' in df.columns:
        try:
            # Emails sent (from) vs received (to)
            ratio = df['email_from_count'] / df['email_to_count'].replace(0, 1)
            
            # Calculate z-score of the ratio
            mean = ratio.mean()
            std = ratio.std()
            if std > 0:
                df['email_sent_received_ratio_z'] = (ratio - mean) / std
            else:
                df['email_sent_received_ratio_z'] = 0
        except Exception:
            df['email_sent_received_ratio_z'] = 0
    
    return df

def add_label_column(df, insiders_path='r4.2/answers/insiders.csv'):
    """
    Add insider threat labels if available
    
    Parameters:
    -----------
    df : DataFrame
        DataFrame to add labels to
    insiders_path : str
        Path to insiders.csv file
        
    Returns:
    --------
    df : DataFrame
        DataFrame with added label column
    """
    print("Adding label column...")
    
    try:
        if os.path.exists(insiders_path):
            # Load insiders data
            insiders_df = pd.read_csv(insiders_path)
            
            # Get list of insider users
            insider_users = insiders_df['user'].unique().tolist()
            
            # Create label column
            df['label'] = df['user'].apply(lambda x: 1 if str(x) in map(str, insider_users) else 0)
            
            print(f"Added labels: {df['label'].sum()} insiders identified out of {len(df)} users")
        else:
            # No insiders file, use dummy labels
            df['label'] = 0
            print("Insiders file not found. Added dummy label column.")
    except Exception as e:
        print(f"Error adding labels: {e}")
        df['label'] = 0
    
    return df

def main():
    """
    Main function to generate all features in combined_features.csv
    """
    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Process CERT r4.2 dataset and generate features for insider threat detection.')
    parser.add_argument('--full-dataset', action='store_true', help='Process the entire dataset instead of just a sample')
    parser.add_argument('--output', default='combined_features.csv', help='Output CSV file name')
    parser.add_argument('--chunk-size', type=int, default=CHUNK_SIZE, help='Chunk size for processing large files')
    args = parser.parse_args()
    
    # Use local variables rather than modifying globals
    chunk_size = args.chunk_size
    
    start_time = time.time()
    print("Starting comprehensive feature generation...")
    
    if args.full_dataset:
        print("*** PROCESSING FULL DATASET (THIS MAY TAKE A LONG TIME AND REQUIRE SIGNIFICANT MEMORY) ***")
        sample_size = None
    else:
        print(f"Processing sample of {SAMPLE_SIZE} rows per dataset")
        sample_size = SAMPLE_SIZE
    
    # Define the dataset files to use
    dataset_files = {
        'device': 'r4.2/device.csv',
        'email': 'r4.2/email.csv',
        'file': 'r4.2/file.csv', 
        'http': 'r4.2/http.csv',
        'logon': 'r4.2/logon.csv'
    }
    
    # Check which files exist
    existing_files = {}
    for dataset_type, file_path in dataset_files.items():
        if os.path.exists(file_path):
            existing_files[dataset_type] = file_path
        else:
            print(f"Warning: File {file_path} not found and will be skipped")
    
    if not existing_files:
        print("Error: No dataset files found. Please check the file paths.")
        return
    
    # Step 1: Load and combine datasets with basic features
    # Pass both sample_size and chunk_size explicitly
    df = combine_datasets(existing_files, sample_size=sample_size, chunk_size=chunk_size)
    print(f"Combined dataset shape after initial processing: {df.shape}")
    
    # Step 2: Calculate z-scores
    df = calculate_z_scores(df)
    
    # Step 3: Calculate ratio to median
    df = calculate_ratio_to_median(df)
    
    # Step 4: Calculate burstiness and consistency metrics
    df = calculate_burstiness_features(df)
    
    # Step 5: Identify extreme outliers
    df = calculate_extreme_outliers(df)
    
    # Step 6: Calculate variability features
    df = calculate_variability_features(df)
    
    # Step 7: Create file/logon and email sent/received ratios
    df = create_file_logon_and_email_ratios(df)
    
    # Step 8: Calculate risk scores and composite indicators
    df = calculate_risk_scores(df)
    
    # Step 9: Add label column if insiders.csv exists
    df = add_label_column(df)
    
    # Save the final combined dataset
    print(f"\nSaving the combined dataset with all features to {args.output}...")
    df.to_csv(args.output, index=False)
    
    # Print dataset statistics
    print(f"\nFinal Dataset Statistics:")
    print(f"Total users/rows: {df.shape[0]}")
    print(f"Total features: {df.shape[1]}")
    
    if 'label' in df.columns:
        insider_rate = df['label'].mean()
        print(f"Insider rate: {insider_rate:.4f} ({int(insider_rate * df.shape[0])} users)")
    
    # Execution time
    execution_time = time.time() - start_time
    execution_minutes = execution_time / 60
    print(f"\nFeature generation completed in {execution_time:.2f} seconds ({execution_minutes:.2f} minutes)")
    print(f"Combined features with all columns saved to {args.output}")

if __name__ == "__main__":
    main()