# vim: set fileencoding=utf-8 :

from gurobipy import *

class ILPUnsolvable (Exception):
    pass

class ILPScheduler:
    def __init__ (self, data):
        """
        :param data: WorkerSlotModel
        """
        self.data = data
        self.model = Model("scheduling")

    def optimize (self):
        """
        Run optimization pass
        """
        data = self.data
        slotNames = data.slotNames
        workerNames = data.workerNames

        availability = { (w, s) : data.available (w, s) for w in workerNames for s in slotNames }

        model = self.model

        # x[(w,s)] == 1 if worker w is assigned to slot s
        x = model.addVars(availability.keys(), ub=availability, vtype=GRB.BINARY, name='x')

        # every worker needs slots
        slotsPerWorker = model.addVars (workerNames, vtype=GRB.INTEGER, name='slotsPerWorker')
        model.addConstrs((x.sum(w, '*') == slotsPerWorker[w] for w in workerNames), name='slotsPerWorkerConstr')

        model.addConstrs((slotsPerWorker[w] >= data.minSlots (w) for w in workerNames), name='minSlotsPerWorkerConstr')
        model.addConstrs((slotsPerWorker[w] <= data.maxSlots (w) for w in workerNames), name='maxSlotsPerWorkerConstr')

        # but the number of worker per slot is limited
        workerPerSlot = model.addVars (slotNames, vtype=GRB.INTEGER, name='workerPerSlot')
        model.addConstrs ((x.sum('*', s) == workerPerSlot[s] for s in slotNames), name='workerPerSlotAssign')

        model.addConstrs ((workerPerSlot[s] <= data.maxWorker (s) for s in slotNames), name='maxWorkerPerSlotConstr')
        model.addConstrs ((workerPerSlot[s] >= data.minWorker (s) for s in slotNames), name='minWorkerPerSlotConstr')

        # also the number of slots is limited
        slotUsed = model.addVars (slotNames, vtype=GRB.BINARY, name='slotUsed')
        for s in slotNames:
            # slotUsed[s] is True iff there is at least one worker per slot
            model.addGenConstrOr (slotUsed[s], [x[(w, s)] for w in workerNames], name='slotUsedConstr')
        totalSlotsUsed = model.addVar (vtype=GRB.INTEGER)
        model.addConstr (totalSlotsUsed == slotUsed.sum ('*'), name='totalSlotsUsedConstr')

        model.addConstr(totalSlotsUsed >= data.minTotalSlots, name='minTotalSlotsUsedConstr')
        model.addConstr(totalSlotsUsed <= data.maxTotalSlots, name='maxTotalSlotsUsedConstr')

        # optimize for these constraints, more important first
        # - similar-sized slots: minimize diff of each workerPerSlot and
        # maxWorkerPerSlot. minimizing maxWorkerPerSlot is not sufficient,
        # because solution 4/4/1 and 3/4/2 are the same
        maxWorkerPerSlot = model.addVar (vtype=GRB.INTEGER)
        # maxWorkerPerSlot=max(workerPerSlot)
        model.addGenConstrMax (maxWorkerPerSlot, workerPerSlot, name='maxWorkerPerSlotConstr')
        workerMaxDiff = model.addVars (slotNames, vtype=GRB.INTEGER, name='workerMaxDiff')
        model.addConstrs ((workerMaxDiff[s] == maxWorkerPerSlot-workerPerSlot[s] for s in slotNames), name='workerMaxDiffConstr')
        cost1 = model.addVar (vtype=GRB.INTEGER)
        model.addGenConstrMax (cost1, workerMaxDiff, name='cost1Constr')

        # - prefer solutions that fulfil most people’s wishes (everyone gets what he
        # wants) by minimizing sum of priority.
        totalPriority = model.addVar (vtype=GRB.INTEGER)
        maxTotalPriority = 0
        terms = []
        for w in workerNames:
            for s in slotNames:
                p = data.priority (w, s)
                if p is not None:
                    terms.append (x[w,s]*p)
                    maxTotalPriority += p
        model.addConstr (totalPriority == quicksum(terms), 'totalPriority')

        # - try to keep worker pairs together
#        pairs = [(self.workers[0], self.workers[1])]
#        terms = []
#        for a, b in pairs:
#            for j, s in enumerate (slotNames):
#                terms.append (
        # XXX: todo

        # combine everything into a single cost function, because multi-objective
        # solving sucks
        totalCost = model.addVar (vtype=GRB.INTEGER)
        # worker per slot is more important, so multiply with max sum of priority
        model.addConstr (totalCost == cost1*maxTotalPriority + totalPriority, 'totalCost')
        model.ModelSense = GRB.MINIMIZE
        model.setObjective (totalCost)

        model.optimize ()
        status = model.Status
        if status in (GRB.Status.INF_OR_UNBD, GRB.Status.INFEASIBLE, GRB.Status.UNBOUNDED):
            raise ILPUnsolvable ()
        if status != GRB.Status.OPTIMAL:
            print ('not optimal')

        print ('total priority is', totalPriority.X)
        print ('max worker per slot is', maxWorkerPerSlot.X)
        print ('cost1 is', cost1.X)
        print ('total cost is', totalCost.X)
        for s in slotNames:
            print ('slotUsed', s, slotUsed[s].X)
        print ('totalSlotsUsed', totalSlotsUsed.X)
        return x, workerPerSlot

def firstNonNull (l):
    for i in range (len (l)):
        if l[i]:
            return i

def mainOptimize (args):
    import sys
    from tabulate import tabulate
    from model import WorkerSlotModel

    # read data
    with open (args.file, 'r') as fd:
        data = WorkerSlotModel.fromYaml (fd)

    sched = ILPScheduler (data)
    try:
        x, groupSize = sched.optimize ()
    except ILPUnsolvable:
        print ('unsolvable')
        return

    tbl = []
    for w in data.workerNames:
        r = [w]
        for s in data.slotNames:
            if groupSize[s].X > 0:
                if x[w,s].X == 1:
                    r.append (data.priority (w, s))
                else:
                    r.append ('')
        tbl.append (r)
    # skip row description, thus [1:]
    tbl.sort (key=lambda x: firstNonNull (x[1:]))
    hdr = []
    for s in data.slotNames:
        if groupSize[s].X > 0:
            hdr.append ('{} ({})'.format (s, int (groupSize[s].X)))

    print (tabulate (tbl, headers=hdr))

    if args.output:
        for s in data.slotNames:
            if groupSize[s].X > 0:
                with open (os.path.join (args.output, s + '.txt'), 'w') as fd:
                    for w in data.workerNames:
                        if x[w,s].X == 1:
                            fd.write ('{}\n'.format (w))

def addParser (subparsers):
    parser = subparsers.add_parser ('optimize', help='Optimize schedule')
    parser.add_argument('-o', '--output', help='Output directory for mapping files')
    parser.add_argument('file', help='input file')
    parser.set_defaults (func=mainOptimize)