preprocess.py

"""
Módulo de preprocesamiento de datos para el modelo
Optimizado para trabajar con Ray Train
"""

import pandas as pd
import numpy as np
import pickle
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import train_test_split
import torch
from torch.utils.data import DataLoader, TensorDataset


class DataPreprocessor:
    """
    Clase para manejar el preprocesamiento de datos de ventas
    """
    
    def __init__(self):
        self.encoders = {}
        self.scaler = StandardScaler()
        self.feature_columns = None
        self.categorical_columns = ['Region', 'Country', 'Item Type', 'Sales Channel', 'Order Priority']
    
    def load_and_preprocess(self, file_path):
        """
        Carga y preprocesa los datos de ventas para Ray Train
        
        Args:
            file_path: ruta al archivo CSV
            
        Returns:
            X_scaled: características normalizadas (numpy array)
            y: valores target (numpy array)
            preprocessing_info: diccionario con información de preprocesamiento
        """
        print("Cargando datos para Ray Train...")
        df = pd.read_csv(file_path)
        
        print(f"Dataset shape: {df.shape}")
        print(f"Columnas: {df.columns.tolist()}")
        
        # Convertir fechas a características numéricas
        df['Order Date'] = pd.to_datetime(df['Order Date'])
        df['Ship Date'] = pd.to_datetime(df['Ship Date'])
        
        # Extraer características de las fechas
        df['Order_Year'] = df['Order Date'].dt.year
        df['Order_Month'] = df['Order Date'].dt.month
        df['Order_Day'] = df['Order Date'].dt.day
        df['Order_DayOfWeek'] = df['Order Date'].dt.dayofweek
        
        # Días entre orden y envío
        df['Days_to_Ship'] = (df['Ship Date'] - df['Order Date']).dt.days
        
        # Eliminar columnas originales de fecha y Order ID
        df = df.drop(['Order Date', 'Ship Date', 'Order ID'], axis=1)
        
        # Separar características y target
        target_column = 'Total Profit'
        self.feature_columns = [col for col in df.columns if col != target_column]
        
        X = df[self.feature_columns]
        y = df[target_column]
        
        # Codificar variables categóricas
        for col in self.categorical_columns:
            if col in X.columns:
                le = LabelEncoder()
                X.loc[:, col] = le.fit_transform(X[col])
                self.encoders[col] = le
        
        # Normalizar características numéricas
        X_scaled = self.scaler.fit_transform(X)
        
        # Crear diccionario con información de preprocesamiento
        preprocessing_info = {
            'encoders': self.encoders,
            'scaler': self.scaler,
            'feature_columns': self.feature_columns,
            'categorical_columns': self.categorical_columns
        }
        
        print("Preprocesamiento completado para Ray Train")
        # Retornar como numpy arrays para compatibilidad con Ray Train
        return X_scaled, y.values, preprocessing_info
    
    def save_preprocessing_info(self, filepath='preprocessing_info.pkl'):
        """
        Guarda la información de preprocesamiento
        
        Args:
            filepath: ruta donde guardar el archivo
        """
        preprocessing_info = {
            'encoders': self.encoders,
            'scaler': self.scaler,
            'feature_columns': self.feature_columns,
            'categorical_columns': self.categorical_columns
        }
        
        with open(filepath, 'wb') as f:
            pickle.dump(preprocessing_info, f)
        
        print(f"Información de preprocesamiento guardada en {filepath}")
    
    def load_preprocessing_info(self, filepath='preprocessing_info.pkl'):
        """
        Carga la información de preprocesamiento
        
        Args:
            filepath: ruta del archivo a cargar
        """
        with open(filepath, 'rb') as f:
            preprocessing_info = pickle.load(f)
        
        self.encoders = preprocessing_info['encoders']
        self.scaler = preprocessing_info['scaler']
        self.feature_columns = preprocessing_info['feature_columns']
        self.categorical_columns = preprocessing_info['categorical_columns']
        
        return preprocessing_info
    
    def preprocess_single_sample(self, sale_data):
        """
        Preprocesa una sola muestra para predicción
        
        Args:
            sale_data: diccionario con los datos de la venta
            
        Returns:
            numpy array con características preprocesadas
        """
        # Crear DataFrame con una sola fila
        df = pd.DataFrame([sale_data])
        
        # Convertir fechas si están presentes
        if 'Order Date' in df.columns:
            df['Order Date'] = pd.to_datetime(df['Order Date'])
            df['Order_Year'] = df['Order Date'].dt.year
            df['Order_Month'] = df['Order Date'].dt.month
            df['Order_Day'] = df['Order Date'].dt.day
            df['Order_DayOfWeek'] = df['Order Date'].dt.dayofweek
        
        if 'Ship Date' in df.columns:
            df['Ship Date'] = pd.to_datetime(df['Ship Date'])
            if 'Order Date' in df.columns:
                df['Days_to_Ship'] = (df['Ship Date'] - df['Order Date']).dt.days
        
        # Eliminar columnas no necesarias
        columns_to_drop = ['Order Date', 'Ship Date', 'Order ID', 'Total Profit']
        df = df.drop([col for col in columns_to_drop if col in df.columns], axis=1)
        
        # Asegurar que tengamos todas las columnas necesarias
        for col in self.feature_columns:
            if col not in df.columns:
                df[col] = 0  # Valor por defecto
        
        # Reordenar columnas
        df = df[self.feature_columns]
        
        # Codificar variables categóricas
        for col, encoder in self.encoders.items():
            if col in df.columns:
                try:
                    df[col] = encoder.transform(df[col])
                except:
                    # Si el valor no está en el encoder, usar el valor más común
                    df[col] = 0
        
        # Normalizar características
        X_scaled = self.scaler.transform(df)
        
        return X_scaled


def create_data_loaders(X, y, batch_size=64, val_split=0.2, test_split=0.1):
    """
    Crea los data loaders para entrenamiento, validación y prueba
    Compatible con Ray Train para procesamiento distribuido
    
    Args:
        X: características (numpy array o tensor)
        y: targets (numpy array o tensor)
        batch_size: tamaño del batch
        val_split: proporción para validación
        test_split: proporción para prueba
        
    Returns:
        train_loader, val_loader, test_loader
    """
    # Dividir en conjunto de entrenamiento, validación y prueba
    X_temp, X_test, y_temp, y_test = train_test_split(
        X, y, test_size=test_split, random_state=42
    )
    
    X_train, X_val, y_train, y_val = train_test_split(
        X_temp, y_temp, test_size=val_split/(1-test_split), random_state=42
    )
    
    # Convertir a tensores de PyTorch para Ray Train
    # Ray Train maneja mejor los FloatTensors para distribución
    X_train_tensor = torch.FloatTensor(X_train)
    y_train_tensor = torch.FloatTensor(y_train).reshape(-1, 1)
    X_val_tensor = torch.FloatTensor(X_val)
    y_val_tensor = torch.FloatTensor(y_val).reshape(-1, 1)
    X_test_tensor = torch.FloatTensor(X_test)
    y_test_tensor = torch.FloatTensor(y_test).reshape(-1, 1)
    
    # Crear datasets
    train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
    val_dataset = TensorDataset(X_val_tensor, y_val_tensor)
    test_dataset = TensorDataset(X_test_tensor, y_test_tensor)
    
    # Crear data loaders - Ray Train los preparará para distribución
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)
    test_loader = DataLoader(test_dataset, batch_size=batch_size)
    
    print(f"DataLoaders creados para Ray Train:")
    print(f"  • Conjunto de entrenamiento: {len(train_dataset)} muestras")
    print(f"  • Conjunto de validación: {len(val_dataset)} muestras")
    print(f"  • Conjunto de prueba: {len(test_dataset)} muestras")
    
    return train_loader, val_loader, test_loader