"""Employee scheduling with Heusor.

Every shift of a hospital's four weeks needs an employee with the shift's
skill. An employee works one shift at a time, at most one a day, with 10
hours' rest between shifts, and not on a day they are unavailable. They would
rather not work on days they marked undesired, would like the days they
marked desired, and the work is shared out evenly.

Run it:  python employee_scheduling.py
"""

from __future__ import annotations

import copy
import random
from dataclasses import dataclass, field
from datetime import date, datetime, time, timedelta
from typing import Annotated

import heusor
from heusor import (
    ConstraintCollectors,
    ConstraintFactory,
    HardSoftScore,
    Joiners,
    PlanningPin,
    PlanningVariable,
    constraint_provider,
    planning_entity,
    planning_solution,
    problem_fact,
)

# ─── The model ──────────────────────────────────────────────────────────────


@problem_fact
@dataclass
class Employee:
    name: str
    skills: set[str] = field(default_factory=set)
    unavailable_dates: set[date] = field(default_factory=set)
    undesired_dates: set[date] = field(default_factory=set)
    desired_dates: set[date] = field(default_factory=set)


@planning_entity
@dataclass
class Shift:
    id: int
    start: datetime
    end: datetime
    location: str
    required_skill: str
    employee: Annotated[Employee | None, PlanningVariable] = None
    pinned: Annotated[bool, PlanningPin] = False


@planning_solution
@dataclass
class Schedule:
    employees: list[Employee]
    shifts: list[Shift]
    score: HardSoftScore | None = None


# ─── The constraints ────────────────────────────────────────────────────────

MINUTE = timedelta(minutes=1)


def minutes_on(shift: Shift, day: date) -> int:
    """Minutes of the shift that fall on ``day``; zero or less when none do."""
    return (
        min(shift.end, datetime.combine(day, time(0)) + timedelta(days=1))
        - max(shift.start, datetime.combine(day, time(0)))
    ) // MINUTE


@constraint_provider
def define_constraints(f: ConstraintFactory):
    return [
        # Hard: the employee has the skill the shift needs.
        f.for_each(Shift)
        .filter(lambda shift: shift.required_skill not in shift.employee.skills)
        .penalize(HardSoftScore.ONE_HARD)
        .as_constraint("Missing required skill"),
        # Hard: one shift at a time — penalized by the minutes of overlap.
        f.for_each_unique_pair(
            Shift,
            Joiners.equal(lambda shift: shift.employee),
            Joiners.overlapping(lambda shift: shift.start, lambda shift: shift.end),
        )
        .penalize(
            HardSoftScore.ONE_HARD,
            lambda first, second: (min(first.end, second.end) - max(first.start, second.start)) // MINUTE,
        )
        .as_constraint("Overlapping shift"),
        # Hard: 10 hours' rest — penalized by the minutes short of it.
        f.for_each(Shift)
        .join(
            Shift,
            Joiners.equal(lambda shift: shift.employee),
            Joiners.less_than_or_equal(lambda shift: shift.end, lambda shift: shift.start),
        )
        .filter(lambda first, second: second.start - first.end < timedelta(hours=10))
        .penalize(HardSoftScore.ONE_HARD, lambda first, second: 600 - (second.start - first.end) // MINUTE)
        .as_constraint("At least 10 hours between 2 shifts"),
        # Hard: at most one shift starting on a day.
        f.for_each_unique_pair(
            Shift,
            Joiners.equal(lambda shift: shift.employee),
            Joiners.equal(lambda shift: shift.start.date()),
        )
        .penalize(HardSoftScore.ONE_HARD)
        .as_constraint("Max one shift per day"),
        # Hard: not on an unavailable day — penalized by the minutes on it.
        f.for_each(Shift)
        .flatten_last(lambda shift: shift.employee.unavailable_dates)
        .filter(lambda shift, day: minutes_on(shift, day) > 0)
        .penalize(HardSoftScore.ONE_HARD, minutes_on)
        .as_constraint("Unavailable employee"),
        # Soft: rather not on an undesired day, gladly on a desired one.
        f.for_each(Shift)
        .flatten_last(lambda shift: shift.employee.undesired_dates)
        .filter(lambda shift, day: minutes_on(shift, day) > 0)
        .penalize(HardSoftScore.ONE_SOFT, minutes_on)
        .as_constraint("Undesired day for employee"),
        f.for_each(Shift)
        .flatten_last(lambda shift: shift.employee.desired_dates)
        .filter(lambda shift, day: minutes_on(shift, day) > 0)
        .reward(HardSoftScore.ONE_SOFT, minutes_on)
        .as_constraint("Desired day for employee"),
        # Soft: share the shifts out evenly, counting employees with none.
        f.for_each(Shift)
        .group_by(lambda shift: shift.employee, ConstraintCollectors.count())
        .complement(Employee)
        .group_by(ConstraintCollectors.load_balance(scale=1000))
        .penalize(HardSoftScore.ONE_SOFT, lambda balance: balance.unfairness)
        .as_constraint("Balance employee shift assignments"),
    ]


# ─── The data ───────────────────────────────────────────────────────────────

FIRST_DAY = date(2026, 2, 2)  # a Monday
DAYS = 28
WARDS = ["Emergency", "Surgery", "Pediatrics"]
# (start hour, length in hours, nurses, doctors) for each shift of a day.
SHIFT_PATTERN = [(6, 8, 2, 1), (14, 8, 2, 1), (22, 8, 2, 0)]
NURSES, DOCTORS = 32, 10
# Each employee marks this many days of the four weeks of each kind.
UNAVAILABLE, UNDESIRED, DESIRED = 3, 3, 3


def generate(seed: int = 0) -> Schedule:
    """Four weeks of a hospital: 3 wards, each staffed day and night (early,
    late and night shifts of 8 hours, 2 nurses each, and a doctor on the early
    and late shifts): 672 shifts, for 32 nurses and 10 doctors. A nurse in
    Pediatrics needs the "Pediatrics" skill, which a third of the nurses
    have. Every employee marks 3 days unavailable, 3 undesired and 3 desired."""
    rng = random.Random(seed)
    days = [FIRST_DAY + timedelta(days=i) for i in range(DAYS)]
    employees = []
    for i in range(NURSES + DOCTORS):
        skills = {"Nurse"} if i < NURSES else {"Doctor"}
        if i < NURSES and i % 3 == 0:
            skills.add("Pediatrics")
        marked = rng.sample(days, UNAVAILABLE + UNDESIRED + DESIRED)
        employees.append(
            Employee(
                f"{'Nurse' if i < NURSES else 'Doctor'} {i + 1 if i < NURSES else i - NURSES + 1}",
                skills,
                set(marked[:UNAVAILABLE]),
                set(marked[UNAVAILABLE : UNAVAILABLE + UNDESIRED]),
                set(marked[UNAVAILABLE + UNDESIRED :]),
            )
        )
    shifts = []
    for day in days:
        for ward in WARDS:
            for hour, length, nurses, doctors in SHIFT_PATTERN:
                start = datetime.combine(day, time(hour))
                end = start + timedelta(hours=length)
                nurse_skill = "Pediatrics" if ward == "Pediatrics" else "Nurse"
                for skill in [nurse_skill] * nurses + ["Doctor"] * doctors:
                    shifts.append(Shift(len(shifts), start, end, ward, skill))
    return Schedule(employees, shifts)


# ─── Solving and explaining ─────────────────────────────────────────────────


def solve(schedule: Schedule, time_limit: float | None = 30, step_limit: int | None = None, seed: int = 0) -> Schedule:
    return heusor.solve(schedule, define_constraints, time_limit=time_limit, step_limit=step_limit, seed=seed)


def explain(schedule: Schedule) -> None:
    explanation = heusor.score(schedule, define_constraints)
    print(f"Score: {explanation.score}")
    for name, score in explanation.constraints.items():
        print(f"  {name:36} {score}")
    worked = {e.name: 0 for e in schedule.employees}
    for shift in schedule.shifts:
        if shift.employee is not None:
            worked[shift.employee.name] += 1
    print(f"  Shifts per employee: {min(worked.values())} to {max(worked.values())}")


# ─── A sick call ────────────────────────────────────────────────────────────


def sick_call(
    schedule: Schedule,
    name: str,
    first_day: date,
    days: int = 3,
    time_limit: float | None = 10,
    step_limit: int | None = None,
    seed: int = 0,
) -> Schedule:
    """``name`` calls in sick from ``first_day`` for ``days`` days. Their shifts
    in that time are given to someone else; every other shift is pinned, so no
    one else's rota changes."""
    plan = copy.deepcopy(schedule)
    sick = next(e for e in plan.employees if e.name == name)
    off = {first_day + timedelta(days=i) for i in range(days)}
    sick.unavailable_dates |= off
    for shift in plan.shifts:
        if shift.employee is sick and shift.start.date() in off:
            shift.employee, shift.pinned = None, False
        else:
            shift.pinned = True
    return heusor.solve(plan, define_constraints, time_limit=time_limit, step_limit=step_limit, seed=seed)


if __name__ == "__main__":
    schedule = solve(generate())
    explain(schedule)
    sick_day = FIRST_DAY + timedelta(days=10)
    # Whoever works most in those three days calls in sick.
    window = {sick_day + timedelta(days=i) for i in range(3)}
    busy = [s.employee.name for s in schedule.shifts if s.start.date() in window]
    name = max(sorted(set(busy)), key=busy.count)
    covered = sick_call(schedule, name, sick_day)
    print(f"\nAfter {name} calls in sick for {sick_day:%a %d %b} and the 2 days after:")
    explain(covered)
    for before, after in zip(schedule.shifts, covered.shifts):
        if before.employee is not after.employee and before.employee.name != after.employee.name:
            print(f"  {before.start:%a %d %b %H:%M} {before.location}: {before.employee.name} → {after.employee.name}")