"""MixIn classes for compatible models with the federated algorithms."""
from __future__ import annotations

from abc import ABC, abstractmethod
import os
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Generic,
    List,
    Mapping,
    MutableMapping,
    Optional,
    Sequence,
    Tuple,
    Union,
    cast,
)

from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
import numpy as np
import pandas as pd

from bitfount.config import BITFOUNT_DEFAULT_BATCHED_EXECUTION
from bitfount.data.dataloaders import BitfountDataLoader
from bitfount.data.datasources.base_source import BaseSource
from bitfount.federated.aggregators.base import _BaseAggregatorFactory
from bitfount.federated.algorithms.model_algorithms.evaluate import ModelEvaluation
from bitfount.federated.algorithms.model_algorithms.federated_training import (
    FederatedModelTraining,
)
from bitfount.federated.algorithms.model_algorithms.inference import ModelInference
from bitfount.federated.authorisation_checkers import IdentityVerificationMethod
from bitfount.federated.exceptions import PSIMultiplePodsError
from bitfount.federated.helper import (
    _check_and_update_pod_ids,
    _create_aggregator,
    _create_message_service,
    _get_idp_url,
)
from bitfount.federated.modeller import _Modeller
from bitfount.federated.protocols.model_protocols.federated_averaging import (
    FederatedAveraging,
)
from bitfount.federated.protocols.psi import (
    PrivateSetIntersection,
    _PSICompatibleAlgoFactory_,
)
from bitfount.federated.protocols.results_only import (
    ResultsOnly,
    _ResultsOnlyCompatibleAlgoFactory_,
)
from bitfount.federated.transport.message_service import _MessageService
from bitfount.hub.helper import _create_bitfounthub

if TYPE_CHECKING:
    from bitfount.data.datastructure import DataStructure
    from bitfount.federated.transport.config import MessageServiceConfig
    from bitfount.hub.api import BitfountHub
    from bitfount.metrics import Metric
    from bitfount.types import (
        _SerializedWeights,
        _Weights,
        _Residuals,
        _DistributedModelTypeOrReference,
    )

from bitfount.federated.logging import _get_federated_logger
from bitfount.types import T_DTYPE

logger = _get_federated_logger(__name__)


class _DistributedModelMixIn(ABC, Generic[T_DTYPE]):
    """A mixin for models used in federated mechanisms.

    An abstract base mixin for models that are compatible with the following
    distributed learning protocols:
        - FederatedAveraging
    """

    datastructure: DataStructure
    # Set on the BaseModel
    epochs: Optional[int] = None
    steps: Optional[int] = None
    # Used to identify the relevant section of the datastructure that applies to
    # the pod/logical pod/datasource that this model is being run against (if any)
    _datastructure_identifier: Optional[str] = None

    @abstractmethod
    def get_param_states(self) -> _Weights:
        """Gets the current states of the trainable parameters of the model.

        Returns:
            A dict of param names to tensors
        """
        raise NotImplementedError

    @abstractmethod
    def apply_weight_updates(self, weight_updates: Sequence[_Residuals]) -> _Weights:
        """Applies weight updates to the weights of this model.

        Apply a sequence of parameter weight updates (mappings of parameter name
        to a tensor describing the weight update) to the parameters of this model.
        Used by Modeller to apply updates received from Workers.

        Args:
            weight_updates (Sequence[_Residuals]): The sequence
                of weight updates

        Returns:
            The updated parameters as a dict of name to tensor.
        """
        raise NotImplementedError

    @abstractmethod
    def update_params(self, new_model_params: _Weights) -> None:
        """Updates the current model parameters to the ones provided.

        Used by Worker to update new parameters received from Modeller.

        Args:
            new_model_params (WeightMapping): The new model parameters
                to update to as a mapping of parameter names to tensors.
        """
        raise NotImplementedError

    @abstractmethod
    def deserialize_params(self, serialized_weights: _SerializedWeights) -> _Weights:
        """Converts serialized model parameters to tensors.

        Used by Worker to convert serialized parameters received from Modeller.

        Args:
            serialized_weights (_SerializedWeights): The model parameters
                to deserialize.
        """
        raise NotImplementedError

    @abstractmethod
    def serialize_params(self, weights: _Weights) -> _SerializedWeights:
        """Serializes model parameters.

        Used by Modeller to serialize model parameters to send to the Worker.

        Args:
            weights (_Weights): The model parameters
                to serialize.
        """
        raise NotImplementedError

    @abstractmethod
    def serialize(self, filename: Union[str, os.PathLike]) -> None:
        """Serialize model to file with provided `filename`.

        Args:
           filename: Path to file to save serialized model.
        """
        raise NotImplementedError

    @abstractmethod
    def deserialize(self, content: Union[str, os.PathLike, bytes]) -> None:
        """Deserialize model.

        Args:
            content: Byte stream or path to file containing serialized model.
        """
        raise NotImplementedError

    @abstractmethod
    def diff_params(self, old_params: _Weights, new_params: _Weights) -> _Residuals:
        """Calculates the difference between two sets of model parameters."""
        raise NotImplementedError

    @abstractmethod
    def set_model_training_iterations(self, iterations: int) -> None:
        """Sets model steps or epochs to the appropriate number between updates."""
        raise NotImplementedError

    @abstractmethod
    def reset_trainer(self) -> None:
        """Resets the trainer to its initial state.

        :::note

        Importantly, calling this method in between `fit` calls allows the caller to
        repeatedly refit the model continuing from the batch after the one that was last
        fit. This only applies to step-wise training.

        :::
        """
        raise NotImplementedError

    @abstractmethod
    def log_(self, name: str, value: Any, **kwargs: Any) -> Any:
        """Logs a metric with a particular value to the user's configured model loggers.

        Args:
            name: The name of the metric to log
            value: The value of the metric to log
            **kwargs: Additional keyword arguments to pass to the logger
        """
        raise NotImplementedError

    @abstractmethod
    def tensor_precision(self) -> T_DTYPE:
        """Gets the floating point precision used by model tensors.

        Typically this will be 32 bits.
        """
        raise NotImplementedError

    @abstractmethod
    def _fit_local(
        self,
        data: BaseSource,
        metrics: Optional[MutableMapping[str, Metric]] = None,
        **kwargs: Any,
    ) -> Dict[str, str]:
        """Fits model locally using data from `datasource`.

        Should be implemented in the final model class that subclasses
        DistributedModelMixIn.
        """
        raise NotImplementedError

    @abstractmethod
    def _get_model(
        self,
        extra_imports: Optional[List[str]] = None,
        hub: Optional[BitfountHub] = None,
        **kwargs: Any,
    ) -> _DistributedModelTypeOrReference:
        """Gets the model to be used in distributed training.

        This method is used by the Modeller to get the model to be used in
        distributed training.

        Args:
            extra_imports: A list of extra imports to add to the model file if the model
                is a custom model.
            hub: A BitfountHub instance to be passed to `BitfountModelReference`.
                Optional.
            **kwargs: Additional keyword arguments to pass to the
                `BitfountModelReference` constructor as hyperparameters.

        Returns:
            Either a `BitfountModelReference` object if the model is a custom model, or
            the model itself if it is a built-in model.
        """
        raise NotImplementedError

    def set_datastructure_identifier(self, datastructure_identifier: str) -> None:
        """Sets the datastructure identifier for the model.

        This is used to identify which elements of the datastructure are referring
        to the current runner and is normally the pod identifier (for single
        datasource pods) or the "logical pod" identifier for pods with multiple
        datasources that have registered as different pods.

        This must be called on the Pod/Worker side in Distributed training because it
        is needed for the model to be able to extract the relevant information from the
        datastructure sent by the Modeller.

        Args:
            datastructure_identifier: The datastructure identifier for the model.
        """
        self._datastructure_identifier = datastructure_identifier

    def _fit_federated(
        self,
        pod_identifiers: List[str],
        username: Optional[str] = None,
        aggregator: Optional[_BaseAggregatorFactory] = None,
        steps_between_parameter_updates: Optional[int] = None,
        epochs_between_parameter_updates: Optional[int] = None,
        bitfounthub: Optional[BitfountHub] = None,
        ms_config: Optional[MessageServiceConfig] = None,
        message_service: Optional[_MessageService] = None,
        pod_public_key_paths: Optional[Mapping[str, Path]] = None,
        pretrained_file: Optional[Union[str, os.PathLike]] = None,
        secure_aggregation: bool = False,
        auto_eval: bool = True,
        identity_verification_method: IdentityVerificationMethod = IdentityVerificationMethod.DEFAULT,  # noqa: B950
        private_key_or_file: Optional[Union[RSAPrivateKey, Path]] = None,
        idp_url: Optional[str] = None,
        model_out: Optional[Union[Path, str]] = None,
        require_all_pods: bool = False,
        extra_imports: Optional[List[str]] = None,
        modeller_checkpointing: bool = True,
        checkpoint_filename: Optional[str] = None,
        project_id: Optional[str] = None,
        run_on_new_data_only: bool = False,
        batched_execution: Optional[bool] = None,
        **kwargs: Any,
    ) -> Optional[Dict[str, str]]:
        """Fits model federated-ly."""
        if not bitfounthub:
            bitfounthub = _create_bitfounthub(username=username)
        if batched_execution is None:
            batched_execution = BITFOUNT_DEFAULT_BATCHED_EXECUTION
        if run_on_new_data_only:
            logger.warning(
                "Running on new data only is only supported for predict method. "
                "Resuming task on the entire datasource."
            )
            run_on_new_data_only = False
        if batched_execution:
            logger.warning(
                "Batched execution is only supported for prediction and evaluation. "
                "Resuming task without batched execution."
            )
            batched_execution = False

        model = self._get_model(extra_imports=extra_imports, hub=bitfounthub, **kwargs)
        algorithm = FederatedModelTraining(
            model=model,
            pretrained_file=pretrained_file,
            modeller_checkpointing=modeller_checkpointing,
            checkpoint_filename=checkpoint_filename,
            project_id=project_id,
        )

        pod_identifiers = _check_and_update_pod_ids(pod_identifiers, bitfounthub)
        if (
            len(pod_identifiers) > 1
            and hasattr(self.datastructure, "table")
            and isinstance(self.datastructure.table, str)
        ):
            raise ValueError(
                "You are attempting to train on multiple pods, but the "
                "provided DataStructure table name is a string. "
                "Please make sure that the `table` argument to the "
                "`DataStructure` is a mapping of Pod names to table names. "
            )
        datastructure_pod_identifiers = self.datastructure.get_pod_identifiers()
        if datastructure_pod_identifiers:
            datastructure_pod_identifiers = _check_and_update_pod_ids(
                datastructure_pod_identifiers, bitfounthub
            )
            self.datastructure._update_datastructure_with_hub_identifiers(
                datastructure_pod_identifiers
            )
        if not message_service:
            message_service = _create_message_service(bitfounthub.session, ms_config)

        if not aggregator:
            aggregator = _create_aggregator(secure_aggregation=secure_aggregation)

        if not idp_url:
            idp_url = _get_idp_url()

        # TODO: [BIT-1098] Manage pods with different schemas
        protocol = FederatedAveraging(
            algorithm=algorithm,
            aggregator=aggregator,
            steps_between_parameter_updates=steps_between_parameter_updates,
            epochs_between_parameter_updates=epochs_between_parameter_updates,
            auto_eval=auto_eval,
        )

        modeller = _Modeller(
            protocol=protocol,
            message_service=message_service,
            bitfounthub=bitfounthub,
            pod_public_key_paths=pod_public_key_paths,
            identity_verification_method=identity_verification_method,
            private_key=private_key_or_file,
            idp_url=idp_url,
        )
        # Run task
        result = modeller.run(
            pod_identifiers,
            require_all_pods=require_all_pods,
            model_out=model_out,
            run_on_new_data_only=run_on_new_data_only,
            batched_execution=batched_execution,
            project_id=project_id,
        )

        if result is False:
            return None
        else:
            return result

    def fit(
        self,
        data: Optional[BaseSource] = None,
        metrics: Optional[Dict[str, Metric]] = None,
        pod_identifiers: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Optional[Dict[str, str]]:
        """Fits model either locally or federated-ly.

        `pod_identifiers` must be provided for federated training.

        :::caution

        For remote training, this method does not use the existing parameters of the
        model unless they are serialized and provided using the `pretrained_file`
        keyword argument.

        :::

        Args:
            data: Datasource for training. Defaults to None.
            metrics: Metrics to calculate for validation. Defaults to None.
            pod_identifiers: List of pod identifiers. If this is provided, the model
                will be trained in a federated manner. Defaults to None.
            **kwargs: Optional keyword arguments passed to the federated fit method.
                Any unrecognized keyword arguments will be interpreted as
                custom model hyperparameters.

        Returns:
            A dictionary of metrics and their values. Optional.

        Raises:
            ValueError: If neither `pod_identifiers` are provided for federated training
                nor `data` is provided for local training.
        """
        if pod_identifiers is not None:
            logger.info(f"Training federated with pods: {', '.join(pod_identifiers)}.")
            return self._fit_federated(pod_identifiers=pod_identifiers, **kwargs)
        elif data is not None:
            logger.info("Training locally using provided datasource.")
            return self._fit_local(data=data, metrics=metrics)

        raise ValueError(
            "Please provide either pod identifiers "
            "or a datasource for model training."
        )

    def evaluate(
        self,
        test_dl: Optional[BitfountDataLoader] = None,
        pod_identifiers: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Union[Tuple[np.ndarray, np.ndarray], Dict[str, float]]:
        """Evaluates model either locally or federated-ly.

        `pod_identifiers` must be provided for federated evaluation.

        :::caution

        For remote evaluation, this method does not use the existing parameters of the
        model unless they are serialized and provided using the `pretrained_file`
        keyword argument.

        :::

        Args:
            test_dl: Optional dataloader to run inference on which takes precedence over
                the dataloader returned by `self.test_dataloader`.
            pod_identifiers: List of pod identifiers. If this is provided, the model
                will be evaluated in a federated manner. Defaults to None.
            **kwargs: Optional keyword arguments passed to the federated evaluate
                method.

        Returns:
            Either:
                - A tuple of numpy arrays containing the predicted and actual values if
                evaluated locally.
                - A dictionary of metrics if evaluated remotely.
        """
        if pod_identifiers is not None:
            logger.info(f"Evaluating model on pods: {', '.join(pod_identifiers)}.")
            return self._evaluate_federated(pod_identifiers=pod_identifiers, **kwargs)

        logger.info("Evaluating model locally")
        return self._evaluate_local(test_dl=test_dl)

    def _evaluate_local(
        self, test_dl: Optional[BitfountDataLoader] = None, **kwargs: Any
    ) -> Tuple[np.ndarray, np.ndarray]:
        """This method runs inference on the test dataloader.

        Args:
            test_dl: Optional dataloader to run inference on which takes precedence over
                the dataloader returned by `self.test_dataloader`.

        Returns:
            A tuple of predictions and targets as numpy arrays.
        """
        raise NotImplementedError

    def _evaluate_federated(
        self,
        pod_identifiers: List[str],
        username: Optional[str] = None,
        aggregator: Optional[_BaseAggregatorFactory] = None,
        bitfounthub: Optional[BitfountHub] = None,
        ms_config: Optional[MessageServiceConfig] = None,
        message_service: Optional[_MessageService] = None,
        pod_public_key_paths: Optional[Mapping[str, Path]] = None,
        pretrained_file: Optional[Union[str, os.PathLike]] = None,
        secure_aggregation: bool = False,
        identity_verification_method: IdentityVerificationMethod = IdentityVerificationMethod.DEFAULT,  # noqa: B950
        private_key_or_file: Optional[Union[RSAPrivateKey, Path]] = None,
        idp_url: Optional[str] = None,
        model_out: Optional[Union[Path, str]] = None,
        require_all_pods: bool = False,
        extra_imports: Optional[List[str]] = None,
        project_id: Optional[str] = None,
        run_on_new_data_only: bool = False,
        batched_execution: Optional[bool] = None,
        **kwargs: Any,
    ) -> Dict[str, float]:
        """Evaluates model remotely on all the specified pods separately.

        Returns:
            A dictionary of metrics.
        """
        if not bitfounthub:
            bitfounthub = _create_bitfounthub(username=username)
        if batched_execution is None:
            batched_execution = BITFOUNT_DEFAULT_BATCHED_EXECUTION
        if run_on_new_data_only:
            logger.warning(
                "Running on new data only is only supported for predict method. "
                "Resuming task on the entire datasource."
            )
            run_on_new_data_only = False
        if len(pod_identifiers) > 1 and batched_execution:
            logger.warning(
                "Batched execution is only supported for single pod evaluation. "
                "Resuming task without batched execution."
            )
            batched_execution = False
        model = self._get_model(extra_imports=extra_imports, hub=bitfounthub, **kwargs)
        algorithm = ModelEvaluation(
            model=model, pretrained_file=pretrained_file, project_id=project_id
        )

        pod_identifiers = _check_and_update_pod_ids(pod_identifiers, bitfounthub)
        if (
            len(pod_identifiers) > 1
            and hasattr(self.datastructure, "table")
            and isinstance(self.datastructure.table, str)
        ):
            raise ValueError(
                "You are attempting to evaluate on multiple pods, but the "
                "provided DataStructure table name is a string. "
                "Please make sure that the `table` argument to the "
                "`DataStructure` is a mapping of Pod names to table names. "
            )
        datastructure_pod_identifiers = self.datastructure.get_pod_identifiers()
        if datastructure_pod_identifiers:
            datastructure_pod_identifiers = _check_and_update_pod_ids(
                datastructure_pod_identifiers, bitfounthub
            )
            self.datastructure._update_datastructure_with_hub_identifiers(
                datastructure_pod_identifiers
            )
        if not message_service:
            message_service = _create_message_service(bitfounthub.session, ms_config)

        if not idp_url:
            idp_url = _get_idp_url()

        protocol = ResultsOnly(
            algorithm=algorithm,
            aggregator=aggregator,
            secure_aggregation=secure_aggregation,
        )

        modeller = _Modeller(
            protocol=protocol,
            message_service=message_service,
            bitfounthub=bitfounthub,
            pod_public_key_paths=pod_public_key_paths,
            identity_verification_method=identity_verification_method,
            private_key=private_key_or_file,
            idp_url=idp_url,
        )

        metrics = modeller.run(
            pod_identifiers,
            require_all_pods=require_all_pods,
            model_out=model_out,
            project_id=project_id,
            batched_execution=batched_execution,
            run_on_new_data_only=run_on_new_data_only,
        )
        return cast(Dict[str, float], metrics)

    def predict(
        self,
        data: Optional[BaseSource] = None,
        pod_identifiers: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Union[List[np.ndarray], Dict[str, List[np.ndarray]]]:
        """Infers model either locally or federated-ly.

        `pod_identifiers` must be provided for federated inference.

        :::caution

        For remote inference, this method does not use the existing parameters of the
        model unless they are serialized and provided using the `pretrained_file`
        keyword argument.

        :::

        Args:
            data: Optional datasource to run inference on if training locally. Defaults
                to None.
            pod_identifiers: List of pod identifiers. If this is provided, the model
                will be inferred in a federated manner. Defaults to None.
            **kwargs: Optional keyword arguments passed to the federated predict
                method.

        Returns:
            Either:
                - A numpy array containing the predicted values if
                inferred locally.
                - A dictionary of numpy arrays if inferred remotely.
        """
        # Prediction only happens on one epoch, and not dependent
        # on the number of steps. To ensure that the tasks are
        # mapped to the same task hash in the database, we modify
        # the hyperparmeters here, since it does not
        if self.epochs:
            self.epochs = 1
        else:
            self.steps = 1
        if data and pod_identifiers:
            raise ValueError(
                "Only one of 'data' and 'pod_identifiers' can be provided. "
                "Received both."
            )
        elif pod_identifiers:
            logger.info(f"Inferring model on pods: {', '.join(pod_identifiers)}.")
            return self._predict_federated(pod_identifiers=pod_identifiers, **kwargs)
        elif data:
            logger.info("Inferring model locally")
            return self._predict_local(data)

        raise ValueError(
            "One of 'data' or 'pod_identifiers' must be provided. Received neither."
        )

    def _predict_local(self, data: BaseSource, **kwargs: Any) -> List[np.ndarray]:
        """This method runs inference on the datasource.

        Args:
            data: DataSource to run inference on.

        Returns:
            Predictions (or whatever the model outputs) as a numpy array.
        """
        raise NotImplementedError

    def _predict_federated(
        self,
        pod_identifiers: List[str],
        username: Optional[str] = None,
        aggregator: Optional[_BaseAggregatorFactory] = None,
        bitfounthub: Optional[BitfountHub] = None,
        ms_config: Optional[MessageServiceConfig] = None,
        message_service: Optional[_MessageService] = None,
        pod_public_key_paths: Optional[Mapping[str, Path]] = None,
        pretrained_file: Optional[Union[str, os.PathLike]] = None,
        secure_aggregation: bool = False,
        identity_verification_method: IdentityVerificationMethod = IdentityVerificationMethod.DEFAULT,  # noqa: B950
        private_key_or_file: Optional[Union[RSAPrivateKey, Path]] = None,
        idp_url: Optional[str] = None,
        model_out: Optional[Union[Path, str]] = None,
        require_all_pods: bool = False,
        extra_imports: Optional[List[str]] = None,
        project_id: Optional[str] = None,
        run_on_new_data_only: bool = False,
        batched_execution: Optional[bool] = None,
        **kwargs: Any,
    ) -> Dict[str, List[np.ndarray]]:
        """Infers the model remotely on all the specified pods separately.

        Returns:
            A dictionary of predictions.
        """
        if not bitfounthub:
            bitfounthub = _create_bitfounthub(username=username)
        if batched_execution is None:
            batched_execution = BITFOUNT_DEFAULT_BATCHED_EXECUTION
        if len(pod_identifiers) > 1 and batched_execution:
            logger.warning(
                "Batched execution is only supported for single pod prediction. "
                "Resuming task without batched execution."
            )
            batched_execution = False

        model = self._get_model(extra_imports=extra_imports, hub=bitfounthub, **kwargs)
        algorithm = ModelInference(
            model=model, pretrained_file=pretrained_file, project_id=project_id
        )

        pod_identifiers = _check_and_update_pod_ids(pod_identifiers, bitfounthub)
        if (
            len(pod_identifiers) > 1
            and hasattr(self.datastructure, "table")
            and isinstance(self.datastructure.table, str)
        ):
            raise ValueError(
                "You are attempting to infer on multiple pods, but the "
                "provided DataStructure table name is a string. "
                "Please make sure that the `table` argument to the "
                "`DataStructure` is a mapping of Pod names to table names. "
            )
        datastructure_pod_identifiers = self.datastructure.get_pod_identifiers()
        if datastructure_pod_identifiers:
            datastructure_pod_identifiers = _check_and_update_pod_ids(
                datastructure_pod_identifiers, bitfounthub
            )
            self.datastructure._update_datastructure_with_hub_identifiers(
                datastructure_pod_identifiers
            )

        if not message_service:
            message_service = _create_message_service(bitfounthub.session, ms_config)

        if not idp_url:
            idp_url = _get_idp_url()

        protocol = ResultsOnly(
            algorithm=algorithm,
            aggregator=aggregator,
            secure_aggregation=secure_aggregation,
        )

        modeller = _Modeller(
            protocol=protocol,
            message_service=message_service,
            bitfounthub=bitfounthub,
            pod_public_key_paths=pod_public_key_paths,
            identity_verification_method=identity_verification_method,
            private_key=private_key_or_file,
            idp_url=idp_url,
        )

        predictions = modeller.run(
            pod_identifiers,
            require_all_pods=require_all_pods,
            model_out=model_out,
            project_id=project_id,
            run_on_new_data_only=run_on_new_data_only,
            batched_execution=batched_execution,
        )
        return cast(Dict[str, List[np.ndarray]], predictions)


class _ModellessAlgorithmMixIn(ABC):
    """A mixin for SQL Algorithms.

    An abstract base mixin for SQL algorithms that are compatible with the
    ResultsOnly protocol.
    """

    def execute(
        self,
        pod_identifiers: List[str],
        username: Optional[str] = None,
        bitfounthub: Optional[BitfountHub] = None,
        ms_config: Optional[MessageServiceConfig] = None,
        message_service: Optional[_MessageService] = None,
        pod_public_key_paths: Optional[Mapping[str, Path]] = None,
        identity_verification_method: IdentityVerificationMethod = IdentityVerificationMethod.DEFAULT,  # noqa: B950
        private_key_or_file: Optional[Union[RSAPrivateKey, Path]] = None,
        idp_url: Optional[str] = None,
        require_all_pods: bool = False,
        aggregator: Optional[_BaseAggregatorFactory] = None,
        project_id: Optional[str] = None,
    ) -> List[pd.DataFrame]:
        """Execute ResultsOnly compatible algorithm.

        Syntactic sugar to allow the modeller to call `.execute(...)` on
        ResultsOnly compatible algorithms.
        """
        if not bitfounthub:
            bitfounthub = _create_bitfounthub(username=username)

        pod_identifiers = _check_and_update_pod_ids(pod_identifiers, bitfounthub)

        if not message_service:
            message_service = _create_message_service(bitfounthub.session, ms_config)

        if not idp_url:
            idp_url = _get_idp_url()

        protocol = ResultsOnly(
            algorithm=cast(_ResultsOnlyCompatibleAlgoFactory_, self),
            aggregator=aggregator,
        )
        modeller = _Modeller(
            protocol=protocol,
            message_service=message_service,
            bitfounthub=bitfounthub,
            pod_public_key_paths=pod_public_key_paths,
            identity_verification_method=identity_verification_method,
            private_key=private_key_or_file,
            idp_url=idp_url,
        )

        result = cast(
            List[pd.DataFrame],
            modeller.run(
                pod_identifiers=pod_identifiers,
                require_all_pods=require_all_pods,
                project_id=project_id,
            ),
        )
        return result


class _PSIAlgorithmsMixIn(ABC):
    """A mixin for PSI Algorithms.

    An abstract base mixin for algorithms that are compatible with the
    PSI protocol.
    """

    def execute(
        self: _PSICompatibleAlgoFactory_,
        pod_identifiers: List[str],
        datasource: BaseSource,
        datasource_columns: Optional[List[str]] = None,
        datasource_table: Optional[str] = None,
        pod_columns: Optional[List[str]] = None,
        pod_table: Optional[str] = None,
        username: Optional[str] = None,
        bitfounthub: Optional[BitfountHub] = None,
        ms_config: Optional[MessageServiceConfig] = None,
        message_service: Optional[_MessageService] = None,
        pod_public_key_paths: Optional[Mapping[str, Path]] = None,
        identity_verification_method: IdentityVerificationMethod = IdentityVerificationMethod.DEFAULT,  # noqa: B950
        private_key_or_file: Optional[Union[RSAPrivateKey, Path]] = None,
        idp_url: Optional[str] = None,
        require_all_pods: bool = False,
        project_id: Optional[str] = None,
    ) -> List[pd.DataFrame]:
        """Execute PSI compatible algorithm.

        Syntactic sugar to allow the modeller to call `.intersect(...)` on
        PrivateSetIntersection compatible algorithms.

        Args:
            pod_identifiers: The pod identifier for the Private Set
                Intersection as a list.
            datasource: The modeller datasource on which the Private
                Set Intersection will be performed on.
            datasource_columns: The modeller's columns from their datasource
                on which the private set intersection will be computed as a
                list of strings. Defaults to None.
            datasource_table: The modeller's table from their datasource,
                if the datasource is multitable, on which the private set
                intersection will be computed as a string. Defaults to None.
            pod_columns: The pod's columns from their datasource
                on which the private set intersection will be computed as a
                list of strings. Defaults to None.
            pod_table: The pod's table from their datasource,
                if the datasource is multitable, on which the private set
                intersection will be computed as a string. Defaults to None.
            username: The modeller's username. Defaults to None.
            bitfounthub: The bifount hub instance. Defaults to None.
            ms_config: Configuration for the message service. Defaults to None.
            message_service: The message service to use. Defaults to None.
            pod_public_key_paths: The path for the pod public key. Used
                when authentication with the pod is done with the public key.
                Defaults to None.
            identity_verification_method: The identity verification method
                to use. Defaults to None.
            private_key_or_file: The private key or path to the private
                key. Defaults to None.
            idp_url: The identity provider url. Defaults to None.
            require_all_pods: Whether all pods are needed for the algorithm.
                Defaults to False.

        Returns:
            The records from the modeller dataset which were found in the intersection.

        Raises:
            PSIMultiplePodsError: If `execute` is called on multiple pods.
        """
        if len(pod_identifiers) > 1:
            raise PSIMultiplePodsError(
                "You are trying to run a PrivateSetIntersection "
                "task on multiple pods. We support "
                "PrivateSetIntersection on a single pod."
            )
        if not bitfounthub:
            bitfounthub = _create_bitfounthub(username=username)

        pod_identifiers = _check_and_update_pod_ids(pod_identifiers, bitfounthub)

        if not message_service:
            message_service = _create_message_service(bitfounthub.session, ms_config)

        if not idp_url:
            idp_url = _get_idp_url()

        protocol = PrivateSetIntersection(
            algorithm=self,
            datasource=datasource,
            datasource_columns=datasource_columns,
            datasource_table=datasource_table,
            pod_columns=pod_columns,
            pod_table=pod_table,
        )
        modeller = _Modeller(
            protocol=protocol,
            message_service=message_service,
            bitfounthub=bitfounthub,
            pod_public_key_paths=pod_public_key_paths,
            identity_verification_method=identity_verification_method,
            private_key=private_key_or_file,
            idp_url=idp_url,
        )
        result = cast(
            List[pd.DataFrame],
            modeller.run(
                pod_identifiers=pod_identifiers,
                require_all_pods=require_all_pods,
                project_id=project_id,
            ),
        )
        return result
