"""
Customer Churn Dataset - Data Preprocessing Script
==================================================
Provided by: RSK World
Website: https://rskworld.in/
Email: help@rskworld.in
Phone: +91 93305 39277
Contact Page: https://rskworld.in/contact.php
"""

import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import train_test_split
from imblearn.over_sampling import SMOTE

def load_data(file_path='../data/customer_churn.csv'):
    """
    Load the customer churn dataset
    
    Args:
        file_path (str): Path to the CSV file
        
    Returns:
        pd.DataFrame: Loaded dataset
    """
    df = pd.read_csv(file_path, comment="#", na_values=["", "NA", "NaN"])
    return df

def handle_missing_values(df):
    """
    Handle missing values in the dataset
    
    Args:
        df (pd.DataFrame): Dataset to preprocess
        
    Returns:
        pd.DataFrame: Dataset with handled missing values
    """
    print("Handling missing values...")
    
    # Fill missing values in numerical columns with median
    numerical_cols = df.select_dtypes(include=[np.number]).columns
    for col in numerical_cols:
        if df[col].isnull().sum() > 0:
            df[col].fillna(df[col].median(), inplace=True)
    
    # Fill missing values in categorical columns with mode
    categorical_cols = df.select_dtypes(include=['object', 'category']).columns
    for col in categorical_cols:
        if df[col].isnull().sum() > 0:
            df[col].fillna(df[col].mode()[0], inplace=True)
    
    return df

def encode_categorical_variables(df):
    """
    Encode categorical variables to numerical format
    
    Args:
        df (pd.DataFrame): Dataset to preprocess
        
    Returns:
        pd.DataFrame: Dataset with encoded categorical variables
        dict: Mapping of encoded labels
    """
    print("Encoding categorical variables...")
    
    df_encoded = df.copy()
    label_encoders = {}
    
    # Columns to encode
    categorical_cols = ['Gender', 'City', 'AccountType', 'PaymentMethod', 'Churn', 'AgeGroup', 'TenureBucket']
    
    for col in categorical_cols:
        if col in df_encoded.columns:
            le = LabelEncoder()
            df_encoded[col] = le.fit_transform(df_encoded[col].fillna("Unknown"))
            label_encoders[col] = le
    
    return df_encoded, label_encoders

def create_features(df):
    """
    Create additional features from existing ones
    
    Args:
        df (pd.DataFrame): Dataset to preprocess
        
    Returns:
        pd.DataFrame: Dataset with new features
    """
    print("Creating new features...")
    
    df_features = df.copy()
    
    # Temporal signals from login activity
    if 'LastLogin' in df_features.columns:
        df_features['LastLogin'] = pd.to_datetime(df_features['LastLogin'], errors='coerce')
        df_features['DaysSinceLastLogin'] = (pd.Timestamp.now() - df_features['LastLogin']).dt.days
        df_features['LastLoginDayOfWeek'] = df_features['LastLogin'].dt.dayofweek
        df_features['LastLoginMonth'] = df_features['LastLogin'].dt.month
    
    # Usage and billing relationships
    if 'MonthlyUsage' in df_features.columns and 'BillingAmount' in df_features.columns:
        denom = df_features['BillingAmount'].replace(0, np.nan).fillna(df_features['BillingAmount'].median())
        df_features['UsagePerDollar'] = df_features['MonthlyUsage'] / (denom + 1e-6)
    
    if 'SupportCalls' in df_features.columns and 'Tenure' in df_features.columns:
        df_features['SupportCallsPerMonth'] = df_features['SupportCalls'] / df_features['Tenure'].replace(0, 1)
    
    if 'MonthlyUsage' in df_features.columns and 'Tenure' in df_features.columns:
        df_features['UsagePerTenureMonth'] = df_features['MonthlyUsage'] / df_features['Tenure'].replace(0, 1)
    
    if 'BillingAmount' in df_features.columns and 'ContractLength' in df_features.columns:
        df_features['BillingPerContractMonth'] = df_features['BillingAmount'] / df_features['ContractLength'].replace(0, 1)
        df_features['IsLongContract'] = (df_features['ContractLength'] >= 24).astype(int)
    
    if 'BillingAmount' in df_features.columns and 'Tenure' in df_features.columns:
        df_features['LifetimeValueEstimate'] = df_features['BillingAmount'] * df_features['Tenure'].clip(lower=1)
    
    if {'MonthlyUsage', 'Tenure', 'SupportCalls'}.issubset(df_features.columns):
        df_features['EngagementScore'] = (
            df_features['MonthlyUsage'].rank(pct=True) * 0.5 +
            df_features['Tenure'].rank(pct=True) * 0.3 -
            df_features['SupportCalls'].rank(pct=True) * 0.2
        )
    
    # Age groups and tenure buckets for non-linear patterns
    if 'Age' in df_features.columns:
        df_features['AgeGroup'] = pd.cut(
            df_features['Age'],
            bins=[0, 30, 40, 50, 60, 120],
            labels=['<30', '30-39', '40-49', '50-59', '60+']
        )
    if 'Tenure' in df_features.columns:
        df_features['TenureBucket'] = pd.cut(
            df_features['Tenure'],
            bins=[-1, 6, 12, 24, 48, 120],
            labels=['<=6m', '6-12m', '12-24m', '24-48m', '48m+']
        )
    
    return df_features

def scale_features(df, feature_cols):
    """
    Scale numerical features
    
    Args:
        df (pd.DataFrame): Dataset to preprocess
        feature_cols (list): List of feature columns to scale
        
    Returns:
        pd.DataFrame: Dataset with scaled features
        StandardScaler: Fitted scaler object
    """
    print("Scaling features...")
    
    scaler = StandardScaler()
    df_scaled = df.copy()
    df_scaled[feature_cols] = scaler.fit_transform(df[feature_cols])
    
    return df_scaled, scaler

def prepare_train_test_split(df, target_col='Churn', test_size=0.2, random_state=42):
    """
    Prepare train and test splits
    
    Args:
        df (pd.DataFrame): Preprocessed dataset
        target_col (str): Name of target column
        test_size (float): Proportion of test set
        random_state (int): Random seed
        
    Returns:
        tuple: X_train, X_test, y_train, y_test
    """
    print("Preparing train/test split...")
    
    # Exclude non-feature columns
    exclude_cols = ['CustomerID', 'LastLogin', 'ChurnDate', target_col]
    feature_cols = [col for col in df.columns if col not in exclude_cols]
    
    X = df[feature_cols]
    y = df[target_col]
    
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=random_state, stratify=y
    )
    
    print(f"Training set size: {X_train.shape}")
    print(f"Test set size: {X_test.shape}")
    
    return X_train, X_test, y_train, y_test

def balance_training_data(X_train, y_train):
    """
    Balance the training data using SMOTE to handle class imbalance.
    
    Args:
        X_train (pd.DataFrame): Training feature set
        y_train (pd.Series): Training labels
        
    Returns:
        tuple: Balanced X_train, y_train
    """
    print("Balancing classes with SMOTE...")
    smote = SMOTE(random_state=42)
    X_resampled, y_resampled = smote.fit_resample(X_train, y_train)
    X_balanced = pd.DataFrame(X_resampled, columns=X_train.columns)
    y_balanced = pd.Series(y_resampled, name=y_train.name)
    print(f"Balanced training set size: {X_balanced.shape}")
    return X_balanced, y_balanced

def main():
    """
    Main function to run preprocessing pipeline
    """
    print("Customer Churn Dataset - Data Preprocessing")
    print("Provided by: RSK World (https://rskworld.in/)")
    print("-" * 60)
    
    # Load data
    df = load_data()
    print(f"\nOriginal dataset shape: {df.shape}")
    
    # Handle missing values
    df = handle_missing_values(df)
    
    # Create new features before encoding to capture richer signals
    df_features = create_features(df)
    df_features = handle_missing_values(df_features)
    
    # Encode categorical variables (including engineered bands)
    df_encoded, label_encoders = encode_categorical_variables(df_features)
    
    # Remove raw datetime columns to keep the dataset numeric and leakage-free
    df_encoded = df_encoded.drop(columns=['LastLogin', 'ChurnDate'], errors='ignore')
    
    # Determine which columns to scale (exclude identifiers and target)
    exclude_cols = ['CustomerID', 'Churn']
    feature_cols = [col for col in df_encoded.columns if col not in exclude_cols]
    
    # Scale features for model-friendly distributions
    df_scaled, scaler = scale_features(df_encoded, feature_cols)
    
    print(f"\nPreprocessed dataset shape: {df_scaled.shape}")
    print(f"\nColumns: {list(df_scaled.columns)}")
    
    # Prepare train/test split using scaled data
    X_train, X_test, y_train, y_test = prepare_train_test_split(df_scaled)
    
    # Balance the training data to mitigate churn class imbalance
    X_train_balanced, y_train_balanced = balance_training_data(X_train, y_train)
    
    # Save preprocessed data
    df_scaled.to_csv('../data/customer_churn_preprocessed.csv', index=False)
    print("\nPreprocessed data saved to '../data/customer_churn_preprocessed.csv'")
    
    print("\n" + "=" * 60)
    print("PREPROCESSING COMPLETE")
    print("=" * 60)
    print("\nFor more information, visit: https://rskworld.in/")
    print("Contact: help@rskworld.in | +91 93305 39277")
    
    return X_train_balanced, X_test, y_train_balanced, y_test

if __name__ == "__main__":
    X_train, X_test, y_train, y_test = main()