JITToolBox/module/picker/schema.py

import random
from typing import Optional

from openpyxl.reader.excel import load_workbook
from openpyxl.workbook.workbook import Workbook
from openpyxl.worksheet.worksheet import Worksheet


class PickerStudent:
    total_time: int = 0

    def __init__(self, name: str, so: str, position: int, scores: list[int]):
        self._name = name
        self._so = str(so)
        self._position = position
        self._scores = scores
        self._modify = False
        self._scores = []
        self._scores.extend(scores)
        self._scores_pointer = -1

        self.init_score()

    def init_score(self):
        self._scores += [0] * (self.total_time - len(self._scores))

        for idx, val in enumerate(self._scores):
            if val is None:
                self._scores[idx] = 0
            if self._scores[idx] == 0:
                self._scores_pointer = idx
                break

    def append_score(self, score: int) -> int:
        if self._scores_pointer > self.total_time or self._scores_pointer < 0:
            raise IndexError
        self._scores[self._scores_pointer] = score
        self._scores_pointer += 1
        self.modified()
        return self._scores_pointer - 1

    def change_score(self, new_score: int, index: int):
        if index < 0 or index > self.total_time:
            raise IndexError
        self._scores[index] = new_score
        self.modified()

    @staticmethod
    def pick(student: list['PickerStudent']) -> Optional['PickerStudent']:
        filtered = [item for item in student if item.weight != 100 and item.weight > 0]
        if not filtered:
            return None

        # 计算倒数权重
        weights = [1 / item.weight for item in filtered]
        total = sum(weights)
        r = random.uniform(0, total)

        cumulative = 0
        for item, w in zip(filtered, weights):
            cumulative += w
            if r <= cumulative:
                return item
        return None

    @classmethod
    def set_total_time(cls, total_time: int):
        cls.total_time = total_time

    @property
    def position(self) -> int:
        return self._position

    @property
    def scores(self) -> list[int]:
        return self._scores

    @property
    def modify(self) -> bool:
        return self._modify

    @property
    def weight(self) -> int:
        return int(sum(1 for x in self._scores if x != 0) / self.total_time * 100)

    @property
    def name(self) -> str:
        return self._name

    @property
    def so(self) -> str:
        return self._so

    def saved(self):
        self._modify = False

    def modified(self):
        self._modify = True

    def __str__(self):
        return f"PickerStudent {self._name} at {self.position}, with scores: {','.join(str(x) for x in self._scores if x)}"

    def __repr__(self):
        return self.__str__()


class PickerExcel:
    wb: Optional[Workbook]
    ws: Optional[Worksheet]
    path: str = ''
    max_time_position = 'M1'
    start_row = 5

    def __init__(self, path: str):
        self.open(path)

    @classmethod
    def open(cls, path: str):
        cls.path = path
        cls.wb = load_workbook(path, keep_vba=True)
        cls.ws = cls.wb.active
        PickerStudent.set_total_time(cls.ws[cls.max_time_position].value)

    @classmethod
    def read_student(cls, path: Optional[str] = None) -> list[PickerStudent]:
        if path:
            cls.open(path)

        if cls.wb and cls.ws:
            ret = []
            for index, row in enumerate(cls.ws.iter_rows(
                    min_row=cls.start_row,
                    max_col=4 + cls.ws[cls.max_time_position].value,
                    values_only=True)
            ):
                if (name := row[2]) is None:
                    break
                ret.append(PickerStudent(name, row[1], cls.start_row + index, row[4:]))
            return ret
        raise Exception('No Workbook or Worksheet')

    @classmethod
    def read_total_time(cls, path: Optional[str] = None) -> int:
        if path:
            cls.open(path)

        if cls.wb and cls.ws:
            val = cls.ws[cls.max_time_position].value
            if isinstance(val, int):
                return val
            raise Exception(f'总次数读取错误，期待类型 {type(1)} 但实际为 {type(val)}\n'
                            f'可能的解决方法：\n'
                            f'1、确保总次数位于 \'{cls.max_time_position}\' 单元格；\n'
                            f'2、确保选择了正确的 Excel 文件。')
        raise Exception('No Workbook or Worksheet')

    @classmethod
    def save_total_time(cls, path: Optional[str] = None, value: Optional[int] = None) -> None:
        if path:
            cls.open(path)

        if cls.wb and cls.ws and value:
            cls.ws[cls.max_time_position].value = value
            cls.save()

    @classmethod
    def write_back(cls, student: PickerStudent):
        start_col = 5  # E列
        row = student.position

        for i, val in enumerate(student.scores):
            cls.ws.cell(row=row, column=start_col + i, value=val if val != 0 else '')

        student.saved()
        cls.save()

    @classmethod
    def write_all(cls, students: list[PickerStudent]):
        for student in students:
            if student.modify:
                cls.write_back(student)

    @classmethod
    def save(cls):
        if cls.ws:
            cls.wb.save(cls.path)


if __name__ == '__main__':
    ...