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

"""
Split Learning adapter for NuSVR

NuSVR is a SUPERVISED support vector machine algorithm.
Model split across parties with collaborative training.
HEU-based secure aggregation of model parameters.

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

import logging
from typing import Dict, Union, Optional

import numpy as np

try:
    from xlearn.svm import NuSVR
    USING_XLEARN = True
except ImportError:
    from sklearn.svm import NuSVR
    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 SLNuSVR:
    """
    Split Learning NuSVR
    
    NuSVR is a supervised support vector machine algorithm.
    
    In SL mode:
    - Each party holds part of the model, collaboratively training
    - Model parameters are securely aggregated via HEU encryption
    - Training is non-iterative (single pass)
    
    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 model parameters: 'mean', 'weighted_mean'
    **kwargs
        Parameters passed to NuSVR
    
    Examples
    --------
    >>> import secretflow as sf
    >>> alice = sf.PYU('alice')
    >>> bob = sf.PYU('bob')
    >>> heu = sf.HEU(sf.HEUConfig(...), ...)
    >>> 
    >>> model = FLNuSVR(
    >>>     devices={'alice': alice, 'bob': bob},
    >>>     heu=heu
    >>> )
    >>> 
    >>> # Fit on federated data
    >>> model.fit(fed_X, fed_y)
    >>> 
    >>> # Predict
    >>> predictions = 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("[SL] SLNuSVR initialized with JAX acceleration")
        else:
            logging.info("[SL] SLNuSVR 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 NuSVR instance"""
        return NuSVR(**kwargs)
    
    def fit(
        self,
        x: 'Union[FedNdarray, VDataFrame]',
        y: 'Union[FedNdarray, VDataFrame]'
    ):
        """
        Fit the federated NuSVR model
        
        This is SUPERVISED learning - labels (y) are required.
        Training is non-iterative (single pass).
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features (data stays local)
        y : FedNdarray or VDataFrame
            Federated labels
        
        Returns
        -------
        self : FLNuSVR
            Fitted model
        """
        if isinstance(x, VDataFrame):
            x = x.values
        if isinstance(y, VDataFrame):
            y = y.values
        
        logging.info("[SL] Starting federated NuSVR training (supervised, non-iterative)")
        logging.info(f"[SL] Partitions: {len(x.partitions)}")
        
        # Each party trains local model (single pass)
        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)  # y might be on one party
                model = self.local_models[party_name]
                
                # Train local model (single pass)
                def _local_fit(model, X, y):
                    model.fit(X, y)
                    # Return model parameters for aggregation
                    params = {}
                    if hasattr(model, 'coef_'):
                        params['coef_'] = model.coef_
                    if hasattr(model, 'intercept_'):
                        params['intercept_'] = model.intercept_
                    n_samples = X.shape[0]
                    n_features = X.shape[1] if len(X.shape) > 1 else 1
                    return params, n_samples, n_features
                
                result = device(_local_fit)(model, X_local, y_local)
                logging.info(f"[SL] Party '{party_name}' completed local training")
        
        # Secure aggregation of model parameters
        if self.heu:
            self._secure_aggregate_parameters()
        else:
            logging.warning("[SL] No HEU provided - using local models only")
        
        self._is_fitted = True
        logging.info("[SL] Federated NuSVR training completed")
        return self
    
    def predict(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Predict using federated model
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features for prediction
        
        Returns
        -------
        predictions : FedNdarray
            Predicted values or labels
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted before prediction")
        
        if isinstance(x, VDataFrame):
            x = x.values
        
        # Each party predicts locally
        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:
            if self.heu:
                return self._secure_aggregate_predictions(predictions_list)
            else:
                logging.warning("[SL] Using non-secure aggregation (HEU not provided)")
                return self._simple_aggregate_predictions(predictions_list)
    
    def _secure_aggregate_parameters(self):
        """
        Securely aggregate model parameters across parties using HEU
        
        Collects coef_, intercept_ from each party's local model and
        performs secure weighted averaging based on sample counts.
        
        Returns
        -------
        aggregated_params : dict
            Dictionary containing securely aggregated model parameters
        """
        logging.info("[SL] Secure parameter aggregation via SecureAggregator")
        
        # Get participating parties
        parties = list(self.devices.values())
        host_party = parties[0]
        
        # Create SecureAggregator
        aggregator = SecureAggregator(device=host_party, participants=parties)
        
        # Collect parameters from each party
        coef_list = []
        intercept_list = []
        sample_counts = []
        
        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_'])
        
        # Securely aggregate parameters
        aggregated = {}
        if coef_list:
            aggregated['coef_'] = aggregator.average(coef_list, axis=0)
        if intercept_list:
            aggregated['intercept_'] = aggregator.average(intercept_list, axis=0)
        
        return aggregated
    
    def _secure_aggregate_predictions(self, predictions_list):
        """Securely aggregate predictions using HEU"""
        logging.info("[SL] Secure prediction aggregation via HEU")
        aggregator = SecureAggregator(device=self.heu)
        return aggregator.average(predictions_list)
    
    def _simple_aggregate_predictions(self, predictions_list):
        """Simple aggregation - for development only"""
        return np.mean(predictions_list, axis=0)
    
    def score(self, x: 'Union[FedNdarray, VDataFrame]', y: 'Union[FedNdarray, VDataFrame]'):
        """
        Compute model score
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features
        y : FedNdarray or VDataFrame
            True labels
        
        Returns
        -------
        score : float
            Model score
        """
        predictions = self.predict(x)
        
        # Compute accuracy/R2 score
        # This is simplified - should use proper scoring
        if isinstance(y, VDataFrame):
            y = y.values
        
        from sklearn.metrics import accuracy_score, r2_score
        
        # Determine if classification or regression
        try:
            return accuracy_score(y, predictions)
        except:
            return r2_score(y, predictions)
