import sys
import os

# Add project root directory to Python path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..', '..')))

from myffe.tools.Loggerv1_1 import log_class
from myffe.tools.color_printer import ColorPrinter
from myffe.tools.evaluation_decorator import with_evaluation

@log_class(type='use')
class Normalizer:
    def __init__(self, shortcut_norm=None):
        self.shortcut_norm = shortcut_norm
    
    @with_evaluation
    def __call__(self, data):
        if self.shortcut_norm is None:
            ColorPrinter.print_error('Error: Missing required parameter shortcut_norm')
            return data
        return self.MY_Normalization(data, norm_params=self.shortcut_norm)
    
    def MY_Normalization(self, data, norm_params=None):
        import pandas as pd
        import numpy as np
        
        if 'normalizer_method_choice' not in norm_params:
            ColorPrinter.print_error('Error: Missing required parameter normalizer_method_choice')
            return data
        
        normalizer_method_choice = norm_params['normalizer_method_choice']
        
        # Get columns to normalize
        normalized_col = norm_params.get('normalized_col')             
        
        # Check data size
        data_size = data.memory_usage(deep=True).sum() / (1024 * 1024)  # Convert to MB
        ColorPrinter.print_info(f'Data size: {data_size:.2f} MB')
        
        # Normalization method mapping
        normalizer_config = {
            'L1 Normalization': {
                'params': {'norm': 'l1'},
                'message': 'Using L1 normalization'
            },
            'L2 Normalization': {
                'params': {'norm': 'l2'},
                'message': 'Using L2 normalization'
            },
            'Max Normalization': {
                'params': {'norm': 'max'},
                'message': 'Using Max normalization'
            },
            'Robust Normalization': {
                'params': {'norm': 'l2'},
                'message': 'Using Robust normalization'
            },
            'Power Normalization': {
                'params': {},
                'message': 'Using Power normalization'
            },
            'Log Normalization': {
                'params': {},
                'message': 'Using Log normalization'
            },
            'Square Root Normalization': {
                'params': {},
                'message': 'Using Square Root normalization'
            }
        }
        
        # Apply different normalization methods to different columns
        ColorPrinter.print_info('Applying different normalization methods to different columns')
        
        try:
            from sklearn.preprocessing import Normalizer
            from sklearn.preprocessing import RobustScaler
            
            # Create data copy
            normalized_data = data.copy()
            
            # Skip processing if no columns specified
            if not normalized_col:
                ColorPrinter.print_info('No normalization columns specified, skipping normalization')
                return data
            
            # Check if lengths match
            if len(normalized_col) != len(normalizer_method_choice):
                ColorPrinter.print_error(f'Error: Column count ({len(normalized_col)}) does not match method count ({len(normalizer_method_choice)})')
                return data
            
            # Apply corresponding normalization method to each column
            for col, method in zip(normalized_col, normalizer_method_choice):
                if col in data.columns and method in normalizer_config:
                    config = normalizer_config[method]
                    ColorPrinter.print_info(f'Applying {method} to column {col}')
                    
                    if method == 'Robust Normalization':
                        # Use RobustScaler for normalization
                        scaler = RobustScaler()
                        # Normalize single column
                        normalized_data[[col]] = pd.DataFrame(
                            scaler.fit_transform(normalized_data[[col]]),
                            columns=[col],
                            index=normalized_data.index
                        )
                    elif method == 'Power Normalization':
                        # Use power normalization
                        # Using square as default power
                        normalized_data[col] = normalized_data[col] ** 2
                        # Then perform L2 normalization
                        normalizer = Normalizer(norm='l2')
                        normalized_data[[col]] = pd.DataFrame(
                            normalizer.fit_transform(normalized_data[[col]]),
                            columns=[col],
                            index=normalized_data.index
                        )
                    elif method == 'Log Normalization':
                        # Use log normalization
                        # Ensure values are positive
                        min_val = normalized_data[col].min()
                        if min_val <= 0:
                            normalized_data[col] = normalized_data[col] - min_val + 0.001
                        normalized_data[col] = np.log10(normalized_data[col])
                    elif method == 'Square Root Normalization':
                        # Use square root normalization
                        # Ensure values are non-negative
                        normalized_data[col] = normalized_data[col].clip(lower=0)
                        normalized_data[col] = np.sqrt(normalized_data[col])
                    else:
                        # Use Normalizer for normalization
                        normalizer = Normalizer(**config['params'])
                        # Normalize single column
                        normalized_data[[col]] = pd.DataFrame(
                            normalizer.fit_transform(normalized_data[[col]]),
                            columns=[col],
                            index=normalized_data.index
                        )
                else:
                    ColorPrinter.print_error(f'Column {col} does not exist or normalization method {method} is invalid')
            
            return normalized_data
        except Exception as e:
            ColorPrinter.print_error(f'Error during normalization: {e}')
            return data


"""
Parameter documentation:

shortcut_norm parameter dictionary:
    - normalizer_method_choice: Normalization method choice
      Type: list
      Options: ['L1 Normalization', 'L2 Normalization', 'Max Normalization', 'Robust Normalization', 'Power Normalization', 'Log Normalization', 'Square Root Normalization']
      Default: None (must be specified)
      Description: Specify normalization method for each column, list length must match normalized_col
        - 'L1 Normalization': Normalize using L1 norm, each sample's feature values divided by L1 norm
        - 'L2 Normalization': Normalize using L2 norm, each sample's feature values divided by L2 norm (most commonly used)
        - 'Max Normalization': Normalize using maximum value, each sample's feature values divided by maximum value
        - 'Robust Normalization': Normalize using RobustScaler, based on median and quartiles, not sensitive to outliers
        - 'Power Normalization': Use square as default power, then perform L2 normalization, suitable for data with large skewness
        - 'Log Normalization': Use log transformation, suitable for positive data with large value range
        - 'Square Root Normalization': Use square root transformation, suitable for non-negative data
    - normalized_col: Columns to normalize
      Type: list
      Options: Column names in the dataset
      Default: None (must be specified)
      Description: Specify columns to normalize, list length must match normalizer_method_choice

Usage example:
    # Apply different normalization methods to different columns
    norm_params = {
        'normalizer_method_choice': ['L1 Normalization', 'L2 Normalization', 'Max Normalization'],
        'normalized_col': ['feature1', 'feature2', 'feature3']
    }
    normalizer = Normalizer(shortcut_norm=norm_params)
    normalized_data = normalizer(data)
"""


if __name__ == "__main__":
    import pandas as pd
    import numpy as np
    
    ColorPrinter.print_any("=" * 60)
    ColorPrinter.print_info("Normalization Processing Test Program")
    ColorPrinter.print_any("=" * 60)
    
    # Create test data
    np.random.seed(42)
    test_data = pd.DataFrame({
        'feature1': np.random.randint(1, 100, 10),
        'feature2': np.random.randint(1, 100, 10),
        'feature3': np.random.randint(1, 100, 10)
    })
    
    ColorPrinter.print_any("\nOriginal data:")
    ColorPrinter.print_any(test_data)
    ColorPrinter.print_any("\nData statistics:")
    ColorPrinter.print_any(test_data.describe())
    
    # Test applying different normalization methods to different columns
    ColorPrinter.print_any(f"\n{'=' * 60}")
    ColorPrinter.print_info("Test: Applying different normalization methods to different columns")
    ColorPrinter.print_any('=' * 60)
    
    norm_params = {
        'normalizer_method_choice': ['L1 Normalization', 'L2 Normalization', 'Max Normalization'],
        'normalized_col': ['feature1', 'feature2', 'feature3']
    }
    
    normalizer = Normalizer(shortcut_norm=norm_params)
    normalized_data = normalizer(test_data.copy())
    
    ColorPrinter.print_any("\nNormalized data:")
    ColorPrinter.print_any(normalized_data)
    ColorPrinter.print_any("\nNormalized data statistics:")
    ColorPrinter.print_any(normalized_data.describe())
    
    ColorPrinter.print_any("\n" + "=" * 60)
    ColorPrinter.print_success("Test completed!")
    ColorPrinter.print_any("=" * 60)
