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

"""
Split Learning adapter for PLSCanonical

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

import logging
from typing import Dict, Union, Optional
import numpy as np

try:
    from xlearn.cross_decomposition import PLSCanonical
    USING_XLEARN = True
except ImportError:
    from sklearn.cross_decomposition import PLSCanonical
    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.security.aggregation import SecureAggregator
    SECRETFLOW_AVAILABLE = True
except ImportError:
    SECRETFLOW_AVAILABLE = False


class SLPLSCanonical:
    """Split Learning PLSCanonical"""
    
    def __init__(self, devices: Dict[str, 'PYU'], heu: Optional['HEU'] = None, **kwargs):
        if not SECRETFLOW_AVAILABLE:
            raise RuntimeError("SecretFlow not installed")
        self.devices = devices
        self.heu = heu
        self.local_models = {}
        for party_name, device in devices.items():
            self.local_models[party_name] = device(lambda **kw: PLSCanonical(**kw))(**kwargs)
        self._is_fitted = False
        if USING_XLEARN:
            logging.info(f"[SL] SLPLSCanonical with JAX acceleration")
    
    def fit(self, x: 'Union[FedNdarray, VDataFrame]', y: 'Union[FedNdarray, VDataFrame]' = None):
        """Fit federated model"""
        if isinstance(x, VDataFrame):
            x = x.values
        if y is not None and 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) if y is not None else None
                model = self.local_models[party_name]
                
                if y_local is not None:
                    device(lambda m, X, y: m.fit(X, y))(model, X_local, y_local)
                else:
                    device(lambda m, X: m.fit(X))(model, X_local)
                
                logging.info(f"[SL] Party '{party_name}' trained")
        
        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]
        if self.heu:
            aggregator = SecureAggregator(device=self.heu)
            return aggregator.average(predictions_list)
        return np.mean(predictions_list, axis=0)
    
    def transform(self, x: 'Union[FedNdarray, VDataFrame]'):
        """Transform (if supported)"""
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted before transformation")
        if isinstance(x, VDataFrame):
            x = x.values
        
        transformed_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]
                
                def _transform(m, X):
                    if hasattr(m, 'transform'):
                        return m.transform(X)
                    raise AttributeError("Model does not support transform")
                
                X_trans = device(_transform)(model, X_local)
                transformed_list.append(X_trans)
        
        if len(transformed_list) == 1:
            return transformed_list[0]
        return transformed_list
