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

"""
Federated Learning adapter for SGDClassifier

SGDClassifier is an ITERATIVE SUPERVISED algorithm.
Data remains in local PYUs, JAX-accelerated local computation,
HEU-based secure aggregation after each epoch.

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

import logging
from typing import Dict, Union, Optional

import numpy as np

try:
    from xlearn.linear_model import SGDClassifier
    USING_XLEARN = True
except ImportError:
    from sklearn.linear_model import SGDClassifier
    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 FLSGDClassifier:
    """
    Federated Learning SGDClassifier
    
    SGDClassifier is an iterative supervised algorithm that supports
    incremental learning via partial_fit.
    
    In FL mode:
    - Each party performs local partial_fit in each epoch
    - Model parameters are securely aggregated via HEU after each epoch
    - Supports online/incremental learning
    
    Parameters
    ----------
    devices : Dict[str, PYU]
        Dictionary mapping party names to PYU devices
    heu : HEU, optional
        Homomorphic encryption unit for secure aggregation
    **kwargs
        Parameters passed to SGDClassifier
    
    Examples
    --------
    >>> model = FLSGDClassifier(
    >>>     devices={'alice': alice, 'bob': bob},
    >>>     heu=heu
    >>> )
    >>> # Iterative training with multiple epochs
    >>> model.fit(fed_X, fed_y, epochs=10)
    >>> predictions = model.predict(fed_X_test)
    """
    
    def __init__(
        self,
        devices: Dict[str, 'PYU'],
        heu: Optional['HEU'] = None,
        **kwargs
    ):
        if not SECRETFLOW_AVAILABLE:
            raise RuntimeError("SecretFlow not installed. Install: pip install secretflow")
        
        self.devices = devices
        self.heu = heu
        self.kwargs = kwargs
        
        self.local_models = {}
        for party_name, device in devices.items():
            self.local_models[party_name] = device(self._create_local_model)(**kwargs)
        
        self._is_fitted = False
        
        if USING_XLEARN:
            logging.info("[FL] FLSGDClassifier initialized with JAX acceleration")
        else:
            logging.info("[FL] FLSGDClassifier initialized with sklearn")
    
    @staticmethod
    def _create_local_model(**kwargs):
        return SGDClassifier(**kwargs)
    
    def fit(
        self,
        x: 'Union[FedNdarray, VDataFrame]',
        y: 'Union[FedNdarray, VDataFrame]',
        epochs: int = 10
    ):
        """
        Iterative training with partial_fit
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features
        y : FedNdarray or VDataFrame
            Federated labels
        epochs : int
            Number of training epochs (default: 10)
        """
        if isinstance(x, VDataFrame):
            x = x.values
        if isinstance(y, VDataFrame):
            y = y.values
        
        logging.info(f"[FL] Starting federated SGDClassifier training (iterative, {epochs} epochs)")
        
        for epoch in range(epochs):
            for party_name, device in self.devices.items():
                if device in x.partitions:
                    X_local = x.partitions[device]
                    y_local = y.partitions.get(device, y)
                    model = self.local_models[party_name]
                    
                    def _local_partial_fit(model, X, y):
                        # Get classes for first call
                        classes = None
                        if not hasattr(model, 'classes_'):
                            classes = np.unique(y)
                            model.partial_fit(X, y, classes=classes)
                        else:
                            model.partial_fit(X, y)
                        
                        # Return parameters
                        params = {}
                        if hasattr(model, 'coef_'):
                            params['coef_'] = model.coef_
                        if hasattr(model, 'intercept_'):
                            params['intercept_'] = model.intercept_
                        return params
                    
                    result = device(_local_partial_fit)(model, X_local, y_local)
            
            # Aggregate parameters after each epoch
            if self.heu:
                self._secure_aggregate_parameters()
            
            logging.info(f"[FL] Epoch {epoch+1}/{epochs} completed")
        
        self._is_fitted = True
        logging.info("[FL] Federated SGDClassifier training completed")
        return self
    
    def partial_fit(
        self,
        x: 'Union[FedNdarray, VDataFrame]',
        y: 'Union[FedNdarray, VDataFrame]',
        classes=None
    ):
        """
        Incremental fit on a batch of samples
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Batch of features
        y : FedNdarray or VDataFrame
            Batch of labels
        classes : array-like, optional
            Classes for classification (required on first call)
        """
        if isinstance(x, VDataFrame):
            x = x.values
        if isinstance(y, VDataFrame):
            y = y.values
        
        for party_name, device in self.devices.items():
            if device in x.partitions:
                X_local = x.partitions[device]
                y_local = y.partitions.get(device, y)
                model = self.local_models[party_name]
                
                def _local_partial_fit(model, X, y, classes):
                    if classes is not None:
                        model.partial_fit(X, y, classes=classes)
                    else:
                        model.partial_fit(X, y)
                    return True
                
                device(_local_partial_fit)(model, X_local, y_local, classes)
        
        # Aggregate parameters
        if self.heu:
            self._secure_aggregate_parameters()
        
        self._is_fitted = True
        return self
    
    def predict(self, x: 'Union[FedNdarray, VDataFrame]'):
        """Predict using federated model"""
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted before prediction")
        
        if isinstance(x, VDataFrame):
            x = x.values
        
        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)
        
        if len(predictions_list) == 1:
            return predictions_list[0]
        else:
            if self.heu:
                aggregator = SecureAggregator(device=self.heu)
                return aggregator.average(predictions_list)
            else:
                return np.mean(predictions_list, axis=0)
    
    def _secure_aggregate_parameters(self):
        """
        Securely aggregate model parameters across parties using HEU
        
        For iterative models (SGD, Perceptron, etc.), called after
        each epoch to synchronize parameters across parties.
        """
        logging.debug("[FL] Secure parameter aggregation via SecureAggregator")
        
        parties = list(self.devices.values())
        host_party = parties[0]
        aggregator = SecureAggregator(device=host_party, participants=parties)
        
        coef_list = []
        intercept_list = []
        
        for party_name, device in self.devices.items():
            model = self.local_models[party_name]
            
            def _extract_params(m):
                params = {}
                if hasattr(m, 'coef_'):
                    params['coef_'] = m.coef_
                if hasattr(m, 'intercept_'):
                    params['intercept_'] = m.intercept_
                return params
            
            params = device(_extract_params)(model)
            if 'coef_' in params:
                coef_list.append(params['coef_'])
            if 'intercept_' in params:
                intercept_list.append(params['intercept_'])
        
        aggregated = {}
        if coef_list:
            aggregated['coef_'] = aggregator.average(coef_list, axis=0)
        if intercept_list:
            aggregated['intercept_'] = aggregator.average(intercept_list, axis=0)
        
        # Update local models
        for party_name, device in self.devices.items():
            model = self.local_models[party_name]
            
            def _update_params(m, params):
                if 'coef_' in params:
                    m.coef_ = params['coef_']
                if 'intercept_' in params:
                    m.intercept_ = params['intercept_']
                return m
            
            device(_update_params)(model, aggregated)
        
        return aggregated
