# Author: Chen Xingqiang
# SPDX-License-Identifier: BSD-3-Clause

"""
Federated Learning adapter for RobustScaler

RobustScaler is an UNSUPERVISED preprocessing algorithm.
Data remains in local PYUs, JAX-accelerated local computation,
HEU-based secure aggregation of statistics.

Mode: Federated Learning (FL)
Generated by: StandaloneAlgorithmMigrator
"""

import logging
from typing import Dict, Union, Optional

import numpy as np

try:
    from xlearn.preprocessing import RobustScaler
    USING_XLEARN = True
except ImportError:
    from sklearn.preprocessing import RobustScaler
    USING_XLEARN = False

try:
    from secretflow.data.ndarray.ndarray import FedNdarray
    from secretflow.data.vertical.dataframe import VDataFrame
    from secretflow.device import PYU, HEU
    from secretflow.device.device.pyu import PYUObject
    from secretflow.security.aggregation import SecureAggregator
    SECRETFLOW_AVAILABLE = True
except ImportError:
    SECRETFLOW_AVAILABLE = False


class FLRobustScaler:
    """
    Federated Learning RobustScaler
    
    RobustScaler is an unsupervised preprocessing algorithm.
    
    In FL mode:
    - Each party computes local statistics on their data partition
    - Statistics are securely aggregated via HEU encryption
    - No labels (y) are needed - this is unsupervised preprocessing
    
    Parameters
    ----------
    devices : Dict[str, PYU]
        Dictionary mapping party names to PYU devices
    heu : HEU, optional
        Homomorphic encryption unit for secure aggregation
    aggregation_method : str, default='mean'
        How to aggregate statistics: 'mean', 'weighted_mean'
    **kwargs
        Parameters passed to RobustScaler
    
    Examples
    --------
    >>> import secretflow as sf
    >>> alice = sf.PYU('alice')
    >>> bob = sf.PYU('bob')
    >>> heu = sf.HEU(sf.HEUConfig(...), ...)
    >>> 
    >>> # Unsupervised preprocessing - no labels needed
    >>> scaler = FLRobustScaler(
    >>>     devices={'alice': alice, 'bob': bob},
    >>>     heu=heu
    >>> )
    >>> 
    >>> # Fit on federated data (no y labels)
    >>> scaler.fit(fed_X)
    >>> 
    >>> # Transform data
    >>> X_scaled = scaler.transform(fed_X_test)
    >>> 
    >>> # Or fit and transform in one step
    >>> X_scaled = scaler.fit_transform(fed_X)
    """
    
    def __init__(
        self,
        devices: Dict[str, 'PYU'],
        heu: Optional['HEU'] = None,
        aggregation_method: str = 'mean',
        **kwargs
    ):
        if not SECRETFLOW_AVAILABLE:
            raise RuntimeError("SecretFlow not installed. Install: pip install secretflow")
        
        self.devices = devices
        self.heu = heu
        self.aggregation_method = aggregation_method
        self.kwargs = kwargs
        
        # Create local preprocessors on each PYU
        self.local_models = {}
        for party_name, device in devices.items():
            self.local_models[party_name] = device(self._create_local_model)(**kwargs)
        
        # Track if models are fitted
        self._is_fitted = False
        
        if USING_XLEARN:
            logging.info("[FL] FLRobustScaler initialized with JAX acceleration")
        else:
            logging.info("[FL] FLRobustScaler initialized with sklearn")
        
        logging.info(f"[FL] Parties: {list(devices.keys())}")
        logging.info(f"[FL] Aggregation: {aggregation_method}")
        logging.info(f"[FL] HEU enabled: {heu is not None}")
    
    @staticmethod
    def _create_local_model(**kwargs):
        """Create local RobustScaler instance"""
        return RobustScaler(**kwargs)
    
    def fit(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Fit the federated RobustScaler
        
        This is UNSUPERVISED preprocessing - no labels (y) are needed.
        Each party computes local statistics on their data partition.
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features (vertically or horizontally partitioned)
            Data stays local on each PYU
        
        Returns
        -------
        self : FLRobustScaler
            Fitted preprocessor
        """
        if isinstance(x, VDataFrame):
            x = x.values
        
        logging.info("[FL] Starting federated RobustScaler fitting (unsupervised)")
        logging.info(f"[FL] Partitions: {len(x.partitions)}")
        
        # Each party fits local preprocessor on their data partition
        for party_name, device in self.devices.items():
            if device in x.partitions:
                X_local = x.partitions[device]
                model = self.local_models[party_name]
                
                # Fit local preprocessor
                def _local_fit(model, X):
                    model.fit(X)
                    # Return statistics for verification
                    n_samples = X.shape[0]
                    n_features = X.shape[1] if len(X.shape) > 1 else 1
                    return n_samples, n_features
                
                result = device(_local_fit)(model, X_local)
                logging.info(f"[FL] Party '{party_name}' completed local fitting")
        
        self._is_fitted = True
        logging.info("[FL] Federated RobustScaler fitting completed")
        return self
    
    def transform(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Transform data using federated preprocessor
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features for transformation
        
        Returns
        -------
        X_transformed : FedNdarray
            Transformed data
        """
        if not self._is_fitted:
            raise RuntimeError("Preprocessor must be fitted before transformation")
        
        if isinstance(x, VDataFrame):
            x = x.values
        
        # Each party transforms locally
        transformed_list = []
        
        for party_name, device in self.devices.items():
            if device in x.partitions:
                X_local = x.partitions[device]
                model = self.local_models[party_name]
                
                X_trans = device(lambda m, X: m.transform(X))(model, X_local)
                transformed_list.append(X_trans)
        
        # Return transformed data (no aggregation needed for preprocessing)
        # Each party keeps their own transformed data
        if len(transformed_list) == 1:
            return transformed_list[0]
        else:
            # For preprocessing, usually no aggregation is needed
            # Each party's data is independently transformed
            return transformed_list
    
    def fit_transform(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Fit the preprocessor and transform data
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features
        
        Returns
        -------
        X_transformed : FedNdarray
            Transformed data
        """
        self.fit(x)
        return self.transform(x)
    
    def inverse_transform(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Inverse transform data (if supported by the algorithm)
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Transformed features
        
        Returns
        -------
        X_original : FedNdarray
            Original scale data
        """
        if not self._is_fitted:
            raise RuntimeError("Preprocessor must be fitted before inverse transformation")
        
        if isinstance(x, VDataFrame):
            x = x.values
        
        # Each party inverse transforms locally
        inverse_list = []
        
        for party_name, device in self.devices.items():
            if device in x.partitions:
                X_local = x.partitions[device]
                model = self.local_models[party_name]
                
                # Check if inverse_transform is available
                def _local_inverse(m, X):
                    if hasattr(m, 'inverse_transform'):
                        return m.inverse_transform(X)
                    else:
                        raise AttributeError(f"{type(m).__name__} does not support inverse_transform")
                
                X_inv = device(_local_inverse)(model, X_local)
                inverse_list.append(X_inv)
        
        # Return inverse transformed data
        if len(inverse_list) == 1:
            return inverse_list[0]
        else:
            return inverse_list
