messyscript/换座位/generate_seating.py

#!/usr/bin/env python3
"""
Generate seating charts from Excel class data.

Required input columns: 班级, 姓名, 性别, 总分
"""

import argparse
import random
import sys
from pathlib import Path

try:
    from openpyxl import Workbook
    from openpyxl.styles import Alignment, Border, Font, PatternFill, Side
    from openpyxl.utils import get_column_letter
except ImportError:
    print("Error: openpyxl is required. Install with: pip install openpyxl")
    sys.exit(1)


def parse_args():
    parser = argparse.ArgumentParser(
        description="根据Excel学生数据生成座位表"
    )
    parser.add_argument(
        "excel_file", type=str,
        help="输入Excel文件路径 (.xlsx)"
    )
    parser.add_argument(
        "-c", "--class", dest="class_name", type=str, required=True,
        help="要生成座位表的班级名称"
    )
    parser.add_argument(
        "-s", "--style", type=str, required=True,
        help="座位表样式，如 '2222'（4组2人桌）、'3332'（3组3人桌+1组2人桌）"
    )
    parser.add_argument(
        "-m", "--allow-mixed", action="store_true",
        help="允许男女生同桌（默认不允许）"
    )
    parser.add_argument(
        "-t", "--strategy", type=str, default="good-with-good",
        choices=["good-with-good", "good-with-bad", "random"],
        help="成绩配对策略：good-with-good（好+好）、good-with-bad（好+差）、random（无关）"
    )
    parser.add_argument(
        "-o", "--output", type=str, default=None,
        help="输出Excel文件路径（默认：<输入文件>_<班级>_座位表.xlsx）"
    )
    parser.add_argument(
        "-n", "--count", type=int, default=1,
        help="生成座位表的个数（默认1个）"
    )
    parser.add_argument(
        "--seed", type=int, default=None,
        help="随机种子（用于 random 策略的可重复性）"
    )
    parser.add_argument(
        "--usage", action="help",
        help="显示帮助信息"
    )
    return parser.parse_args()


def read_students(filepath, class_name):
    """从Excel读取学生数据，按班级筛选。自动查找表头行。"""
    from openpyxl import load_workbook

    wb = load_workbook(filepath, data_only=True)
    ws = wb.active

    required = ["班级", "姓名", "性别", "总分"]
    indices = {}
    header_row = None

    # 自动查找包含所有必需列的表头行（最多扫描前10行）
    for r in range(1, min(ws.max_row + 1, 11)):
        row_vals = []
        for col in range(1, ws.max_column + 1):
            val = ws.cell(row=r, column=col).value
            row_vals.append(str(val).strip() if val else "")
        if all(col_name in row_vals for col_name in required):
            for col_name in required:
                indices[col_name] = row_vals.index(col_name)
            header_row = r
            break

    if header_row is None:
        # 输出前几行帮助诊断
        print("错误：Excel文件中未找到 '班级'、'姓名'、'性别'、'总分' 列。")
        for r in range(1, min(ws.max_row + 1, 6)):
            preview = [str(ws.cell(row=r, column=c).value or "") for c in range(1, ws.max_column + 1)]
            print(f"  第{r}行：{preview}")
        sys.exit(1)

    # 规范化班级名：去除"班"后缀以便灵活匹配
    class_name_clean = class_name.rstrip("班")
    data_start = header_row + 1

    students = []
    for row in range(data_start, ws.max_row + 1):
        class_val = str(ws.cell(row=row, column=indices["班级"] + 1).value or "").strip()
        class_val_clean = class_val.rstrip("班")
        if class_val_clean != class_name_clean:
            continue

        name = str(ws.cell(row=row, column=indices["姓名"] + 1).value or "").strip()
        gender = str(ws.cell(row=row, column=indices["性别"] + 1).value or "").strip()
        score_str = str(ws.cell(row=row, column=indices["总分"] + 1).value or "").strip()

        if not name:
            continue

        try:
            score = float(score_str)
        except ValueError:
            score = 0.0

        students.append({"name": name, "gender": gender, "score": score})

    wb.close()
    return students


def interleave_high_low(items):
    """将按分数降序排列的列表重新排列为 高-低-高-低 交错顺序。"""
    result = []
    left, right = 0, len(items) - 1
    while left <= right:
        result.append(items[left])
        left += 1
        if left <= right:
            result.append(items[right])
            right -= 1
    return result


def build_desk_groups(style_str, num_students):
    """根据样式字符串构建桌组大小列表。"""
    sizes = [int(c) for c in style_str]
    for s in sizes:
        if s < 1:
            print(f"错误：样式中的数字必须大于0，发现 '{s}'")
            sys.exit(1)

    row_capacity = sum(sizes)
    num_rows = max(1, (num_students + row_capacity - 1) // row_capacity)

    desk_groups = []
    for _ in range(num_rows):
        for size in sizes:
            desk_groups.append(size)

    return desk_groups


def is_male(gender):
    return gender in ("男", "M", "m", "male", "Male")


def is_female(gender):
    return gender in ("女", "F", "f", "female", "Female")


def allocate_students(students, desk_sizes, strategy, allow_mixed, seed):
    """将学生分配到桌组。"""
    if seed is not None:
        random.seed(seed)

    males = [s for s in students if is_male(s["gender"])]
    females = [s for s in students if is_female(s["gender"])]
    others = [s for s in students if not is_male(s["gender"]) and not is_female(s["gender"])]

    # 无法识别性别的归入人数多的一方
    if others:
        if len(males) >= len(females):
            males.extend(others)
        else:
            females.extend(others)

    # --- 允许混坐：所有学生一起排 ---
    if allow_mixed:
        all_students = males + females
        if strategy == "good-with-good":
            all_students.sort(key=lambda s: s["score"], reverse=True)
        elif strategy == "good-with-bad":
            all_students.sort(key=lambda s: s["score"], reverse=True)
            all_students = interleave_high_low(all_students)
        else:
            random.shuffle(all_students)

        groups = []
        idx = 0
        for size in desk_sizes:
            group = all_students[idx:idx + size]
            if group:
                groups.append(group)
            idx += size
            if idx >= len(all_students):
                break
        return groups

    # --- 性别分开排 ---
    if strategy == "good-with-good":
        males.sort(key=lambda s: s["score"], reverse=True)
        females.sort(key=lambda s: s["score"], reverse=True)
    elif strategy == "good-with-bad":
        males.sort(key=lambda s: s["score"], reverse=True)
        males = interleave_high_low(males)
        females.sort(key=lambda s: s["score"], reverse=True)
        females = interleave_high_low(females)
    else:
        random.shuffle(males)
        random.shuffle(females)

    groups = []
    m_idx, f_idx = 0, 0

    for size in desk_sizes:
        m_rem = len(males) - m_idx
        f_rem = len(females) - f_idx

        if m_rem >= f_rem:
            take = min(size, m_rem)
            group = males[m_idx:m_idx + take]
            m_idx += take
        else:
            take = min(size, f_rem)
            group = females[f_idx:f_idx + take]
            f_idx += take

        if group:
            groups.append(group)

    return groups


def generate_excel(groups, style_str, output_path, class_name):
    """生成座位表Excel文件。"""
    wb = Workbook()
    ws = wb.active
    ws.title = f"{class_name}座位表"

    sizes = [int(c) for c in style_str]
    num_groups_per_row = len(sizes)
    num_rows = (len(groups) + num_groups_per_row - 1) // num_groups_per_row

    total_cols = sum(sizes) + (num_groups_per_row - 1)

    thin_border = Border(
        left=Side(style="thin"), right=Side(style="thin"),
        top=Side(style="thin"), bottom=Side(style="thin"),
    )
    center_align = Alignment(horizontal="center", vertical="center", wrap_text=True)
    male_fill = PatternFill(start_color="D6EAF8", end_color="D6EAF8", fill_type="solid")
    female_fill = PatternFill(start_color="FADBD8", end_color="FADBD8", fill_type="solid")
    podium_fill = PatternFill(start_color="D5D8DC", end_color="D5D8DC", fill_type="solid")
    podium_font = Font(bold=True, size=14)
    name_font = Font(size=14)
    # score_font kept for reference but not used separately

    # --- 讲台行 ---
    ws.merge_cells(start_row=1, start_column=1, end_row=1, end_column=total_cols)
    podium_cell = ws.cell(row=1, column=1, value="讲  台")
    podium_cell.font = podium_font
    podium_cell.alignment = center_align
    podium_cell.fill = podium_fill
    for c in range(1, total_cols + 1):
        ws.cell(row=1, column=c).fill = podium_fill
        ws.cell(row=1, column=c).border = thin_border
    ws.row_dimensions[1].height = 35

    # --- 空行 ---
    # (直接将学生放在第3行及以下)

    # --- 学生座位 ---
    group_idx = 0
    for row_num in range(num_rows):
        excel_row = row_num + 3  # row 1 = 讲台, row 2 = blank
        col_offset = 1

        for g in range(num_groups_per_row):
            size = sizes[g]

            if group_idx < len(groups):
                group = groups[group_idx]
                group_idx += 1

                for i in range(size):
                    col = col_offset + i
                    cell = ws.cell(row=excel_row, column=col)

                    if i < len(group):
                        student = group[i]
                        cell.value = student['name']
                        cell.font = name_font
                        cell.alignment = center_align
                        if is_male(student["gender"]):
                            cell.fill = male_fill
                        elif is_female(student["gender"]):
                            cell.fill = female_fill
                    else:
                        cell.value = ""

                    cell.border = thin_border
                    ws.column_dimensions[get_column_letter(col)].width = 13
            else:
                for i in range(size):
                    col = col_offset + i
                    cell = ws.cell(row=excel_row, column=col)
                    cell.border = thin_border
                    ws.column_dimensions[get_column_letter(col)].width = 13

            col_offset += size

            # 过道
            if g < num_groups_per_row - 1:
                aisle_col = col_offset
                ws.column_dimensions[get_column_letter(aisle_col)].width = 4
                col_offset += 1

        ws.row_dimensions[excel_row].height = 36

    wb.save(output_path)
    print(f"座位表已保存至：{output_path}")


def main():
    args = parse_args()

    # 验证样式
    if not args.style.isdigit() or len(args.style) < 2:
        print("错误：样式必须为纯数字字符串，至少2位，如 '2222' 或 '3332'")
        sys.exit(1)

    input_path = Path(args.excel_file)
    if not input_path.exists():
        print(f"错误：文件不存在：{args.excel_file}")
        sys.exit(1)

    if args.count < 1:
        print("错误：生成个数必须大于0")
        sys.exit(1)

    # 确定输出路径模板
    if args.output:
        output_base = args.output
    else:
        output_base = f"{input_path.stem}_{args.class_name}_座位表.xlsx"

    # 读取学生（只读一次）
    print(f"正在读取 {args.excel_file} ...")
    students = read_students(args.excel_file, args.class_name)

    if not students:
        print(f"警告：班级 '{args.class_name}' 中没有找到任何学生，将生成空座位表。")
    else:
        m_count = sum(1 for s in students if is_male(s["gender"]))
        f_count = sum(1 for s in students if is_female(s["gender"]))
        print(f"班级 '{args.class_name}' 共有 {len(students)} 名学生（男：{m_count}，女：{f_count}）")

    # 构建桌组（只算一次）
    desk_sizes = build_desk_groups(args.style, len(students))

    # 确定基础种子
    base_seed = args.seed if args.seed is not None else random.randint(0, 999999)

    strategy_names = {
        "good-with-good": "好+好",
        "good-with-bad": "好+差",
        "random": "与成绩无关",
    }

    # 生成多个座位表
    for i in range(args.count):
        seed = base_seed + i

        # 输出文件名：多个时加序号后缀
        if args.count > 1:
            p = Path(output_base)
            output_path = str(p.parent / f"{p.stem}_{i + 1}{p.suffix}")
        else:
            output_path = output_base

        groups = allocate_students(students, desk_sizes, args.strategy, args.allow_mixed, seed)

        generate_excel(groups, args.style, output_path, args.class_name)

    # 摘要
    print(f"\n座位表配置：")
    print(f"  样式：{args.style}（{len(args.style)}组，每组{'-'.join(args.style)}人）")
    print(f"  策略：{strategy_names[args.strategy]}")
    print(f"  男女生同桌：{'允许' if args.allow_mixed else '不允许'}")
    print(f"  生成个数：{args.count}")
    print(f"  总桌组数：{len(groups)}，总排数：{(len(groups) + len(args.style) - 1) // len(args.style)}")


if __name__ == "__main__":
    main()