"""
Autoencoder Preprocessing

Implements autoencoder-based missing value imputation
"""

import numpy as np
import pandas as pd
from myffe.advance_preprocessing.preprocessing_all.autoencoder.autoencoder import AutoEncoder
from myffe.tools.color_printer import ColorPrinter


class AutoencoderPreprocessing:
    """
    Autoencoder Preprocessing for missing value imputation
    """
    
    def __init__(self, shortcut_autoencoder):
        """
        Initialize AutoencoderPreprocessing
        
        Args:
            shortcut_autoencoder (dict): Autoencoder configuration
        """
        self.shortcut_autoencoder = shortcut_autoencoder
        self.autoencoder = None
        
        if shortcut_autoencoder == 'skip':
            ColorPrinter.print_info("Autoencoder preprocessing skipped")
        else:
            ColorPrinter.print_info("Autoencoder preprocessing initialized")
            ColorPrinter.print_info(f"Autoencoder configuration: {shortcut_autoencoder}")
    
    def __call__(self, data, data_label=None):
        """
        Apply autoencoder preprocessing to data
        
        Args:
            data (DataFrame): Input data with missing values
            data_label (str): Label column name (if any)
            
        Returns:
            DataFrame: Data with missing values imputed
        """
        if self.shortcut_autoencoder == 'skip':
            ColorPrinter.print_info("Skipping autoencoder preprocessing")
            return data
        
        try:
            # Separate features and label if label is provided
            if data_label and data_label in data.columns:
                label_data = data[[data_label]].copy()
                feature_data = data.drop(columns=[data_label])
            else:
                label_data = None
                feature_data = data.copy()
            
            # Convert DataFrame to numpy array
            feature_np = feature_data.to_numpy()
            
            # Create autoencoder instance
            self.autoencoder = AutoEncoder(self.shortcut_autoencoder)
            
            # Apply autoencoder
            imputed_feature_np = self.autoencoder(feature_np)
            
            # Convert back to DataFrame
            imputed_feature_data = pd.DataFrame(imputed_feature_np, columns=feature_data.columns, index=feature_data.index)
            
            # Combine with label if it exists
            if label_data is not None:
                imputed_data = pd.concat([imputed_feature_data, label_data], axis=1)
            else:
                imputed_data = imputed_feature_data
            
            # Preserve original data types
            for col in data.columns:
                if data[col].dtype != imputed_data[col].dtype:
                    try:
                        imputed_data[col] = imputed_data[col].astype(data[col].dtype)
                    except:
                        pass
            
            ColorPrinter.print_success("OK Autoencoder preprocessing completed")
            return imputed_data
        except Exception as e:
            ColorPrinter.print_error(f"Error in autoencoder preprocessing: {str(e)}")
            import traceback
            traceback.print_exc()
            return data
