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

"""
Split Learning adapter for LogisticRegression

LogisticRegression is a SUPERVISED cross-validation model.
Model split across parties with collaborative training.
HEU-based secure aggregation.

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

import logging
from typing import Dict, Union, Optional

import numpy as np

try:
    from xlearn.linear_model import LogisticRegression
    USING_XLEARN = True
except ImportError:
    from sklearn.linear_model import LogisticRegression
    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 SLLogisticRegression:
    """
    Split Learning LogisticRegression
    
    Cross-validation model with automatic hyperparameter tuning.
    
    In SL mode:
    - Each party performs local CV training
    - Best parameters are securely aggregated via HEU
    - Non-iterative training (CV internally iterates)
    
    Parameters
    ----------
    devices : Dict[str, PYU]
        Dictionary mapping party names to PYU devices
    heu : HEU, optional
        Optional HEU for secure model part communication
    **kwargs
        Parameters passed to LogisticRegression
    
    Examples
    --------
    >>> model = FLLogisticRegression(
    >>>     devices={'alice': alice, 'bob': bob},
    >>>     heu=heu,
    >>>     cv=5
    >>> )
    >>> model.fit(fed_X, fed_y)
    >>> 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("[SL] SLLogisticRegression initialized with JAX acceleration")
        else:
            logging.info("[SL] SLLogisticRegression initialized with sklearn")
    
    @staticmethod
    def _create_local_model(**kwargs):
        return LogisticRegression(**kwargs)
    
    def fit(self, x: 'Union[FedNdarray, VDataFrame]', y: 'Union[FedNdarray, VDataFrame]'):
        """Fit with cross-validation"""
        if isinstance(x, VDataFrame):
            x = x.values
        if isinstance(y, VDataFrame):
            y = y.values
        
        logging.info("[SL] Starting federated LogisticRegression training (CV)")
        
        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_fit(model, X, y):
                    model.fit(X, y)
                    params = {}
                    if hasattr(model, 'coef_'):
                        params['coef_'] = model.coef_
                    if hasattr(model, 'intercept_'):
                        params['intercept_'] = model.intercept_
                    if hasattr(model, 'alpha_'):
                        params['alpha_'] = model.alpha_
                    return params
                
                result = device(_local_fit)(model, X_local, y_local)
                logging.info(f"[SL] Party '{party_name}' completed CV training")
        
        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)
