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

"""
Federated Learning adapter for TruncatedSVD

TruncatedSVD is an UNSUPERVISED dimensionality reduction algorithm.
Data remains in local PYUs, JAX-accelerated local computation,
HEU-based secure aggregation.

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

import logging
from typing import Dict, Union, Optional

import numpy as np

try:
    from xlearn.decomposition import TruncatedSVD
    USING_XLEARN = True
except ImportError:
    from sklearn.decomposition import TruncatedSVD
    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 FLTruncatedSVD:
    """
    Federated Learning TruncatedSVD
    
    TruncatedSVD is an unsupervised dimensionality reduction algorithm.
    
    In FL mode:
    - Each party performs local dimensionality reduction on their data partition
    - Results are securely aggregated via HEU encryption
    - No labels (y) are needed - this is unsupervised learning
    
    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 results: 'mean', 'weighted_mean'
    **kwargs
        Parameters passed to TruncatedSVD
    
    Examples
    --------
    >>> import secretflow as sf
    >>> alice = sf.PYU('alice')
    >>> bob = sf.PYU('bob')
    >>> heu = sf.HEU(sf.HEUConfig(...), ...)
    >>> 
    >>> # Unsupervised dimensionality reduction - no labels needed
    >>> model = FLTruncatedSVD(
    >>>     devices={'alice': alice, 'bob': bob},
    >>>     heu=heu
    >>> )
    >>> 
    >>> # Fit on federated data (no y labels)
    >>> model.fit(fed_X)
    >>> 
    >>> # Transform data
    >>> X_transformed = model.transform(fed_X_test)
    """
    
    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 models 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] FLTruncatedSVD initialized with JAX acceleration")
        else:
            logging.info("[FL] FLTruncatedSVD 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 TruncatedSVD instance"""
        return TruncatedSVD(**kwargs)
    
    def fit(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Fit the federated TruncatedSVD model
        
        This is UNSUPERVISED learning - no labels (y) are needed.
        Each party performs local dimensionality reduction on their data partition.
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features (vertically or horizontally partitioned)
            Data stays local on each PYU
        
        Returns
        -------
        self : FLTruncatedSVD
            Fitted model
        """
        if isinstance(x, VDataFrame):
            x = x.values
        
        logging.info("[FL] Starting federated TruncatedSVD dimensionality reduction (unsupervised)")
        logging.info(f"[FL] Partitions: {len(x.partitions)}")
        
        # Each party performs local dimensionality reduction 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]
                
                # Train local model
                def _local_fit(model, X):
                    model.fit(X)
                    # Return model statistics for verification
                    n_samples = X.shape[0]
                    n_features = X.shape[1] if len(X.shape) > 1 else 1
                    n_components = getattr(model, 'n_components', None) or getattr(model, 'n_components_', None)
                    return n_samples, n_features, n_components
                
                result = device(_local_fit)(model, X_local)
                logging.info(f"[FL] Party '{party_name}' completed local dimensionality reduction")
        
        self._is_fitted = True
        logging.info("[FL] Federated TruncatedSVD dimensionality reduction completed")
        return self
    
    def transform(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Transform data using federated model
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features for transformation
        
        Returns
        -------
        X_transformed : FedNdarray
            Transformed data
        """
        if not self._is_fitted:
            raise RuntimeError("Model 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)
        
        # Aggregate transformed data
        if len(transformed_list) == 1:
            return transformed_list[0]
        else:
            # Use HEU for secure aggregation if available
            if self.heu:
                return self._secure_aggregate_transform(transformed_list)
            else:
                logging.warning("[FL] Using non-secure aggregation (HEU not provided)")
                return self._simple_aggregate_transform(transformed_list)
    
    def fit_transform(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Fit the model 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 space data
        """
        if not self._is_fitted:
            raise RuntimeError("Model 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)
        
        # Aggregate
        if len(inverse_list) == 1:
            return inverse_list[0]
        else:
            if self.heu:
                return self._secure_aggregate_transform(inverse_list)
            else:
                return self._simple_aggregate_transform(inverse_list)
    
    def _secure_aggregate_transform(self, transform_list):
        """
        Securely aggregate transformed data using HEU encryption
        """
        logging.info("[FL] Secure transformation aggregation via HEU")
        aggregator = SecureAggregator(device=self.heu)
        if self.aggregation_method == 'mean':
            return aggregator.average(transform_list)
        elif self.aggregation_method == 'weighted_mean':
            # Weighted average based on sample counts from each party
            weights = []
            for party_name, device in self.devices.items():
                if device in x.partitions:
                    X_local = x.partitions[device]
                    n_samples = device(lambda X: X.shape[0])(X_local)
                    weights.append(n_samples)
            
            total = sum(weights)
            normalized_weights = [w / total for w in weights]
            
            weighted_results = []
            for t, w in zip(transform_list, normalized_weights):
                weighted_t = t * w
                weighted_results.append(weighted_t)
            
            return aggregator.sum(weighted_results, axis=0)
        else:
            raise ValueError(f"Unsupported aggregation method: {self.aggregation_method}")
    
    def _simple_aggregate_transform(self, transform_list):
        """
        Simple (non-secure) transformation aggregation - for development/testing only
        
        WARNING: This does not preserve privacy. Use only for testing.
        """
        if len(transform_list) == 1:
            return transform_list[0]
        
        # Simple average
        return np.mean(transform_list, axis=0)
    
    @property
    def components_(self):
        """
        Get components from local models (for PCA, NMF, etc.)
        
        Note: In federated setting, this returns aggregated components
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted first")
        
        # Collect components from all parties
        components_list = []
        for party_name, device in self.devices.items():
            model = self.local_models[party_name]
            components = device(lambda m: getattr(m, 'components_', None))(model)
            if components is not None:
                components_list.append(components)
        
        if not components_list:
            return None
        
        # Aggregate components
        if self.heu:
            aggregator = SecureAggregator(device=self.heu)
            return aggregator.average(components_list)
        else:
            return np.mean(components_list, axis=0)
    
    @property
    def explained_variance_(self):
        """
        Get explained variance (for PCA, etc.)
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted first")
        
        # Collect from all parties
        variance_list = []
        for party_name, device in self.devices.items():
            model = self.local_models[party_name]
            variance = device(lambda m: getattr(m, 'explained_variance_', None))(model)
            if variance is not None:
                variance_list.append(variance)
        
        if not variance_list:
            return None
        
        # Aggregate
        if self.heu:
            aggregator = SecureAggregator(device=self.heu)
            return aggregator.average(variance_list)
        else:
            return np.mean(variance_list, axis=0)
    
    @property
    def explained_variance_ratio_(self):
        """
        Get explained variance ratio (for PCA, etc.)
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted first")
        
        # Collect from all parties
        ratio_list = []
        for party_name, device in self.devices.items():
            model = self.local_models[party_name]
            ratio = device(lambda m: getattr(m, 'explained_variance_ratio_', None))(model)
            if ratio is not None:
                ratio_list.append(ratio)
        
        if not ratio_list:
            return None
        
        # Aggregate
        if self.heu:
            aggregator = SecureAggregator(device=self.heu)
            return aggregator.average(ratio_list)
        else:
            return np.mean(ratio_list, axis=0)
