"""
01/12/2010 Initial formulation for Roster Redesign exploration
"""

from coinor.pulp import *
import coinor.dippy as dippy
import math
from numpy import array, concatenate

SHIFTS = ["A", "B", "C", "D","N", "Z", "P", "T", "X"]

REGISTRAR = ["W1", "W2", "W3", "G1", "G2", "G3", "R1", "R2", "R3", "B1", "B2", "B3"]

DAY = range(1,85)


# Function to determine if it is a weekend or not
# Returns 1 if day is mon-thurs
# returns 5 if day is friday
# returns 6 if day is saturday
# returns 7 if day is a sunday
def weekend(day):
    if (day+1)%7 == 0:
        weekend=6
    elif day%7==0:
        weekend=7
    elif (day+6)%7 ==0 or (day+5)%7==0 or (day+4)%7==0  or (day+3)%7==0:
        weekend =1 #day is monday to thurs
    else:
        weekend = 5 #day is a friday
    return weekend

"""#test code for weekend function
for d in DAYS:
    print weekend(d)
"""


DEMAND={}
#demand       
for d in DAY:
    DEMAND[d] = {}
    day=weekend(d)
    DEMAND[d]['A']=1
    if day == 1: #mon - thurs
        DEMAND[d]['B']=1
        DEMAND[d]['C']=1
        DEMAND[d]['D']=1
        DEMAND[d]['N']=1
        DEMAND[d]['Z']=0
        DEMAND[d]['P']=0
        DEMAND[d]['X']=0
        DEMAND[d]['T']=7
    elif day ==5: #friday
        DEMAND[d]['B']=1
        DEMAND[d]['C']=1
        DEMAND[d]['D']=1
        DEMAND[d]['N']=0
        DEMAND[d]['Z']=1
        DEMAND[d]['P']=0
        DEMAND[d]['X']=0
        DEMAND[d]['T']=7
    elif day ==6: #saturday
        DEMAND[d]['B']=0
        DEMAND[d]['C']=0
        DEMAND[d]['D']=0
        DEMAND[d]['N']=0
        DEMAND[d]['Z']=1
        DEMAND[d]['P']=2
        DEMAND[d]['X']=8
        DEMAND[d]['T']=0
    else: #sunday
        DEMAND[d]['B']=0
        DEMAND[d]['C']=0
        DEMAND[d]['D']=0
        DEMAND[d]['N']=1
        DEMAND[d]['Z']=1
        DEMAND[d]['P']=2
        DEMAND[d]['X']=7
        DEMAND[d]['T']=0



""" # can print out entire dictionary
for s in SHIFTS:
    for record in DEMAND.values():
        print record, record[s]
"""

#Linear problem
roster_prob = dippy.DipProblem("Roster", LpMinimize)

#Create Variables
REGISTRARS = [(r, s, d) for r in REGISTRAR \
                        for s in SHIFTS \
                        for d in DAY]
regVars = LpVariable.dicts("registrar", REGISTRARS, cat = LpBinary)


#objective function
roster_prob += 1, "arbitary objective function"

#constraints
#Allocate each registrar to one shift
for r in REGISTRAR:
    for d in DAY:
        roster_prob += lpSum([regVars[(r,s,d)] for s in SHIFTS]) == 1, "Registrar_%s_do_shift_%d"%(r,d)

#Meet demand
for s in SHIFTS:
    for d in DAY:
        roster_prob += lpSum([regVars[(r,s,d)] for r in REGISTRAR]) == DEMAND[d][s], "meet_shift_demand%s_on_day%d"%(s,d)
        

#Solve
dippy.Solve(roster_prob)

"""
#Display Roster
for d in DAY:
    for r in REGISTRAR:
        for s in SHIFTS:
            if regVars[(r,s,d)].varValue==1:
                print r,s,d
"""


def output_solution (regVars):
    #A file is opened
    rosterout=open('roster_solution.txt','w')
    for d in DAY:
        for r in REGISTRAR:
            for s in SHIFTS:
                if regVars[(r,s,d)].varValue==1:
                    rosterout.write(r + " ")
                    rosterout.write(s + " ")
                    rosterout.write(str(d))

            rosterout.write("\n")
        #rosterout.write("\n")

    rosterout.close()


output_solution (regVars)

def solution_table(regVars):
    #A file is opened
    rosterout=open('solution_table.html','w')

    

    #table
    rosterout.write("<table>")
    rosterout.write("<tr><td>Team</td>")
    rosterout.write("<tr><td>Day</td>")

    for d in DAY:        
        rosterout.write("<td>" + str(d) + "</td>")
    rosterout.write("</tr>")
        
    
    
    for r in REGISTRAR:
        rosterout.write("<tr><td>"+r+"</td>")
        for d in DAY:
            for s in SHIFTS:
                if regVars[(r,s,d)].varValue==1:
                    if s =='A':
                        rosterout.write("<td bgcolor=#0066FF>")
                    elif s =='B':
                        rosterout.write("<td bgcolor=#009933>")
                    elif s =='N':
                        rosterout.write("<td bgcolor=#FFFF00>")
                    elif s =='P' or s =='X' or s =='Z':
                        rosterout.write("<td bgcolor=#999999>")
                    else:
                        rosterout.write("<td>") 
                    rosterout.write(s)
                    rosterout.write("   </td>")
        rosterout.write("</tr>")

    rosterout.write("</table>")

    rosterout.close()

solution_table(regVars)