OperatorSchedule/models.py

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