from ml4ir.base.model.losses.loss_base import RelevanceLossBase
import ml4ir.applications.ranking.model.losses.pointwise_losses as pointwise_losses
import ml4ir.applications.ranking.model.losses.pairwise_losses as pairwise_losses
import ml4ir.applications.ranking.model.losses.listwise_losses as listwise_losses
from ml4ir.applications.ranking.config.keys import LossKey


def get_loss(loss_key: str,
             scoring_type: str,
             output_name: str = "score") -> RelevanceLossBase:
    """
    Factor method to get Loss function object

    Parameters
    ----------
    loss_key : str
        Name of the loss function as specified by LossKey
    scoring_type : str
        Type of scoring function - pointwise, pairwise, groupwise
    output_name: str
        Name of the output node for the predicted scores

    Returns
    -------
    RelevanceLossBase
        RelevanceLossBase object that applies the final activation layer
        and computes the loss function from the model score
    """
    if loss_key == LossKey.SIGMOID_CROSS_ENTROPY:
        return pointwise_losses.SigmoidCrossEntropy(loss_key=loss_key,
                                                    scoring_type=scoring_type,
                                                    output_name=output_name)
    elif loss_key == LossKey.RANK_ONE_LISTNET:
        return listwise_losses.RankOneListNet(loss_key=loss_key,
                                              scoring_type=scoring_type,
                                              output_name=output_name)
    elif loss_key == LossKey.SOFTMAX_CROSS_ENTROPY:
        return listwise_losses.SoftmaxCrossEntropy(loss_key=loss_key,
                                                   scoring_type=scoring_type,
                                                   output_name=output_name)
    elif loss_key == LossKey.AUX_ONE_HOT_CROSS_ENTROPY:
        return listwise_losses.AuxiliaryOneHotCrossEntropy(loss_key=loss_key,
                                                           scoring_type=scoring_type,
                                                           output_name=output_name)
    elif loss_key == LossKey.AUX_SOFTMAX_CROSS_ENTROPY:
        return listwise_losses.AuxiliarySoftmaxCrossEntropy(loss_key=loss_key,
                                                            scoring_type=scoring_type,
                                                            output_name=output_name)
    else:
        raise NotImplementedError