'''
    File: num_evaluation_method.py
    Author: Wes Holliday (wesholliday@berkeley.edu) and Eric Pacuit (epacuit@umd.edu)
    Date: June 2, 2024
    
    The NumEvaluationMethod class and helper functions for numerical evaluation methods.
'''

import functools
import numpy as np
from pref_voting.mappings import Utility
from pref_voting.rankings import Ranking

class NumEvaluationMethod(object): 
    """
    A class to add functionality to numerical evaluation methods 

    Args:
        nem (function): An implementation of a numerical evaluation method. The function should accept any type of profile, and a keyword parameter ``curr_cands`` to find the numerical evaluation after restricting to ``curr_cands``. 
        name (string): The human-readable name of the numerical evaluation function.
    
    Returns: 
        A utility function (Utility) that represents the numerical evaluation of each candidate.

    """
    def __init__(self, nem, name = None): 
        
        self.nem = nem
        self.name = name
        functools.update_wrapper(self, nem)   

    def __call__(self, edata, curr_cands = None, **kwargs):

        if (curr_cands is not None and len(curr_cands) == 0) or len(edata.candidates) == 0: 
            return Utility({})
        return self.nem(edata, curr_cands = curr_cands, **kwargs)
        
    def ranking(self, edata, curr_cands = None, **kwargs): 
        """
        Return the ranking generated by the numerical evaluation. 
        """

        if (curr_cands is not None and len(curr_cands) == 0) or len(edata.candidates) == 0: 
            return Ranking({})
        ev =  self.nem(edata, curr_cands = curr_cands, **kwargs)
        return ev.ranking()
        
    def display(self, edata, curr_cands = None, cmap = None, **kwargs): 
        """
        Display the winning set of candidates.
        """
 
        cmap = cmap if cmap is not None else edata.cmap

        ev = self.__call__(edata, curr_cands = curr_cands, **kwargs)

        if ev is None:  # some voting methods may return None if, for instance, it is taking long to compute the winner.
            print(f"{self.name} numerical evaluation is not available")
        else: 
            w_str = f"{self.name} evaluation is " 
            print(w_str + str(ev))
        
    def set_name(self, new_name):
        """Set the name of the social welfare function."""

        self.name = new_name

    def __str__(self): 
        return f"{self.name}"

def nem(name = None):
    """
    A decorator used when creating a numerical evaluation method. 
    """
    def wrapper(f):
        return NumEvaluationMethod(f, name=name)
    return wrapper
