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

"""
Split Learning adapter for VarianceThreshold

VarianceThreshold is an UNSUPERVISED feature selection algorithm.
Model split across parties with collaborative training.
HEU-based secure aggregation of variance statistics.

Mode: Split Learning (SL)
Generated by: StandaloneAlgorithmMigrator
"""

import logging
from typing import Dict, Union, Optional

import numpy as np

try:
    from xlearn.feature_selection import VarianceThreshold
    USING_XLEARN = True
except ImportError:
    from sklearn.feature_selection import VarianceThreshold
    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 SLVarianceThreshold:
    """
    Split Learning VarianceThreshold
    
    VarianceThreshold is an unsupervised feature selection algorithm that
    removes features with low variance.
    
    In SL mode:
    - Each party holds part of the model, collaboratively training
    - Variances are securely aggregated via HEU encryption
    - No labels (y) are needed - this is unsupervised feature selection
    
    Parameters
    ----------
    devices : Dict[str, PYU]
        Dictionary mapping party names to PYU devices
    heu : HEU, optional
        Optional HEU for secure model part communication
    aggregation_method : str, default='mean'
        How to aggregate variances: 'mean', 'weighted_mean'
    **kwargs
        Parameters passed to VarianceThreshold (threshold, etc.)
    
    Examples
    --------
    >>> import secretflow as sf
    >>> alice = sf.PYU('alice')
    >>> bob = sf.PYU('bob')
    >>> heu = sf.HEU(sf.HEUConfig(...), ...)
    >>> 
    >>> # Unsupervised feature selection - no labels needed
    >>> selector = FLVarianceThreshold(
    >>>     devices={'alice': alice, 'bob': bob},
    >>>     heu=heu,
    >>>     threshold=0.1
    >>> )
    >>> 
    >>> # Fit on federated data (no y labels)
    >>> selector.fit(fed_X)
    >>> 
    >>> # Transform data (select features)
    >>> X_selected = selector.transform(fed_X_test)
    >>> 
    >>> # Or fit and transform in one step
    >>> X_selected = selector.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 feature selectors 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("[SL] SLVarianceThreshold initialized with JAX acceleration")
        else:
            logging.info("[SL] SLVarianceThreshold initialized with sklearn")
        
        logging.info(f"[SL] Parties: {list(devices.keys())}")
        logging.info(f"[SL] Aggregation: {aggregation_method}")
        logging.info(f"[SL] HEU enabled: {heu is not None}")
    
    @staticmethod
    def _create_local_model(**kwargs):
        """Create local VarianceThreshold instance"""
        return VarianceThreshold(**kwargs)
    
    def fit(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Fit the federated VarianceThreshold
        
        This is UNSUPERVISED feature selection - no labels (y) are needed.
        Each party holds part of the model, collaboratively training.
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features (vertically or horizontally partitioned)
            Data stays local on each PYU
        
        Returns
        -------
        self : FLVarianceThreshold
            Fitted feature selector
        """
        if isinstance(x, VDataFrame):
            x = x.values
        
        logging.info("[SL] Starting federated VarianceThreshold fitting (unsupervised)")
        logging.info(f"[SL] Partitions: {len(x.partitions)}")
        
        # Each party computes local variances
        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 feature selector
                def _local_fit(model, X):
                    model.fit(X)
                    # Return statistics
                    n_samples = X.shape[0]
                    n_features_in = X.shape[1] if len(X.shape) > 1 else 1
                    variances = getattr(model, 'variances_', None)
                    return n_samples, n_features_in, variances
                
                result = device(_local_fit)(model, X_local)
                logging.info(f"[SL] Party '{party_name}' completed variance computation")
        
        self._is_fitted = True
        logging.info("[SL] Federated VarianceThreshold fitting completed")
        return self
    
    def transform(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Select features based on variance threshold
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features for selection
        
        Returns
        -------
        X_selected : FedNdarray
            Selected features
        """
        if not self._is_fitted:
            raise RuntimeError("VarianceThreshold must be fitted before transformation")
        
        if isinstance(x, VDataFrame):
            x = x.values
        
        # Each party selects features locally
        selected_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_selected = device(lambda m, X: m.transform(X))(model, X_local)
                selected_list.append(X_selected)
        
        # Return selected features
        if len(selected_list) == 1:
            return selected_list[0]
        else:
            # For feature selection, each party keeps their selected features
            return selected_list
    
    def fit_transform(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Fit and select features
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features
        
        Returns
        -------
        X_selected : FedNdarray
            Selected features
        """
        self.fit(x)
        return self.transform(x)
    
    def inverse_transform(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Inverse transform (restore removed features with zeros)
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Selected features
        
        Returns
        -------
        X_original : FedNdarray
            Original shape data with removed features as zeros
        """
        if not self._is_fitted:
            raise RuntimeError("VarianceThreshold 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]
                
                X_inv = device(lambda m, X: m.inverse_transform(X))(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
    
    @property
    def variances_(self):
        """
        Get feature variances from local models
        
        Note: In federated setting, this returns aggregated variances
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted first")
        
        # Collect variances from all parties
        variances_list = []
        for party_name, device in self.devices.items():
            model = self.local_models[party_name]
            variances = device(lambda m: getattr(m, 'variances_', None))(model)
            if variances is not None:
                variances_list.append(variances)
        
        if not variances_list:
            return None
        
        # Aggregate variances
        if self.heu:
            aggregator = SecureAggregator(device=self.heu)
            return aggregator.average(variances_list)
        else:
            return np.mean(variances_list, axis=0)
