294 lines
12 KiB
Python
294 lines
12 KiB
Python
from __future__ import annotations
|
||
from dataclasses import dataclass, field
|
||
from datetime import datetime
|
||
from enum import Enum
|
||
from typing import List, Dict, Set, Union, Optional
|
||
from ortools.sat.python import cp_model
|
||
from config import YEAR, MONTH, SHIFTS, OPERATORS
|
||
|
||
class PreferenceType(Enum):
|
||
"""Типы предпочтений"""
|
||
WORK_ON_DATE = "work_on_date"
|
||
WORK_ON_DAY_OF_WEEK = "work_on_day_of_week"
|
||
WORK_ONLY_SHIFT = "work_only_shift"
|
||
WORK_ONLY_DATE = "work_only_date"
|
||
OFF_ON_DATE = "off_on_date"
|
||
|
||
class DayOfWeek(Enum):
|
||
"""Дни недели"""
|
||
MONDAY = 0
|
||
TUESDAY = 1
|
||
WEDNESDAY = 2
|
||
THURSDAY = 3
|
||
FRIDAY = 4
|
||
SATURDAY = 5
|
||
SUNDAY = 6
|
||
|
||
@dataclass
|
||
class OperatorPreferences:
|
||
"""
|
||
Класс для управления предпочтениями оператора.
|
||
Использует паттерн Builder для удобного создания.
|
||
"""
|
||
operator_name: str
|
||
work_on_dates: List[int] = field(default_factory=list) # Дни месяца (1-31)
|
||
work_on_day_of_week: List[DayOfWeek] = field(default_factory=list)
|
||
work_only_shift: Optional[int] = None # Индекс смены (0-3)
|
||
off_on_dates: List[int] = field(default_factory=list) # Дни, когда оператор НЕ работает
|
||
|
||
def __post_init__(self):
|
||
if self.work_only_shift is not None and not (0 <= self.work_only_shift < len(SHIFTS)):
|
||
raise ValueError(f"Индекс смены {self.work_only_shift} вне диапазона 0-{len(SHIFTS)-1}")
|
||
|
||
@classmethod
|
||
def create(cls, operator_name: str) -> 'PreferencesBuilder':
|
||
"""Статический метод для начала построения предпочтений"""
|
||
return PreferencesBuilder(operator_name)
|
||
|
||
def to_dict(self) -> Dict[str, Union[int, str, List[int]]]:
|
||
"""
|
||
Конвертирует объект предпочтений в формат, пригодный для config.py.
|
||
Возвращает словарь с ключами, понятными для парсера в create_model.
|
||
"""
|
||
result = {}
|
||
|
||
if self.work_only_shift is not None:
|
||
result['shift'] = self.work_only_shift
|
||
|
||
if self.work_on_day_of_week:
|
||
day_names = [d.name.lower() for d in self.work_on_day_of_week]
|
||
# Если только один день недели, можно хранить как строку, иначе как список
|
||
if len(day_names) == 1:
|
||
result['day_of_week'] = day_names[0]
|
||
else:
|
||
result['day_of_week'] = day_names
|
||
|
||
if self.work_on_dates:
|
||
result['dates'] = self.work_on_dates
|
||
|
||
if self.off_on_dates:
|
||
result['off_dates'] = self.off_on_dates
|
||
|
||
return result
|
||
|
||
def __str__(self) -> str:
|
||
"""Метод для отладки"""
|
||
parts = []
|
||
if self.work_only_shift is not None:
|
||
parts.append(f"Shift {self.work_only_shift}")
|
||
if self.work_on_day_of_week:
|
||
parts.append(f"Days: {', '.join([d.name for d in self.work_on_day_of_week])}")
|
||
if self.work_on_dates:
|
||
parts.append(f"Dates: {self.work_on_dates}")
|
||
if self.off_on_dates:
|
||
parts.append(f"Off: {self.off_on_dates}")
|
||
|
||
return f"OperatorPreferences({self.operator_name}: {', '.join(parts)})"
|
||
|
||
class PreferencesBuilder:
|
||
"""Builder для создания предпочтений оператора"""
|
||
|
||
def __init__(self, operator_name: str):
|
||
self.operator_name = operator_name
|
||
self.prefs = OperatorPreferences(operator_name)
|
||
|
||
def work_on_date(self, day: int) -> 'PreferencesBuilder':
|
||
"""Добавить день месяца для работы"""
|
||
if not 1 <= day <= 31:
|
||
raise ValueError(f"День {day} вне диапазона 1-31")
|
||
self.prefs.work_on_dates.append(day)
|
||
return self
|
||
|
||
def work_on_dates(self, days: List[int]) -> 'PreferencesBuilder':
|
||
"""Добавить список дней месяца для работы"""
|
||
for day in days:
|
||
self.work_on_date(day)
|
||
return self
|
||
|
||
def work_on_day_of_week(self, day: DayOfWeek) -> 'PreferencesBuilder':
|
||
"""Добавить день недели для работы"""
|
||
if day not in self.prefs.work_on_day_of_week:
|
||
self.prefs.work_on_day_of_week.append(day)
|
||
return self
|
||
|
||
def work_on_days_of_week(self, days: List[DayOfWeek]) -> 'PreferencesBuilder':
|
||
"""Добавить список дней недели для работы"""
|
||
for day in days:
|
||
self.work_on_day_of_week(day)
|
||
return self
|
||
|
||
def work_only_shift(self, shift_idx: int) -> 'PreferencesBuilder':
|
||
"""Установить предпочтительную смену (работает только в этой смене)"""
|
||
self.prefs.work_only_shift = shift_idx
|
||
return self
|
||
|
||
def off_on_date(self, day: int) -> 'PreferencesBuilder':
|
||
"""Установить день, когда оператор НЕ работает"""
|
||
if not 1 <= day <= 31:
|
||
raise ValueError(f"День {day} вне диапазона 1-31")
|
||
self.prefs.off_on_dates.append(day)
|
||
return self
|
||
|
||
def off_on_dates(self, days: List[int]) -> 'PreferencesBuilder':
|
||
"""Установить дни, когда оператор НЕ работает"""
|
||
for day in days:
|
||
self.off_on_date(day)
|
||
return self
|
||
|
||
def build(self) -> OperatorPreferences:
|
||
"""Завершить построение и вернуть объект предпочтений"""
|
||
return self.prefs
|
||
|
||
def get_days() -> List[str]:
|
||
"""Генерирует список дат месяца"""
|
||
return [f"{YEAR}-{MONTH:02d}-{i:02d}" for i in range(1, 31)]
|
||
|
||
def create_model() -> Tuple[cp_model.CpModel, Dict, Set]:
|
||
"""
|
||
Создаёт и настраивает CP-модель.
|
||
|
||
Возвращает:
|
||
model: CP-модель
|
||
variables: словарь переменных работы
|
||
operators_set: множество операторов для удобства
|
||
"""
|
||
model = cp_model.CpModel()
|
||
variables = {}
|
||
days = get_days()
|
||
|
||
# Определяем переменные: works[(day, shift_idx, op)]
|
||
for day in days:
|
||
for shift_idx in range(len(SHIFTS)):
|
||
for op in OPERATORS:
|
||
var_name = f"{op}_{day}_shift_{shift_idx}"
|
||
variables[(day, shift_idx, op)] = model.NewBoolVar(var_name)
|
||
|
||
# Определяем множество операторов
|
||
operators_set = set(OPERATORS)
|
||
|
||
# 1. Учёт предпочтений операторов
|
||
# Ожидаем, что OPERATOR_PREFERENCES теперь содержит объекты OperatorPreferences или их словари
|
||
for op, prefs_data in OPERATOR_PREFERENCES.items():
|
||
# Если это словарь (из config.py), преобразуем в объект
|
||
if isinstance(prefs_data, dict):
|
||
prefs = OperatorPreferences(op)
|
||
if 'shift' in prefs_data:
|
||
prefs.work_only_shift = prefs_data['shift']
|
||
if 'day_of_week' in prefs_data:
|
||
dow = prefs_data['day_of_week']
|
||
if isinstance(dow, str):
|
||
try:
|
||
prefs.work_on_day_of_week.append(DayOfWeek[dow.upper()])
|
||
except KeyError:
|
||
pass
|
||
elif isinstance(dow, list):
|
||
for d in dow:
|
||
try:
|
||
prefs.work_on_day_of_week.append(DayOfWeek[d.upper()])
|
||
except KeyError:
|
||
pass
|
||
if 'dates' in prefs_data:
|
||
prefs.work_on_dates = prefs_data['dates']
|
||
if 'off_dates' in prefs_data:
|
||
prefs.off_on_dates = prefs_data['off_dates']
|
||
else:
|
||
prefs = prefs_data
|
||
|
||
# Применяем предпочтения
|
||
if prefs.work_only_shift is not None:
|
||
# Оператор предпочитает конкретную смену
|
||
shift_idx = prefs.work_only_shift
|
||
for day in days:
|
||
model.AddHint(variables[(day, shift_idx, op)], 1)
|
||
|
||
if prefs.work_on_day_of_week:
|
||
# Оператор работает только в определенные дни недели
|
||
for day in days:
|
||
date_obj = datetime(int(day.split('-')[0]), int(day.split('-')[1]), int(day.split('-')[2]))
|
||
if date_obj.weekday() in [d.value for d in prefs.work_on_day_of_week]:
|
||
# Если есть preference_shift, то только в эту смену
|
||
if prefs.work_only_shift is not None:
|
||
model.AddHint(variables[(day, prefs.work_only_shift, op)], 1)
|
||
else:
|
||
# Иначе в любую смену (но это уже обрабатывается другими ограничениями)
|
||
# Для hint можно добавить в первую смену как ориентир
|
||
model.AddHint(variables[(day, 0, op)], 1)
|
||
|
||
if prefs.work_on_dates:
|
||
# Оператор предпочитает конкретные дни
|
||
for pref_day_num in prefs.work_on_dates:
|
||
target_day = f"{YEAR}-{MONTH:02d}-{pref_day_num:02d}"
|
||
if target_day in days:
|
||
if prefs.work_only_shift is not None:
|
||
model.AddHint(variables[(target_day, prefs.work_only_shift, op)], 1)
|
||
else:
|
||
model.AddHint(variables[(target_day, 0, op)], 1)
|
||
|
||
if prefs.off_on_dates:
|
||
# Оператор НЕ работает в эти дни
|
||
for off_day_num in prefs.off_on_dates:
|
||
target_day = f"{YEAR}-{MONTH:02d}-{off_day_num:02d}"
|
||
if target_day in days:
|
||
# Запрещаем работу в этот день
|
||
for shift_idx in range(len(SHIFTS)):
|
||
model.Add(variables[(target_day, shift_idx, op)] == 0)
|
||
|
||
# Условие 1: На каждую смену должен быть назначен оператор
|
||
for day in days:
|
||
for shift_idx in range(len(SHIFTS)):
|
||
model.Add(sum(variables[(day, shift_idx, op)] for op in OPERATORS) == 1)
|
||
|
||
# Условие 2: У каждого оператора должно быть 8 выходных (20 рабочих дней)
|
||
for op in OPERATORS:
|
||
model.Add(sum(variables[(day, shift_idx, op)] for day in days for shift_idx in range(len(SHIFTS))) == 20)
|
||
|
||
# Условие 3: Оператор не должен работать больше 5 дней подряд
|
||
for op in OPERATORS:
|
||
for day_idx in range(len(days) - 4):
|
||
model.Add(sum(variables[(days[day_idx + i], shift_idx, op)]
|
||
for i in range(5)
|
||
for shift_idx in range(len(SHIFTS))) <= 4)
|
||
|
||
# Условие 4: В день оператор может работать только одну смену
|
||
for day in days:
|
||
for op in OPERATORS:
|
||
model.Add(sum(variables[(day, shift_idx, op)] for shift_idx in range(len(SHIFTS))) <= 1)
|
||
|
||
return model, variables, operators_set
|
||
|
||
def solve_model() -> bool:
|
||
"""
|
||
Решает модель.
|
||
|
||
Возвращает:
|
||
True - если решение найдено
|
||
False - если не удалось составить график
|
||
"""
|
||
model, variables, operators_set = create_model()
|
||
|
||
solver = cp_model.CpSolver()
|
||
status = solver.Solve(model)
|
||
|
||
return status in (cp_model.OPTIMAL, cp_model.FEASIBLE)
|
||
|
||
def get_solution() -> Dict:
|
||
"""
|
||
Получает решение модели.
|
||
|
||
Возвращает:
|
||
Словарь с расписанием: {день: {смена: {оператор: True/False}}}
|
||
"""
|
||
model, variables, operators_set = create_model()
|
||
solver = cp_model.CpSolver()
|
||
solver.Solve(model)
|
||
|
||
schedule = {}
|
||
for day in get_days():
|
||
schedule[day] = {
|
||
shift_idx: {op: solver.Value(variables[(day, shift_idx, op)])
|
||
for op in OPERATORS}
|
||
for shift_idx in range(len(SHIFTS))
|
||
}
|
||
|
||
return schedule
|