"""
November-Feburary 2010-2011
Amelia White Summer Project
Supervisors:
Mike O'Sullivan and Cameron Walker

Functions for roster_rotate_1_4.py

Initial exploratoray analysis into formulating a roster for General
Medicine registrars at Auckland Hostipal
"""

#Import the required libraries
import scipy
import matplotlib.pyplot as plt
import numpy as np
import math

#Can't compare variables directly. Although variables are decleared as binary
#They seem to be stored as floating point numbers
TOL = 1e-5



"""
# Function to determine if it is a weekend or not
# Inputs:
#   day = day from 0-84 in 12 week roster
# Outputs:
#   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


"""
# Function to return a numerical reference of the day of week
# Inputs:
#   day = day from 0-84 in 12 week roster
# Outputs:
#   day_num = day of week(mon = 1 ... sun = 7)
"""
def what_day(day):
    if (day+1)%7 == 0:
        day_num=6   # Saturday
    elif day%7==0: 
        day_num=7   # Sunday
    elif (day+6)%7 ==0:
        day_num=1   # monday
    elif (day+5)%7==0:
        day_num = 2 # Tuesday
    elif (day+4)%7==0:
        day_num= 3  # Wednesday
    elif (day+3)%7==0:
        day_num =4  # Thursday
    else:
        day_num = 5 # Friday
    return day_num


"""
# Function to generate number of patients admitted to general medicine
# Inputs:
#   day = day from 1-7
# Outputs:
#   patient_num = number of patients admitted to GM wards
"""
def admitting_nums(day):
    if day == 1:
        patient_num = 28  # monday
    elif day == 2: 
        patient_num = 28 # Tuesday  
    elif day == 3:
        patient_num = 28 # wednesday  
    elif day == 4:
        patient_num = 32 # Thursday
    elif day == 5:
        patient_num= 32  # Friday
    elif day == 6:
        patient_num =28*4  # Saturday = 28*4 to model shift A taking all the patients
    else:
        patient_num = 20*4 # Sunday = 20*4 to model shift A taking all the patients
    return patient_num


"""
# Solution Table
# Writes rotating schedule to html table (one schedule)
# Inputs:
#   regVars: variables X(r,s,d)=1 if registrar r works shift s on day d, 0 otherwise
#   DAY:     List of days '1' to '8' 
#   WEEKS:   List of weeks '1' to length of roster
#   SHIFTS:  List of possible shifts
# Outputs:
#   The roster written to the file: "solution_table_rotate.html" this is the schedule
#   that each registrar does.
"""
def solution_table(regVars,DAY,WEEKS,SHIFTS):
    
    #A file is opened
    rosterout=open('solution_table_rotate.html','w')

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

    #write day header
    for d in DAY:        
        rosterout.write("<td>" + str(d) + "</td>")
    rosterout.write("</tr>")

    #write schedule       
    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-TOL: 
                    rosterout.write("<td>")
                    rosterout.write(s)
                    rosterout.write("   </td>")
        rosterout.write("</tr>")

    rosterout.write("</table>")

    rosterout.close()


"""
# What_doing
# takes registrar and the day (0-84) and works out what shift they are doing
# Based on the roster.
# Inputs:
#   registrar: registrar to find out what shift they are doing
#   DAY:       List of days '1' to '8'
#   ORDER:     List of the registrars in order of the rotating roster is assigned in
#   ROSTER:    12 week roster for one registrar
# Outputs:
#   shift:      the shift the registrar is doing of day d
"""

def what_doing(registrar,day,ORDER,ROSTER): 
    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



"""
# lp_what_doing
# takes registrar and the day (0-84) and works out what shift they are doing
# Based on the roster.
# Inputs:
#   registrar: registrar to find out what shift they are doing
#   DAY:       List of days '1' to '8'
#   ORDER:     List of the registrars in order of the rotating roster is assigned in
#   ROSTER:    12 week roster for one registrar
# Outputs:
#   shift:      the shift the registrar is doing of day d
"""

def lp_what_doing(registrar,day,ORDER): 
    index_order = ORDER.index(registrar)
    start_week = index_order #week of the rotating schedule they start in
    day_same_as = (day+7*index_order) % 84
    if day_same_as == 0:
        day_same_as = 84
    return day_same_as




"""
# full_schedule
# Writes rotating schedule to html table
# Inputs:
#   regVars: variables X(r,s,d)=1 if registrar r works shift s on day d, 0 otherwise
#   DAY:     List of days '1' to '8' 
#   WEEKS:   List of weeks '1' to length of roster
#   SHIFTS:  List of possible shifts
#   REGISTRAR: List of registrars
#   ORDER:     List of the registrars in order of the rotating roster is assigned in
# Outputs:
#   The full roster written to the file: "solution_table_rotate_full.html"
#   ROSTER:      List of the order shifts are done in
#   FULL_ROSTER: Dictionary conatining the shifts for each registrar on each day
"""

def full_schedule(regVars,DAY,WEEKS,SHIFTS,REGISTRAR,ORDER):
    ROSTER=[]
    # creates list of the order that shifts are done in for the first registrar
    for w in WEEKS:
        for d in DAY:        
            for s in SHIFTS:
                if math.fabs(regVars[(w,s,d)].varValue-1)<=TOL:
                    ROSTER.append(s)
                    print w, d, s
                    
    
    # A file is opened
    rosterout=open('solution_table_rotate_full.html','w')

    # HTML table tags
    rosterout.write("<table cellpadding='4' cellspacing='0'>")
    rosterout.write("<tr><td>Team</td>")
    rosterout.write("<tr><td>Day</td>")

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

    #dictionary to store full roster
    FULL_ROSTER={}
    
    #Loop to write roster to HTML table and build full roster dictionary
    i=1
    for r in REGISTRAR:
        #if night shift write border
        week_start=ORDER.index(r)+1
        night_start=week_start*7-7
        #print night_start
        night_finish=night_start+7
        #print night_finish
        rosterout.write("<tr><td>"+ r +"</td>")
        FULL_ROSTER[r]={}

        #if r == REGISTRAR[0]:
        #    #rosterout.write(" <td colspan=8><table border='4' cellpadding='4' bordercolor=yellow cellspacing='0'><tr>")
        
        for d in range(1,85):
            #If first registrar then write roster otherwise write rotated roster
            if r == REGISTRAR[0]:
                #print d
                current_shift = ROSTER[d-1]
                
                
            else:
                current_shift= what_doing(r,d-1,ORDER,ROSTER)#call function here
            #Write color code
            #if d == night_start:
            #    rosterout.write(" <td colspan=8><table border='4' cellpadding='4' bordercolor=yellow cellspacing='0'><tr>")
                
            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 to html and dictionary and end cell tag
            rosterout.write(current_shift)
            FULL_ROSTER[r][d-1]=current_shift
            rosterout.write("     </td>")

            #if d == night_finish:
            #    rosterout.write(" </td></tr></table></td>")
                
        rosterout.write("</tr>")
        i+=1

    rosterout.write("</table>")
    rosterout.close()
    
    return (ROSTER, FULL_ROSTER)

"""
# patient_nums
# function to determine number of patients in each ward based on the schedule
# Inputs:
#   WARD_NUMBERS:   dictionary with initial patient numbers 
#   REGISTRAR:      List of registrars
#   FULL_ROSTER:    Dictionary of which shift each registrar is doing each day
#   PATIENT_DATA:   list of number of patients admitted to GM each day  
#   WEIGHTS:        Fraction of patients that get alloctaed to each shift
#   DISCHARGE_RATE: Fraction of patients that are discharged each day
#   SHIFTS:         List of shifts
Outputs:
#   WARD_NUMBERS:   Full dictionary over registrar, Day for patient numbers
"""
def patient_nums(WARD_NUMBERS, patientVars):

    
    for d in range(1,85):
        WARD_NUMBERS[('W', d)]=patientVars[('W',d)].varValue
        WARD_NUMBERS[('B', d)]=patientVars[('B',d)].varValue
        WARD_NUMBERS[('R', d)]=patientVars[('R',d)].varValue
        WARD_NUMBERS[('G', d)]=patientVars[('G',d)].varValue
        
    return WARD_NUMBERS
    


"""
# plot_patients
# Function to plot ward numbers
# uses matplotlib see https://matplotlib.sourceforge.net/
# Inputs:
#   WARD_NUMBERS:   Full dictionary over registrar, Day for patient numbers
#   REGISTRAR:      List of registrars
# Outputs:
#   Plots patient numbers for each ward
"""
def plot_patients(WARDS, WARD_NUMBERS):

    totals_W=[]
    totals_B=[]
    totals_R=[]
    totals_G=[]
    for d in range(0,85):
        totals_W.append(WARD_NUMBERS[('W',d)])
        totals_R.append(WARD_NUMBERS[('R',d)])
        totals_G.append(WARD_NUMBERS[('G',d)])
        totals_B.append(WARD_NUMBERS[('B',d)])


    #print totals_W
    # Make plot
    plt.subplot(111, axisbg='lightslategray')

    #for w in
    #Plot colours for each ward
    plt.plot(range(0,85),totals_W,'w',label='White')
    plt.plot(range(0,85),totals_B,'k',label='Black')
    plt.plot(range(0,85),totals_G,'y',label='Gold')
    plt.plot(range(0,85),totals_R,'r',label='Red')




    # Plot commands to make pretty
    plt.xlabel('Day')
    plt.ylabel('Number of Patients')
    plt.title('Estimate of Patient numbers across wards')
    plt.legend(('White','Black','Gold','Red'),loc='best')
    plt.xlim([0,21])
    plt.ylim([0,55])
    plt.grid(True)

    # Make legend have a grey background
    leg = plt.gca().get_legend()
    frame  = leg.get_frame()  
    frame.set_facecolor('0.80')  
   
    # Show plot
    plt.show()