import streamlit as st
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

st.title('护士排班 NRP problem')

st.subheader('请输入您的护士编号，即可获得您本周的排班表')

res = st.number_input('',0,100)   #用于显示数字输入小部件

if res:
    res = int(res)
    df = pd.read_excel('result.xlsx')

    key1 = df['a'].to_list()
    key = [eval(k) for k in key1]
    value = df['b'].to_list()

    dic = dict(zip(key,value))

    J = list(range(7))  # 星期几 1 2 3 4 5 6 7
    I = list(range(6))  # 班次 {1:6-10,2:10-14,3:14-18,4:18-22,5:22-2,6:2-6}
    nurse = [[0 for col in J] for row in I]
    for k,v in dic.items():
        if k[2] == res:
            nurse[k[0]][k[1]] += v    

    week = ['Mon','Tus','Wed','Ths','Fri','Sat','Sun']

    table = pd.DataFrame(nurse, columns= week)   #将字典转换成为数据框
    table.index = ['2:00-6:00','6:00-10:00','10:00-14:00','14:00-18:00','18:00-22:00','22:00-2:00']

    print(f'nurse{res}\'s table')

    plt.rcParams.update({"font.size":15})#此处必须添加此句代码方可改变标题字体大小

    plt.figure(figsize=(12,8))
    ax = sns.heatmap(table,linewidths=0.05,linecolor='black', cmap="YlGnBu")
    ax.xaxis.tick_top()
    ax.tick_params(axis='both', which='both', length=0)  # 删除横纵坐标轴上的刻度线

    plt.yticks( rotation=0)
    plt.savefig(f'timetable_{res}.png',dpi=600)

    st.caption('请查收您的排班表')    #用于写字幕
    st.image(f'timetable_{res}.png')      #用于显示图像

    with open(f"timetable_{res}.png", "rb") as file:
        btn = st.download_button(
             label="下载排班表",
             data=file,
             file_name="timetable.png",
           )