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

"""
Secret Sharing adapter for LinearSVC

LinearSVC is a SUPERVISED support vector machine algorithm for classification.
Data aggregated to SPU with full MPC protection.
SPU-based secure aggregation of model parameters.

Mode: Secret Sharing (SS)
Generated by: StandaloneAlgorithmMigrator
"""

import logging
from typing import Union, Optional

import numpy as np

try:
    from xlearn.svm import LinearSVC
    USING_XLEARN = True
except ImportError:
    from sklearn.svm import LinearSVC
    USING_XLEARN = False

try:
    from secretflow.data.ndarray.ndarray import FedNdarray
    from secretflow.data.vertical.dataframe import VDataFrame
    from secretflow.device import SPU
    from secretflow.device.device.spu import SPUObject
    SECRETFLOW_AVAILABLE = True
except ImportError:
    SECRETFLOW_AVAILABLE = False


class SSLinearSVC:
    """
    Secret Sharing LinearSVC
    
    LinearSVC is a supervised support vector machine algorithm for classification.
    
    In SS mode:
    - All data aggregated to SPU (Secure Processing Unit) for encrypted computation
    - Model training happens in MPC environment on SPU
    - Full privacy protection via secret sharing
    
    Parameters
    ----------
    spu : SPU
        Secure Processing Unit for encrypted computation
    **kwargs
        Parameters passed to LinearSVC
    
    Examples
    --------
    >>> import secretflow as sf
    >>> alice = sf.PYU('alice')
    >>> bob = sf.PYU('bob')
    >>> spu = sf.SPU(sf.SPUConfig(...), ...)
    >>> 
    >>> model = SSLinearSVC(spu=spu)
    >>> 
    >>> # Fit on federated data
    >>> model.fit(fed_X, fed_y)
    >>> 
    >>> # Predict
    >>> predictions = model.predict(fed_X_test)
    """
    
    def __init__(
        self,
        spu: 'SPU',
        **kwargs
    ):
        if not SECRETFLOW_AVAILABLE:
            raise RuntimeError("SecretFlow not installed. Install: pip install secretflow")
        
        self.spu = spu
        self.kwargs = kwargs
        
        # Model will be created on SPU
        self.model = None
        self._is_fitted = False
        
        if USING_XLEARN:
            logging.info("[SS] SSLinearSVC initialized with JAX acceleration")
        else:
            logging.info("[SS] SSLinearSVC initialized with sklearn")
        
        logging.info("[SS] Using SPU for secure computation")
    
    @staticmethod
    def _create_model(**kwargs):
        """Create LinearSVC instance"""
        return LinearSVC(**kwargs)
    
    def fit(
        self,
        x: 'Union[FedNdarray, VDataFrame]',
        y: 'Union[FedNdarray, VDataFrame]'
    ):
        """
        Fit the LinearSVC model in SPU
        
        This is SUPERVISED learning - labels (y) are required.
        All computation happens in encrypted SPU environment.
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features (will be aggregated to SPU)
        y : FedNdarray or VDataFrame
            Federated labels (will be aggregated to SPU)
        
        Returns
        -------
        self : SSLinearSVC
            Fitted model
        """
        if isinstance(x, VDataFrame):
            x = x.values
        if isinstance(y, VDataFrame):
            y = y.values
        
        logging.info("[SS] Starting LinearSVC training in SPU (secure, encrypted)")
        
        # Define training function to run on SPU
        def _spu_fit(X, y, **kwargs):
            """Training happens entirely in SPU's encrypted environment"""
            model = LinearSVC(**kwargs)
            model.fit(X, y)
            return model
        
        # Aggregate data to SPU and train
        # Convert FedNdarray partitions to SPU objects

        x_parts = [x.partitions[pyu].to(self.spu) for pyu in x.partitions]

        y_parts = [y.partitions[pyu].to(self.spu) for pyu in y.partitions]

        

        self.model = self.spu(_spu_fit)(x_parts, y_parts, **self.kwargs)
        
        self._is_fitted = True
        logging.info("[SS] LinearSVC training completed in SPU")
        return self
    
    def predict(self, x: 'Union[FedNdarray, VDataFrame]'):
        """
        Predict using model in SPU
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features for prediction
        
        Returns
        -------
        predictions : SPUObject
            Predicted class labels (encrypted in SPU)
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted before prediction")
        
        if isinstance(x, VDataFrame):
            x = x.values
        
        logging.info("[SS] Making predictions in SPU (secure, encrypted)")
        
        # Aggregate data to SPU
        X_spu = x.to(self.spu)
        
        # Predict in SPU
        predictions = self.spu(lambda m, X: m.predict(X))(self.model, X_spu)
        
        return predictions
    
    def score(self, x: 'Union[FedNdarray, VDataFrame]', y: 'Union[FedNdarray, VDataFrame]'):
        """
        Compute model score in SPU
        
        Parameters
        ----------
        x : FedNdarray or VDataFrame
            Federated features
        y : FedNdarray or VDataFrame
            True labels
        
        Returns
        -------
        score : SPUObject
            Model accuracy score (encrypted in SPU)
        """
        predictions = self.predict(x)
        
        if isinstance(y, VDataFrame):
            y = y.values
        
        # Aggregate y to SPU
        y_spu = y.to(self.spu)
        
        # Compute accuracy score in SPU
        def _compute_accuracy(y_true, y_pred):
            from sklearn.metrics import accuracy_score
            return accuracy_score(y_true, y_pred)
        
        score = self.spu(_compute_accuracy)(y_spu, predictions)
        
        return score

