import pandas as pd
from datetime import datetime
import random

# 对于表格每一列：
# i = 1    -> 时间戳
# i = 6    -> 姓名
# i = 7    -> 部门编号
# i = 8    -> 老人 / 小朋友
# i = 9~28 -> 共 20 个班次，每个班次的意愿度
# i = 29   -> 想要值班的次数
col_timestamp = 1
col_name = 6
col_department = 7
col_member = 8
col_dutystart = 9
col_dutyend = 28
col_dutyfreq = 29

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 的数组，表示每位学生是否是电脑部或电器部成员
    is_hr_array = []             # 长度为 N 的数组，表示每位学生是否是人资部成员

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

    # 遍历过滤后的数据，组建待返回的信息
    for index, (name, info_list) in enumerate(all_data.items()):
        # 学生姓名
        index_to_name_dict[index] = name
        # 意愿度
        preference_mat.append(info_list[col_dutystart:col_dutyend + 1])
        # 想要值班的次数
        want_num_array.append(info_list[col_dutyfreq])
        # 是否是小朋友
        is_new_array.append(index_to_type[info_list[col_member]] == "小朋友")
        # 是否是电脑部或电器部成员
        is_tech_array.append(index_to_departments[info_list[col_department]] in ["电脑部", "电器部"])
        # 是否是人资部成员
        is_hr_array.append(index_to_departments[info_list[col_department]] == "人资部")

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

def save_to_excel(variables, all_data, index_to_name_dict, 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
        hr_new_index = []
        none_tech_new_index = []
        new_index = []
        all_index = []
        _cnt = 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"] = index_to_departments[all_data[index_to_name_dict[i]][col_department]]
                single_stu_info["type"] = index_to_type[all_data[index_to_name_dict[i]][col_member]]
                single_stu_info["duty_monitor"] = False
                on_duty_list.append(single_stu_info)
                single_class_num += 1

                all_index.append(_cnt)
                if single_stu_info["type"] == "小朋友":
                    new_index.append(_cnt)
                    if single_stu_info["department"] in ["人资部", "财外部", "文宣部"]:
                        none_tech_new_index.append(_cnt)
                        if single_stu_info["department"] == 3:
                            hr_new_index.append(_cnt)
                _cnt += 1

        # 为每一班的值班组长打上标记
        duty_monitor_index = -1
        if hr_new_index:
            duty_monitor_index = random.choice(hr_new_index)
        elif none_tech_new_index:
            duty_monitor_index = random.choice(none_tech_new_index)
        elif new_index:
            duty_monitor_index = random.choice(new_index)
        else:
            duty_monitor_index = random.choice(all_index)
        on_duty_list[duty_monitor_index]["duty_monitor"] = True

        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 PatternFill

    # 创建一个新的工作簿
    wb = Workbook()
    ws = wb.active
    if ws is None:
        ws = wb.create_sheet(title="Sheet1")

    # 写入表头
    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']} {stu['department']} {stu['type']}"
            if stu['duty_monitor']:
                ws[col + str(row)].fill = PatternFill(fill_type='solid', start_color='FFFF00', end_color='FFFF00')
            ws[col + str(row)] = str_info

    # 保存文件
    wb.save(file_path)