
"""
Core classes and functions for multiagent resource allocation.
"""

# Imports ###################################################################

import itertools
import numpy as np
import solving


# Helper classes ###########################################################

class Instance():
    """
    An instance of the additive multiagent resource allocation problem.
    """
    def __init__(self, w, agents, objects):
        """
        Initializes a MARA instance using a matrix of weights. A numpy
        matrix is expected as argument, as well as the agent names and
        the object names.
        """
        self.w = w
        self.nb_agents = w.shape[0]
        self.nb_objects = w.shape[1]
        self.agents = agents
        self.objects = objects
        self.num_agents = {agents[i]: i for i in range(self.nb_agents)}
        self.num_objects = {objects[i]: i for i in range(self.nb_objects)}

    def allocations(self):
        """
        Iterator on the set of allocations. Just use an allocation and
        modifies it in place!
        """
        yield from self._allocation_aux(Allocation(self), 0)

    def _allocation_aux(self, allocation, obj):
        if obj == self.nb_objects:
            yield allocation
        else:
            for agent in range(self.nb_agents):
                allocation.give(agent, obj)
                yield from self._allocation_aux(allocation, obj + 1)
                allocation.take_back(agent, obj)

    def u_max(self, agent):
        """Returns the utility of agent a for the whole cake."""
        return sum(self.w[agent])

    def w_max(self, agent, subset=None):
        """Returns the maximum utility of agent a."""
        if subset is None:
            subset = [True] * self.nb_objects
        return max(self.w[agent][obj] for obj in range(self.nb_objects)
                   if subset[obj])

    def best_obj(self, agent, subset):
        """Returns the best object for agent a amongst a given subset."""
        subset2 = subset
        subset = [False] * self.nb_objects
        if type(agent) != int:
            agent = self.num_agents[agent]
        for obj in subset2:
            if type(obj) != int:
                obj = self.num_objects[obj]
            subset[obj] = True
        w_max = self.w_max(agent, subset)
        return [self.objects[obj] for obj in range(self.nb_objects)
                if self.w[agent][obj] == w_max][0]

    def best(self, agent, subset=None):
        """Returns the best objects for agent a amongst a given subset."""
        if subset is None:
            subset = [True] * self.nb_objects
        if type(subset) == set:
            subset2 = subset
            subset = [False] * self.nb_objects
            for obj in subset2:
                if type(obj) != int:
                    obj = self.num_objects[obj]
                subset[obj] = True
        w_max = self.w_max(agent, subset)
        return [obj for obj in range(self.nb_objects)
                if self.w[agent][obj] == w_max]

    def utility(self, agent, share):
        """Returns the utility of agent a for a given share."""
        return np.inner(self.w[agent], share)

    def weight(self, agent, obj):
        if type(agent) != int:
            agent = self.num_agents[agent]
        if type(obj) != int:
            obj = self.num_objects[obj]
        return self.w[agent][obj]

    def _pad(self, s, nb):
        if len(s) > nb:
            return s[:nb - 1] + "."
        return s + (" " * (nb - len(s)))

    def __str__(self):
        return (
            " " * 11 + " ".join(self._pad(o, 4) for o in self.objects) + "\n"
            + " " * 11 + (self._pad("↓", 4) + " ") * self.nb_objects + "\n"
            + "\n".join(
                "{} → {}".format(
                    self._pad(self.agents[i], 8),
                    " ".join(self._pad(str(u), 4) for u in self.w[i])
                )
                for i in range(self.nb_agents))
        )


class Allocation():
    """
    Represents a particular allocation for a MARA instance.
    """
    def __init__(self, instance, pi=None):
        """
        Initializes an allocation. The instance defines the number of agents
        and the number of objects. A boolean matrix is expected. If no matrix
        is provided, then the allocation is initialized to the empty one.
        """
        self.instance = instance
        self.nb_agents = instance.nb_agents
        self.nb_objects = instance.nb_objects
        self.pi = (np.zeros((self.nb_agents, self.nb_objects), dtype=np.bool)
                   if pi is None else pi)

    def give(self, agent, obj):
        """
        Gives an object to an agent. Precondition: the object
        should be unallocated (otherwise the preemption constraint will not
        be satisfied).
        """
        if type(agent) != int:
            agent = self.instance.num_agents[agent]
        if type(obj) != int:
            obj = self.instance.num_objects[obj]
        self.pi[agent][obj] = True

    def take_back(self, agent, obj):
        """
        Takes back an object from the agent's share. If the object
        indeed was in this agent's share, it becomes unallocated (provided
        that the premption constraint was satisfied).
        """
        if type(agent) != int:
            agent = self.instance.num_agents[agent]
        if type(obj) != int:
            obj = self.instance.num_objects[obj]
        self.pi[agent][obj] = False

    def objects(self, agent):
        """Returns the set of objects allocated to the agent."""
        if type(agent) != int:
            agent = self.instance.num_agents[agent]
        return (obj for obj in range(self.nb_objects)
                if self.pi[agent][obj])

    def utility_vect(self):
        """Returns the utility vector."""
        return np.diagonal(np.inner(self.instance.w, self.pi))

    def envy_vector(self):
        """Returns the vector of envy."""
        utilities = np.inner(self.instance.w, self.pi)
        envy = utilities - np.diagonal(utilities).reshape(-1, 1)
        return envy.max(0)

    def sequence(self):
        """Returns a sequence generating the allocation, or None
        if no such sequence exists."""
        seq, subset = [], [True] * self.nb_objects
        for _ in range(self.nb_objects):
            flag = False
            for agent in range(self.nb_agents):
                best_objects = self.instance.best(agent, subset)
                intersection = [obj for obj in best_objects
                                if self.pi[agent][obj]]
                if intersection:
                    flag = True
                    seq.append(agent)
                    subset[intersection[0]] = False
                    break
            if not flag:
                return None
        return seq

    # Efficiency tests ######################################################

    def is_pareto_dominated_by(self, other):
        """Returns true iff self → other is a Pareto improvement"""
        u_1 = self.utility_vect()
        u_2 = other.utility_vect()
        strict_improvement = False
        for agent in range(self.nb_agents):
            if u_1[agent] > u_2[agent]:
                return False
            elif u_1[agent] < u_2[agent]:
                strict_improvement = True
        return strict_improvement

    def is_pareto_optimal(self, method="LP"):
        """
        Returns true iff the allocation is Pareto-optimal.

        Keyword args:
        method -- "LP" if the solver should use an external LP solver (default)
        """
        if method == "LP":
            return solving.is_pareto_optimal(self)
        return solving.find_solution(
            self.instance,
            self.is_pareto_dominated_by
        ) is None

    def is_sequenceable(self):
        """Returns true iff the allocation is sequenceable"""
        return self.sequence() is not None

    def is_swap_optimal(self):
        """Returns true iff the allocation is strictly improving swap-optimal."""
        for agent1, agent2 in itertools.combinations(range(self.nb_agents), 2):
            for obj1, obj2 in itertools.product(self.objects(agent1),
                                                self.objects(agent2)):
                if (
                        self.instance.w[agent1][obj2] > self.instance.w[agent1][obj1] and
                        self.instance.w[agent2][obj1] > self.instance.w[agent2][obj2]
                ):
                    return False
        return True

    # Fairness tests ##########################################################

    def _all_utilities_greater_than(self, threshold):
        u_i = self.utility_vect()
        for agent in range(self.nb_agents):
            if u_i[agent] < threshold[agent]:
                return False
        return True

    def is_max_min_share(self):
        """Returns true iff the allocation satisfies the max-min share
        criterion."""
        return self._all_utilities_greater_than(
            [solving.max_min_share(self.instance, agent)
             for agent in range(self.nb_agents)])

    def is_proportional(self):
        """Returns true iff the allocation is proportional."""
        return self._all_utilities_greater_than(
            [self.instance.u_max(agent) / self.nb_agents
             for agent in range(self.nb_agents)])

    def is_min_max_share(self):
        """Returns true iff the allocation satisfies the min-max share
        criterion."""
        return self._all_utilities_greater_than(
            [solving.min_max_share(self.instance, agent)
             for agent in range(self.nb_agents)])

    def is_envy_free(self):
        """Returns true iff the allocation is envy-free."""
        envy_max = self.envy_vector().max()
        return True if envy_max == 0 else False

    def is_ceei(self):
        """Returns true iff the allocation is a Competitive
        Equilibrium from Equal Incomes."""
        return solving.compute_ceei_prices(self) is not None

    def __str__(self):
        return ('\n'.join(
            self.instance.agents[agent] + ' reçoit {' +
            ', '.join(self.instance.objects[obj]
                      for obj in self.objects(agent)) + '}'
            for agent in range(self.nb_agents)))
