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

"""
Federated Learning adapter for KMeans

KMeans is an UNSUPERVISED clustering algorithm.
Data remains in local PYUs, JAX-accelerated local computation,
HEU-based secure aggregation of cluster assignments.

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

import logging
from typing import Dict, Union, Optional

import numpy as np

try:
    from xlearn.cluster import KMeans
    USING_XLEARN = True
except ImportError:
    from sklearn.cluster import KMeans
    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 FLKMeans:
    """
    Federated Learning KMeans
    
    KMeans is an unsupervised clustering algorithm.
    
    In FL mode:
    - Each party performs local clustering on their data partition
    - Cluster centers/assignments 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 cluster centers: 'mean', 'weighted_mean'
    **kwargs
        Parameters passed to KMeans
    
    Examples
    --------
    >>> import secretflow as sf
    >>> alice = sf.PYU('alice')
    >>> bob = sf.PYU('bob')
    >>> heu = sf.HEU(sf.HEUConfig(...), ...)
    >>> 
    >>> # Unsupervised clustering - no labels needed
    >>> model = FLKMeans(
    >>>     devices={'alice': alice, 'bob': bob},
    >>>     heu=heu
    >>> )
    >>> 
    >>> # Fit on federated data (no y labels)
    >>> model.fit(fed_X)
    >>> 
    >>> # Predict cluster assignments
    >>> labels = model.predict(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] FLKMeans initialized with JAX acceleration")
        else:
            logging.info("[FL] FLKMeans 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 KMeans instance"""
        return KMeans(**kwargs)
    
    def fit(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Fit the federated KMeans model
        
        This is UNSUPERVISED learning - no labels (y) are needed.
        Each party performs local clustering on their data partition.
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features (vertically or horizontally partitioned)
            Data stays local on each PYU
        
        Returns
        -------
        self : FLKMeans
            Fitted model
        """
        if isinstance(x, VDataFrame):
            x = x.values
        
        logging.info("[FL] Starting federated KMeans clustering (unsupervised)")
        logging.info(f"[FL] Partitions: {len(x.partitions)}")
        
        # Each party performs local clustering 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 clustering 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_clusters = getattr(model, 'n_clusters', None) or len(getattr(model, 'cluster_centers_', []))
                    return n_samples, n_features, n_clusters
                
                result = device(_local_fit)(model, X_local)
                logging.info(f"[FL] Party '{party_name}' completed local clustering")
        
        self._is_fitted = True
        logging.info("[FL] Federated KMeans clustering completed")
        return self
    
    def predict(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Predict cluster assignments using federated model
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features for prediction
        
        Returns
        -------
        labels : FedNdarray
            Cluster assignments
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted before prediction")
        
        if isinstance(x, VDataFrame):
            x = x.values
        
        # Each party predicts locally, then we aggregate
        predictions_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]
                
                pred = device(lambda m, X: m.predict(X))(model, X_local)
                predictions_list.append(pred)
        
        # Aggregate predictions
        if len(predictions_list) == 1:
            return predictions_list[0]
        else:
            # Use HEU for secure aggregation if available
            if self.heu:
                return self._secure_aggregate_labels(predictions_list)
            else:
                logging.warning("[FL] Using non-secure aggregation (HEU not provided)")
                return self._simple_aggregate_labels(predictions_list)
    
    def fit_predict(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Fit the model and predict cluster assignments
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features
        
        Returns
        -------
        labels : FedNdarray
            Cluster assignments
        """
        self.fit(x)
        return self.predict(x)
    
    def _secure_aggregate_labels(self, labels_list):
        """
        Securely aggregate cluster labels using HEU encryption
        
        This implements proper privacy-preserving label aggregation.
        """
        logging.info("[FL] Secure label aggregation via HEU")
        # Use HEU for secure voting/averaging
        aggregator = SecureAggregator(device=self.heu)
        return aggregator.majority_vote(labels_list)
    
    def _simple_aggregate_labels(self, labels_list):
        """
        Simple (non-secure) label aggregation - for development/testing only
        
        WARNING: This does not preserve privacy. Use only for testing.
        """
        # Simple majority voting
        if len(labels_list) == 1:
            return labels_list[0]
        
        # Stack and take mode (most common label)
        labels_array = np.stack([np.array(labels) for labels in labels_list])
        # Use mode along axis 0
        from scipy import stats
        result, _ = stats.mode(labels_array, axis=0, keepdims=False)
        return result
    
    @property
    def cluster_centers_(self):
        """
        Get cluster centers from local models
        
        Note: In federated setting, this returns aggregated centers
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted first")
        
        # Collect centers from all parties
        centers_list = []
        for party_name, device in self.devices.items():
            model = self.local_models[party_name]
            centers = device(lambda m: getattr(m, 'cluster_centers_', None))(model)
            if centers is not None:
                centers_list.append(centers)
        
        if not centers_list:
            return None
        
        # Aggregate centers
        if self.heu:
            aggregator = SecureAggregator(device=self.heu)
            return aggregator.average(centers_list)
        else:
            # Simple average
            return np.mean(centers_list, axis=0)
    
    @property
    def labels_(self):
        """
        Get labels from last fit
        
        Note: In federated setting, this returns aggregated labels
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted first")
        
        # Collect labels from all parties
        labels_list = []
        for party_name, device in self.devices.items():
            model = self.local_models[party_name]
            labels = device(lambda m: getattr(m, 'labels_', None))(model)
            if labels is not None:
                labels_list.append(labels)
        
        if not labels_list:
            return None
        
        # Aggregate labels
        if len(labels_list) == 1:
            return labels_list[0]
        else:
            return self._simple_aggregate_labels(labels_list)
