import pandas as pd
from datetime import datetime
from random import choice

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

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

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

def read_excel(file_path):
    df = pd.read_excel(file_path)

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

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

    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])
        depart_mat.append([1 if info_list[col_department] == i else 0 for i in range(1, 6)])
        want_num_array.append(info_list[col_dutyfreq])
        is_new_array.append(index_to_type[info_list[col_member]] == "小朋友")

    return all_data, index_to_name_dict, preference_mat, depart_mat, want_num_array, is_new_array

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

    # 收集每个班次的值班人员信息
    all_result = []
    max_single_class_num = 0
    
    for j in range(len(variables[0])):
        on_duty_list = []
        hr_new_index = []
        none_tech_new_index = []
        new_index = []
        all_index = []
        
        for i in range(len(variables)):
            if variables[i][j] == 1:
                name = index_to_name_dict[i]
                department = index_to_departments[all_data[name][col_department]]
                member_type = index_to_type[all_data[name][col_member]]
                
                single_stu_info = {
                    "name": name,
                    "department": department,
                    "type": member_type,
                    "duty_monitor": False
                }
                on_duty_list.append(single_stu_info)
                
                idx = len(all_index)
                all_index.append(idx)
                if member_type == "小朋友":
                    new_index.append(idx)
                    if department in ["人资部", "财外部", "文宣部"]:
                        none_tech_new_index.append(idx)
                        if department == "人资部":
                            hr_new_index.append(idx)

        # 选择值班组长（优先级：人资部小朋友 > 非技术部小朋友 > 小朋友 > 所有人）
        if hr_new_index:
            duty_monitor_index = choice(hr_new_index)
        elif none_tech_new_index:
            duty_monitor_index = choice(none_tech_new_index)
        elif new_index:
            duty_monitor_index = choice(new_index)
        else:
            duty_monitor_index = choice(all_index)
        on_duty_list[duty_monitor_index]["duty_monitor"] = True

        on_duty_list.sort(key=lambda x: x["department"])
        all_result.append(on_duty_list)
        max_single_class_num = max(max_single_class_num, len(on_duty_list))
    

    from openpyxl import Workbook
    from openpyxl.styles import PatternFill

    wb = Workbook()
    ws = wb.active if wb.active else wb.create_sheet(title="Sheet1")

    # 设置表头
    weekdays = ["星期一", "星期二", "星期三", "星期四", "星期五", "星期六", "星期日"]
    for idx, day in enumerate(weekdays, start=2):
        col = chr(ord('A') + idx - 1)
        ws[f'{col}1'] = day
        ws.column_dimensions[col].width = 25

    # 设置班次标题和起始行
    shift_names = ["第一班", "第二班", "第三班"]
    start_index = []
    for shift_idx, shift_name in enumerate(shift_names):
        start_row = 2 + shift_idx * (max_single_class_num + 1)
        end_row = start_row + max_single_class_num
        ws.merge_cells(f'A{start_row}:A{end_row}')
        ws[f'A{start_row}'] = shift_name
        start_index.append(start_row)

    # 填写值班人员信息
    yellow_fill = PatternFill(fill_type='solid', start_color='FFFF00', end_color='FFFF00')
    for duty_index, duty in enumerate(all_result):
        col = chr(ord('B') + duty_index // 3)
        for stu_index, stu in enumerate(duty):
            row = start_index[duty_index % 3] + stu_index
            cell = ws[f'{col}{row}']
            cell.value = f"{stu['name']} {stu['department']} {stu['type']}"
            if stu['duty_monitor']:
                cell.fill = yellow_fill

    wb.save(file_path)