#!/bin/python
import glpk
import os
import re
import toml
import sys
import math
import random

random.seed(1409)
# 1 input parameter, the input filename

if len(sys.argv) < 2:
    print("Usage: assignment.py <configuration.toml>")
    sys.exit(-1)

with open(sys.argv[1]) as f:
    tm = toml.load(f)

if not ("students" in tm):
    print(f"Missing table students with preferences in {sys.argv[1]}, e.g.\n[students]\n# Student a@abc.de preferes topic 2 over 6 over 4.\n\"a@abc.de\": [2,6,4]")
    sys.exit(-1)

students = tm['students']
print("#Students = ",len(students))
m = 0
mlen = 0
for s in students.keys():
    m = max(m,max(students[s])+1)
    mlen = max(mlen, len(students[s])-1)

topics = [x for x in range(0,m)]

interest_decay = lambda x,y: x/math.pow(2,y)
factor = math.pow(2,mlen)
function="exponential"
if "settings" in tm:
    if "interest_decay" in tm["settings"]:
        if tm["settings"]["interest_decay"] == "lin":
            interest_decay = lambda x,y: x-y
            factor = mlen+1
            function="linear"


lp = glpk.LPX()                         
lp.name = 'assignment'                  
lp.obj.maximize = True                  
lp.rows.add(len(students)+len(topics))
objective=[]
constraints=[]
for idx,(s,r) in enumerate(zip(students.keys(),lp.rows[:len(students)])):
    r.name = 'student %s' % s
    # Each student gets exactly one topic
    r.bounds = 1.0,1.0
    last = 0
    beg = 0
    while beg<idx*len(topics):
        constraints.append(0.0)
        beg+=1
    comp = [(idx,val) for idx,val in enumerate(students[s])]
    comp = sorted(comp, key=lambda x: x[1])
    for (dv,k) in comp:
        while last < k:
            objective.append(0.0)
            constraints.append(0.0)
            last+=1
        objective.append(interest_decay(factor, dv))
        constraints.append(1.0)
        last+=1
    while last < len(topics):
        objective.append(0.0)
        constraints.append(0.0)
        last+=1
    beg = (idx+1)*len(topics)
    while beg < len(students)*len(topics):
        constraints.append(0.0)
        beg+=1

for e,r in enumerate(lp.rows[len(students):]):                       
    r.name = 'topic %s' % e
    # subject must be chosen at most 1 time.
    r.bounds = 0.0,1.0
    beg = 0
    while beg < len(students)*len(topics):
        if beg%len(topics) == e:
            constraints.append(1.0)
        else:
            constraints.append(0.0)
        beg+=1
lp.cols.add(len(students)*len(topics))                          
for c in lp.cols:                       
    c.name = 'x_%d_%d' % (c.index//len(topics),c.index%len(topics))            
    c.bounds = 0.0, 1.0
lp.obj[:] = objective
iteration = 0
optimal_solutions = []
best_solution = 0
while True:
    lp.matrix = constraints
    retval = lp.simplex()
    assert retval is None
    if lp.status != 'opt':
        break
    for col in lp.cols:
        col.kind = int

    retval = lp.integer()

    assert retval is None

    if lp.status != 'opt':
        break
    if lp.obj.value < best_solution:
        break
    else:
        best_solution = lp.obj.value
    next_constraint = []
    current_solution = []
    for c in lp.cols:
        if c.primal == 1:
            first = c.name.find("_")
            end = c.name.find("_", first+1)
            idx = int(c.name[first+1:end])
            next_constraint.append(1.0)
            current_solution.append((list(students.keys())[idx], c.name[end+1:]))
        else:
            next_constraint.append(0.0)
    lp.rows.add(1)
    r = lp.rows[-1]
    r.name = 'exclude %s' % iteration
    r.bounds = 0.0,len(students)-1
    iteration+=1
    # Add cutting plane
    constraints.extend(next_constraint)
    optimal_solutions.append(current_solution)

if len(optimal_solutions) == 0:
    print("Unsolvable assignment!")
    sys.exit(-1)
print("#Optimal Solutions = {0}; objective weight = {1:.6}; {2} scoring".format(len(optimal_solutions), best_solution, function))
print("Uniform choice among optimal solutions:")
opt = random.choice(optimal_solutions)
for (stud,choice) in opt:
    print("{0:<40} Topic {1}".format(stud, choice))