from ortools.sat.python import cp_model

import csv

# Определяем даты апреля 2026 года
days = [
    "2026-04-{:02d}".format(i) for i in range(1, 31)
]

# Смены
shifts = [
    (6, 14.5),   # с 06:00 до 14:30
    (8, 16.5),     # с 08:00 до 16:30
    (11.5, 20),    # с 11:30 до 20:00
    (12.5, 21)      # с 12:30 до 21:00
]

# Операторы и их пожелания
operators = [
    "Иванов",
    "Алексеев",
    "Сидоров",
    "Козобородов",
    "Трудяшкина",
    "Петрова"
]

operator_preferences = {
    "Иванов": ["2026-04-10", "2026-04-13"],
    "Алексеев": [(i, 0) for i in days if (int(i[-2:]) - 1) % 7 == 0],  # Понедельники
    "Сидоров": ["2026-04-25", "2026-04-26"],
    "Козобородов": [(i, 2) for i in days],  # Предпочитает смену 11:30–20:00
    "Трудяшкина": ["2026-04-04"],
    "Петрова": ["2026-04-08", "2026-04-29"]
}

# Создаем модель
model = cp_model.CpModel()

# Переменные: определяем, работает ли оператор на смене в день
works = {}
for day in days:
    for shift_idx in range(len(shifts)):
        for op in operators:
            works[(day, shift_idx, op)] = model.NewBoolVar(f"{op}_{day}_shift_{shift_idx}")

# # Пожелания операторов
for op, prefs in operator_preferences.items():
    if isinstance(prefs, list):
        for pref in prefs:
            if isinstance(pref, tuple):
                day, shift = pref
                model.AddHint(works[day, shift, op], 1)
            else:
                day = pref
                model.Add(sum(works[day, shift_idx, op] for shift_idx in range(len(shifts))) == 0)

# Условие 1: На каждую смену должен быть назначен оператор
for day in days:
    for shift_idx in range(len(shifts)):
        model.Add(sum(works[day, shift_idx, op] for op in operators) == 1)

# Условие 2: У каждого оператора должно быть 8 выходных
for op in operators:
    model.Add(sum(works[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(works[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(works[day, shift_idx, op] for shift_idx in range(len(shifts))) <= 1)

# Решаем модель
solver = cp_model.CpSolver()
status = solver.Solve(model)

if status != cp_model.OPTIMAL and status != cp_model.FEASIBLE:
    print("Невозможно составить график")
else:
    # Создаем CSV файл с результатами
    with open('shift_schedule.csv', 'w', newline='', encoding='utf-8') as file:
        writer = csv.writer(file)
        writer.writerow(['Day'] + [f'{i+1} shift' for i in range(len(shifts))])
        
        for day in days:
            row = [day]
            for shift_idx in range(len(shifts)):
                op = None
                for op_name in operators:
                    if solver.Value(works[day, shift_idx, op_name]) == 1:
                        op = op_name
                        break
                row.append(op)
            writer.writerow(row)

    print("График успешно составлен и сохранен в 'shift_schedule.csv'")