#! /usr/bin/env python3

"""A fair division instance generator."""


#############################################################################
# imports ###################################################################

import sys
import getopt
import textwrap
import time
import random
import numpy as np
from itertools import combinations

#############################################################################
# command line ##############################################################

DEFAULTS = {
    'nb_agents':   3,
    'nb_objects':  10,
    'max_weight':  100,
    'seed':        int(time.time() * 1000),
    'normalized':  False
}


#############################################################################
# helper classes ############################################################

class SPUPGenerator():
    """The random sample Single-Peaked Impartial Culture."""

    def __init__(self, nb_candidates, nb_voters, nb_samples):
        self.nb_candidates = nb_candidates
        self.nb_voters = nb_voters
        self.nb_samples = nb_samples

    def profiles(self):
        """Generates a given number of random profiles. If random.randint
        is correct, the probability distribution corresponds to
        impartial culture for Single-Peaked preferences."""
        for _ in range(self.nb_samples):
            yield [self._profile_gen([-1] * self.nb_candidates)
                   for _ in range(self.nb_voters)]

    def _profile_gen(self, order):
        first = random.randint(0, self.nb_candidates - 1)
        order[0] = first
        return self._profile_gen_aux(first - 1, first + 1, order, 1)

    def _profile_gen_aux(self, left, right, order, rank):
        if left < 0:
            return order[:rank] + list(range(right, self.nb_candidates))
        else:
            if right >= self.nb_candidates:
                return order[:rank] + list(range(left, -1, -1))
            else:
                if random.randint(0, 1):
                    order[rank] = left
                    return self._profile_gen_aux(left - 1, right, order,
                                                 rank + 1)
                else:
                    order[rank] = right
                    return self._profile_gen_aux(left, right + 1, order,
                                                 rank + 1)

    def nb_profiles(self):
        """Returns the number of profiles for this culture."""
        return self.nb_samples


#############################################################################
# generation ################################################################

def generate_uniform(nb_agents, nb_objects, max_weight,
                     normalized=True, strict=True):
    """Generates a new instance using the uniform model. Weights
    are randomly chosen between 0 and max_weight. If normalized
    is True, then all the rows have the same sum. If strict is True,
    then all the agents have strict preferences"""
    matrix = np.zeros((nb_agents, nb_objects), dtype=np.float)

    for agent in range(0, nb_agents):
        for obj in range(0, nb_objects):
            matrix[agent][obj] = random.random()

    if normalized:
        matrix = matrix / matrix.sum(axis=1)[:, np.newaxis]

    matrix = matrix * (max_weight / matrix.max())
    matrix = matrix.astype(int, copy=False)

    if normalized:
        # It could be the case that the rounding breaks the
        # normalization assumption, in which case we have
        # to correct that
        sums = matrix.sum(axis=1)
        max_sum = sums.max()
        sums = max_sum - sums
        # The sums matrix represents the difference between
        # the sum of the maximum row and the sum of the current
        # row. If it is strictly positive, then we add 1 point
        # of utility to a random object until the sum of utilities
        # for this row is equal to the max sum.
        for agent, diff in enumerate(sums):
            while diff:
                obj = random.randrange(nb_objects)
                if matrix[agent][obj] < max_weight:
                    matrix[agent][obj] += 1
                    diff -= 1
    if strict:
        # This could break normalization assumption...
        for agent in range(nb_agents):
            free_weights = set(range(max_weight)) - set(matrix[agent])
            for o1, o2 in combinations(range(nb_objects), 2):
                if matrix[agent][o1] == matrix[agent][o2]:
                    matrix[agent][o2] = free_weights.pop()

    return matrix

def generate_single_peaked(nb_agents, nb_objects, max_weight,
                           normalized=True):
    """Generates a new instance using the uniform single peaked
    model. Weights are randomly chosen between 0 and max_weight.
    If normalized is True, then all the rows have the same sum."""
    matrix = generate_uniform(nb_agents, nb_objects, max_weight,
                              normalized)
    matrix = [sorted(pref_ag, reverse=True) for pref_ag in matrix]
    profile = next(SPUPGenerator(nb_objects, nb_agents, 1).profiles())
    matrix = [[matrix[ag][profile[ag][obj]] for obj in range(nb_objects)]
              for ag in range(nb_agents)]
    return np.array(matrix)


def _exit_usage(name, message=None, code=0):
    if message:
        sys.stderr.write('%s\n' % message)
    sys.stderr.write(textwrap.dedent('''\
Usage: %(name)s [-hA:O:W:s] [filename]
    -h  --help             print this help message then exit
    -A  --nb_agents <a>    number of agents (defaults to %(nb_agents)r)
    -O  --nb_objects <o>   number of objects (defaults to %(nb_objects)r)
    -W  --max_weight <w>   maximum weight (defaults to %(max_weight)r)
    -s  --seed <s>         seed of the random generator
                           (defaults to current system time)
    -n  --normalized       normalized (defaults to %(normalized)r)
    ''') % dict(name=name, **DEFAULTS))
    sys.exit(code)


def _main():
    name, *args = sys.argv
    try:
        options, args = getopt.getopt(args, 'hA:O:W:s',
                                      ['help', 'nb_agents=',
                                       'nb_objects=', 'max_weight=',
                                       'seed=', 'normalized'])
    except getopt.GetoptError as message:
        _exit_usage(name, message, 1)

    nb_agents = DEFAULTS['nb_agents']
    nb_objects = DEFAULTS['nb_objects']
    max_weight = DEFAULTS['max_weight']
    seed = DEFAULTS['seed']
    normalized = DEFAULTS['normalized']

    for opt, value in options:
        if opt in ['-h', '--help']:
            _exit_usage(name)
        elif opt in ['-A', '--nb_agents']:
            nb_agents = int(value)
        elif opt in ['-O', '--nb_objects']:
            nb_objects = int(value)
        elif opt in ['-W', '--max_weights']:
            max_weight = int(value)
        elif opt in ['-s', '--seed']:
            seed = int(value)
        elif opt in ['-normalized', '--normalized']:
            normalized = True
    random.seed(seed)
    matrix = generate_uniform(nb_agents, nb_objects,
                              max_weight, normalized)
    print('\n'.join(' '.join(str(int(x)) for x in row) for row in matrix))


if __name__ == '__main__':
    _main()
