"""
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"]

WEEKS = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12"]

#Order the rotating schedule is allocated in
ORDER = ["W1", "G1", "R1", "B1", "W2", "G2", "R2", "B2", "W3", "G3", "R3", "B3"]

#List to store generated schedule in
ROSTER=[]

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

#Number of registrars
num_reg = len(REGISTRAR)


#Error check number of registrars is the same as the number of registrars in the order list

#ORDER=REGISTRAR

#Set up days for one week - length of schedule that is rotated
DAY = range(1,8)


# 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)
"""

#Set up demand for shifts
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


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

#Create Variables
REGISTRARS = [(w, s, d) for w in WEEKS \
                        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 w in WEEKS:
    for d in DAY:
        roster_prob += lpSum([regVars[(w,s,d)] for s in SHIFTS]) == 1, "Registrar_%s_do_shift_%d"%(w,d)

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

#Solve
dippy.Solve(roster_prob)

#------------------------------------------------------------------#


#Write rotating schedule to html table (one schedule)
def solution_table(regVars):
    #A file is opened
    rosterout=open('solution_table_rotate.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 w in WEEKS:
        rosterout.write("<tr><td>"+w+"</td>")
        for d in DAY:
            for s in SHIFTS:
                if regVars[(w,s,d)].varValue==1:
                    rosterout.write("<td>")
                    rosterout.write(s)
                    rosterout.write("   </td>")
        rosterout.write("</tr>")

    rosterout.write("</table>")

    rosterout.close()

solution_table(regVars)
#-------------------------------------------------------------------#

#takes registrar and the day (0-84 and works out what shift they are doing)
def what_doing(registrar,day): 
    index_order = ORDER.index(registrar)
    start_week = index_order #week of the rotating schedule they start in
    day_same_as = divmod(day-7*index_order,84)[1]    
    shift = ROSTER[day_same_as]
    return shift
#-------------------------------------------------------------------#




# writes rotating schedule to html table
# creates the list ROSTER which containts the schedules in order as elements of this list
# writes full schedule to solution_table_rotate_full.html
# May need to modify so that prints in order of teams rather than order that schedules are allocated
def full_schedule(regVars):

    #creates first schedule
    for w in WEEKS:
        for d in DAY:        
            for s in SHIFTS:
                if regVars[(w,s,d)].varValue==1:
                    ROSTER.append(s)
                    
    
    #print ROSTER
   
    #A file is opened
    rosterout=open('solution_table_rotate_full.html','w')

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

    #Write day headers - could modify to display dates
    for d in range(0,85):        
        rosterout.write("<td>" + str(d) + "</td>")
    rosterout.write("</tr>")

    #Loop to write roster to HTML table   
    i=0
    for r in REGISTRAR:
        rosterout.write("<tr><td>"+ r +"</td>")
        for d in range(0,84):
            #If first registrar then write roster otherwise write rotated roster
            if r == REGISTRAR[0]:
                current_shift = ROSTER[d]
            else:
                current_shift= what_doing(r,d)#call function here
            #Write color code
            if current_shift=='A':
                rosterout.write("<td bgcolor=#0066FF>")
            elif current_shift=='B':
                rosterout.write("<td bgcolor=#009933>")
            elif current_shift=='N':
                rosterout.write("<td bgcolor=#FFFF00>")
            elif current_shift=='P' or current_shift=='X' or current_shift=='Z':
                rosterout.write("<td bgcolor=#999999>")
            else:
                rosterout.write("<td>") 
            #Write shift and end cell tag
            rosterout.write(current_shift)
            rosterout.write("     </td>")
        rosterout.write("</tr>")
        i+=1

    rosterout.write("</table>")

    rosterout.close()
    
    return ROSTER
#-------------------------------------------------------------------#

#Call function to write full schedule to HTML
full_schedule(regVars)