data_loader.py

"""
Data Loader Module
Loads and preprocesses the customer churn dataset for Globomantics.
"""
import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler


class DataLoader:
    """Loads and preprocesses customer churn data for model training."""
    
    def __init__(self):
        """Initialize data loader with configuration parameters."""
        self.test_size = 0.2
        self.random_state = 42
        
        self.scaler = StandardScaler()
        self.X_train = None
        self.X_test = None
        self.y_train = None
        self.y_test = None
    
    def load_data(self):
        """Generate synthetic customer churn dataset."""
        X, y = make_classification(
            n_samples=1000,
            n_features=20,
            n_informative=15,
            n_redundant=3,
            n_repeated=0,
            n_clusters_per_class=3,
            n_classes=2,
            weights=[0.6, 0.4],
            class_sep=1.0,
            flip_y=0.03,
            random_state=self.random_state
        )
        
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(
            X, y,
            test_size=self.test_size,
            random_state=self.random_state,
            stratify=y
        )
        
        self.X_train = self.scaler.fit_transform(self.X_train)
        self.X_test = self.scaler.transform(self.X_test)
        
        return self.X_train, self.X_test, self.y_train, self.y_test
    
    def get_data_info(self):
        """Print information about the loaded dataset."""
        if self.X_train is None:
            print("Data not loaded yet. Call load_data() first.")
            return
        
        print("=" * 50)
        print("Dataset Information")
        print("=" * 50)
        print(f"Training samples: {len(self.X_train)}")
        print(f"Test samples: {len(self.X_test)}")
        print(f"Number of features: {self.X_train.shape[1]}")
        print(f"Class distribution (train): {np.bincount(self.y_train)}")
        print("=" * 50)


if __name__ == "__main__":
    loader = DataLoader()
    X_train, X_test, y_train, y_test = loader.load_data()
    loader.get_data_info()