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

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

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

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

WARDS = ["W", "G", "R", "B"]
registrars = {}
registrars["W"] = ["W1", "W2", "W3"]
registrars["G"] = ["G1", "G2", "G3"]
registrars["R"] = ["R1", "R2", "R3"]
registrars["B"] = ["B1", "B2", "B3"]

#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 = [r for w in WARDS for r in registrars[w]]

#Number of registrars
num_reg = len(REGISTRAR)

#Maximum number of patients allowable per registrar
max_patients=33


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


#INITIAL NUMBER OF PATIENTS IN WARDS AT START OF ROSTER
WARD_NUMBERS={}
WARD_NUMBERS[('W', 0)]=18
WARD_NUMBERS[('B', 0)]=22
WARD_NUMBERS[('R', 0)]=23
WARD_NUMBERS[('G', 0)]=19

#Fraction of patients that get allocated to each shift
#need to change to consider weekends
WEIGHTS={}
WEIGHTS['A']=0.25
WEIGHTS['B']=0.25
WEIGHTS['C']=0.25
WEIGHTS['D']=0.25
WEIGHTS['N']=0
WEIGHTS['Z']=0
WEIGHTS['P']=0
WEIGHTS['T']=0
WEIGHTS['X']=0

#Data of average admissions:
#need to change this to include weekends etc
PATIENT_DATA={}
DISCHARGE_RATE={}
for d in range(0,85):
    PATIENT_DATA[d]=admitting_nums(what_day(d))
    DISCHARGE_RATE[d]=0.33


PATIENT_DATA[0]=0
#print PATIENT_DATA



#Set up demand for shifts
# THIS NEEDS A BIT OF TYDING UP TO MAKE INTO NICE FUNCTION
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']=0
        DEMAND[d]['Z']=0
        DEMAND[d]['P']=0
        DEMAND[d]['X']=0
        DEMAND[d]['T']=8
    elif day ==5: #friday
        DEMAND[d]['B']=1
        DEMAND[d]['C']=1
        DEMAND[d]['D']=1
        DEMAND[d]['N']=0
        DEMAND[d]['Z']=0
        DEMAND[d]['P']=0
        DEMAND[d]['X']=0
        DEMAND[d]['T']=8
    elif day ==6: #saturday
        DEMAND[d]['B']=0
        DEMAND[d]['C']=0
        DEMAND[d]['D']=0
        DEMAND[d]['N']=0
        DEMAND[d]['Z']=0
        DEMAND[d]['P']=2
        DEMAND[d]['X']=9
        DEMAND[d]['T']=0
    else: #sunday
        DEMAND[d]['B']=0
        DEMAND[d]['C']=0
        DEMAND[d]['D']=0
        DEMAND[d]['N']=0
        DEMAND[d]['Z']=0
        DEMAND[d]['P']=2
        DEMAND[d]['X']=9
        DEMAND[d]['T']=0


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

#Create Variables
#Shift variables
REGISTRARS = [(w, s, d) for w in WEEKS \
                        for s in SHIFTS \
                        for d in DAY]
regVars = LpVariable.dicts("registrar", REGISTRARS, cat = LpBinary)
#patient  number variables
# day zero initial number of patients
PATIENT_NUMS = [(w, d) for w in WARDS for d in range(0,85)]

patientVars = LpVariable.dicts("patient_numbers", PATIENT_NUMS, lowBound=0, cat='Continuous')


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


#roster_prob += max_patients, "minimize_max_patients"

#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)
        
#Friday A does P on saturday and sunday
for w in WEEKS:
    roster_prob += regVars[(w,'A',5)] <= regVars[(w,'P',6)], "do_A_on_fri_do_P_sat_%s"%w
for w in WEEKS:
    roster_prob += regVars[(w,'A',5)] <= regVars[(w,'P',7)], "do_A_on_fri_do_P_sun_%s"%w


#Friday B does P on saturday and A on sunday
for w in WEEKS:
    roster_prob += regVars[(w,'B',5)] <= regVars[(w,'P',6)], "do_B_on_fri_do_P_sat_%s"%w
for w in WEEKS:
    roster_prob += regVars[(w,'B',5)] <= regVars[(w,'A',7)], "do_B_on_fri_do_A_sat_%s"%w

#A on saturday does P on Sunday
for w in WEEKS:
    roster_prob += regVars[(w,'A',6)] <= regVars[(w,'P',7)], "do_A_on_Sat_do_P_sun_%s"%w

#Patient Numbers constraints
for w in WARDS:
    roster_prob += patientVars[(w, 0)] == WARD_NUMBERS[(w, 0)],"patient_numbers_init_ward_%s"%w
    
for w in WARDS:
    for d in range(1, 85):
        conExpr = patientVars[(w, d-1)] * (1-DISCHARGE_RATE[d])
        for r in registrars[w]:
            myDay = lp_what_doing(r, d, ORDER)
            (week, day) = divmod(myDay, 7)
            if day == 0:
                day = 7
            else:
                week += 1
##            print d, r, ORDER.index(r), myDay, week, day
            conExpr += lpSum([regVars[(str(week), s, day)]*WEIGHTS[s]*PATIENT_DATA[d] for s in SHIFTS])

##        print conExpr        
        roster_prob += patientVars[(w, d)] == conExpr, "patient_numbers_ward_%s_day_%d"%(w, d)


for w in WARDS:
    for d in range(0,84):
        roster_prob += patientVars[(w, d)] <= max_patients, "patient_numbers_less_than_max_ward_%s_day_%d"%(w, d)


#TEST TEST TEST!!!!
#roster_prob += regVars[('1','A',3)] >= 1, "TEST"



# First week constraint to allow night shift
# can on do one long shift (A or B) in a seven day period
# week one of the roster has been arbitarly selected to be the night shift week
# where a relief registrar covers
# change d to n here
for d in DAY:
    roster_prob += (lpSum([regVars[('2','A',n)] for n in range(1,d)])\
                    +lpSum([regVars[('2','B',n)] for n in range(1,d)])\
                    +lpSum([regVars[('1','A',n)] for n in range(d,8)])\
                    +lpSum([regVars[('1','B',n)] for n in range(d,8)]))<=1,"one-long_shift_in_any_7_days_%d"%d


#No more than 2 long shifts in a 7 day period
for w in range(1,13):
    if w==12:
        nextweek=1
    else:
        nextweek=w+1
    for d in DAY:
        roster_prob += (lpSum([regVars[(str(nextweek),'A',n)] for n in range(1,d)])\
                        +lpSum([regVars[(str(nextweek),'B',n)] for n in range(1,d)])\
                        +lpSum([regVars[(str(w),'A',n)] for n in range(d,8)])\
                        +lpSum([regVars[(str(w),'B',n)] for n in range(d,8)]))<=2,"only_2_long_week_%d_day_%d"%(w,d)



#relief roster doesn't work weekends
roster_prob += regVars[('1','X',6)]>=1,"rel_dont_work_sat"
roster_prob += regVars[('1','X',7)]>=1,"rel_dont_work_sun"


#Solve
#dippy.Solve(roster_prob)
roster_prob.solve()

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

#WRITE SINGLE ROSTER
solution_table(regVars,DAY,WEEKS,SHIFTS)

#Call function to write full schedule to HTML
(ROSTER, FULL_ROSTER) = full_schedule(regVars,DAY,WEEKS,SHIFTS,REGISTRAR,ORDER)

"""
for r in REGISTRAR:
    for d in DAY:
        print FULL_ROSTER[r][d]
"""
    
patient_nums(WARD_NUMBERS, patientVars)
"""
for w in WARDS:
    for d in range(0,84):
        print w, d, patientVars[(w,d)].varValue
    
print  WARD_NUMBERS      
""" 



print '----------------------------------'
#print WARD_NUMBERS[('W',1)]
plot_patients(WARDS, WARD_NUMBERS)