import pandas as pd
from datetime import datetime

index_to_departments = {     # 部门编号和部门名称的对应关系
    1: "电脑部",
    2: "电器部",
    3: "人资部",
    4: "财外部",
    5: "文宣部"
}

index_to_type = {            # 老人/小朋友和对应的编号
    1: "老人",
    2: "小朋友"
}   

def read_excel(file_path):
    # 读取 Excel 文件
    df = pd.read_excel(file_path)

    # 待返回的所有信息。下令 N 为学生人数，M 为班次数
    index_to_name_dict = {}      # 长度为 N 的字典，包含数组下标和学生姓名的对应关系
    preference_mat = []          # N x M 的二维数组，表示每位学生对每个班次的满意度
    want_num_array = []          # 长度为 N 的数组，表示每位学生想要值班的次数
    is_new_array = []            # 长度为 N 的数组，表示每位学生是否是小朋友
    is_tech_array = []           # 长度为 N 的数组，表示每位学生是否是电脑部或电器部成员 

    '''
        遍历每一行，对于每一行：
            i=1    -> 时间戳  
            i=6    -> 姓名
            i=7    -> 部门编号
            i=8    -> 老人/小朋友
            i=9~28 -> 共20个班次，每个班次的意愿度
            i=29   -> 想要值班的次数

        考虑到原始填表信息中可能有某位同学多次提交的记录，先过滤一下冗余信息
    '''
    all_data = {}  # 储存过滤后的数据
    for index, row in df.iterrows():
        # 读取学生姓名，比较时间戳
        name = row[6]
        format = "%Y/%m/%d %H:%M:%S"
        time_this = datetime.strptime(row[1], format)
        if (name not in all_data) or (time_this > datetime.strptime(all_data[name][1], format)):
            info_list = df.iloc[index].tolist()
            all_data[name] = info_list

    # 遍历过滤后的数据，组建待返回的信息
    for index, (name, info_list) in enumerate(all_data.items()):
        # 学生姓名
        index_to_name_dict[index] = name

        # 意愿度
        preference = info_list[9:29]
        preference_mat.append(preference)

        # 想要值班的次数
        want_num = info_list[29]
        want_num_array.append(want_num)

        # 是否是小朋友
        is_new_array.append(index_to_type[info_list[8]] == "小朋友")

        # 是否是电脑部或电器部成员
        is_tech_array.append(index_to_departments[info_list[7]] in ["电脑部", "电器部"])

    return all_data, index_to_name_dict, preference_mat, want_num_array, is_new_array, is_tech_array

def save_to_excel(variables, all_data, index_to_name_dict, preference_mat, file_path):

    # 用一个list储存每一班的值班人员，该list中每个元素是一个存有若干dict的list，每个dict表示某一班的某一值班人员信息
    all_result = []
    max_single_class_num = 0 # 用于记录最大的班次人数，以便后写入 excel
    for j in range(len(variables[0])):
        on_duty_list = []
        single_class_num = 0
        for i in range(len(variables)):
            if variables[i][j] == 1:
                single_stu_info = {}
                single_stu_info["name"] = index_to_name_dict[i]
                single_stu_info["department"] = all_data[index_to_name_dict[i]][7]
                single_stu_info["type"] = all_data[index_to_name_dict[i]][8]
                single_stu_info["adjust"] = preference_mat[i][j] <= 0
                on_duty_list.append(single_stu_info)
                single_class_num += 1
        on_duty_list = sorted(on_duty_list, key=lambda x: x["department"]) # 按部门编号升序排序
        all_result.append(on_duty_list)
        max_single_class_num = max(max_single_class_num, single_class_num)
    

    from openpyxl import Workbook
    from openpyxl.styles import Font, Alignment, Border, Side, PatternFill

    # 创建一个新的工作簿
    wb = Workbook()
    ws = wb.active

    # 写入表头
    width = 25
    ws['B1'] = "星期一"
    ws.column_dimensions['B'].width = width
    ws['C1'] = "星期二"
    ws.column_dimensions['C'].width = width
    ws['D1'] = "星期三"
    ws.column_dimensions['D'].width = width
    ws['E1'] = "星期四"
    ws.column_dimensions['E'].width = width
    ws['F1'] = "星期五"
    ws.column_dimensions['F'].width = width
    ws['G1'] = "星期六"
    ws.column_dimensions['G'].width = width
    ws['H1'] = "星期日"
    ws.column_dimensions['H'].width = width

    start_index = []
    ws.merge_cells(f'A2:A{2+max_single_class_num}')
    ws['A2'] = "第一班"
    start_index.append(2)
    ws.merge_cells(f'A{2+max_single_class_num+1}:A{2+2*max_single_class_num+1}')
    ws[f'A{2+max_single_class_num+1}'] = "第二班"
    start_index.append(2+max_single_class_num+1)
    ws.merge_cells(f'A{2+2*max_single_class_num+2}:A{2+3*max_single_class_num+2}')
    ws[f'A{2+2*max_single_class_num+2}'] = "第三班"
    start_index.append(2+2*max_single_class_num+2)

    for duty_index, duty in enumerate(all_result):
        for stu_index, stu in enumerate(duty):
            col = chr(ord('B')+duty_index//3)
            row = start_index[duty_index%3]+stu_index
            str_info = f"{stu['name']} {index_to_departments[stu['department']]} {index_to_type[stu['type']]}"
            if stu['adjust']:
                ws[col+str(row)].fill = PatternFill(fill_type='solid', start_color='FF0000', end_color='FF0000')
            ws[col+str(row)] = str_info
        

    # 保存文件
    wb.save(file_path)