src/maintenance_scheduling/domain.py

from dataclasses import dataclass, field
from datetime import date, timedelta
from typing import List, Optional, Annotated, Union, Set

from solverforge_legacy.solver import SolverStatus
from solverforge_legacy.solver.score import HardSoftScore
from solverforge_legacy.solver.domain import (
    planning_entity,
    planning_solution,
    PlanningId,
    PlanningVariable,
    PlanningEntityCollectionProperty,
    ProblemFactCollectionProperty,
    ProblemFactProperty,
    ValueRangeProvider,
    PlanningScore,
)
from .json_serialization import JsonDomainBase
from pydantic import Field


# ************************************************************************
# Utility functions for working day calculations
# ************************************************************************


def calculate_end_date(start_date: Optional[date], duration_in_days: int) -> Optional[date]:
    """
    Calculate the end date by adding working days to the start date, skipping weekends.

    The end date is exclusive (like Java's implementation).

    Args:
        start_date: The start date (inclusive)
        duration_in_days: Number of working days the job takes

    Returns:
        The end date (exclusive), or None if start_date is None
    """
    if start_date is None:
        return None

    # Skip weekends. Does not work for holidays.
    # Keep in sync with create_start_date_range().
    # Formula: For every 5 working days plus the weekday offset, add 2 weekend days
    # Python weekday(): Monday=0, Tuesday=1, ..., Sunday=6
    weekend_padding = 2 * ((duration_in_days + start_date.weekday()) // 5)
    return start_date + timedelta(days=duration_in_days + weekend_padding)


def count_working_days_between(start: date, end: date) -> int:
    """
    Count the number of working days (Mon-Fri) between two dates.

    Args:
        start: Start date (inclusive)
        end: End date (exclusive)

    Returns:
        Number of working days between the dates
    """
    if start >= end:
        return 0

    count = 0
    current = start
    while current < end:
        if current.weekday() < 5:  # Monday=0 to Friday=4
            count += 1
        current += timedelta(days=1)
    return count


def create_start_date_range(from_date: date, to_date: date) -> List[date]:
    """
    Generate a list of working days (Mon-Fri) within the date range.

    Args:
        from_date: Start of range (inclusive)
        to_date: End of range (exclusive)

    Returns:
        List of working days in the range
    """
    dates = []
    current = from_date
    while current < to_date:
        if current.weekday() < 5:  # Monday=0 to Friday=4
            dates.append(current)
        current += timedelta(days=1)
    return dates


# ************************************************************************
# Domain classes
# ************************************************************************


@dataclass
class WorkCalendar:
    """Defines the planning window for the schedule."""
    id: Annotated[str, PlanningId]
    from_date: date  # Inclusive
    to_date: date    # Exclusive


@dataclass
class Crew:
    """A maintenance crew that can be assigned to jobs."""
    id: Annotated[str, PlanningId]
    name: str

    def __hash__(self):
        return hash(self.id)

    def __eq__(self, other):
        if isinstance(other, Crew):
            return self.id == other.id
        return False

    def __str__(self):
        return f"{self.name}({self.id})"


@planning_entity
@dataclass
class Job:
    """
    A maintenance job that needs to be scheduled.

    Planning variables:
    - crew: The crew assigned to this job
    - start_date: When the job starts (inclusive)

    The end_date is computed from start_date and duration_in_days.
    """
    id: Annotated[str, PlanningId]
    name: str
    duration_in_days: int
    min_start_date: date    # Inclusive - earliest the job can start
    max_end_date: date      # Exclusive - latest the job can end
    ideal_end_date: date    # Exclusive - preferred end date
    tags: Set[str] = field(default_factory=set)

    # Planning variables
    crew: Annotated[Optional[Crew], PlanningVariable] = None
    start_date: Annotated[Optional[date], PlanningVariable] = None

    def get_end_date(self) -> Optional[date]:
        """Calculate the end date based on start_date and duration."""
        return calculate_end_date(self.start_date, self.duration_in_days)

    @property
    def end_date(self) -> Optional[date]:
        """End date property for convenient access."""
        return self.get_end_date()

    def calculate_overlap(self, other: "Job") -> int:
        """
        Calculate the number of overlapping working days with another job.

        Args:
            other: The other job to compare with

        Returns:
            Number of overlapping working days
        """
        if self.start_date is None or other.start_date is None:
            return 0

        self_end = self.get_end_date()
        other_end = other.get_end_date()

        if self_end is None or other_end is None:
            return 0

        # Calculate overlap range
        overlap_start = max(self.start_date, other.start_date)
        overlap_end = min(self_end, other_end)

        if overlap_start >= overlap_end:
            return 0

        return count_working_days_between(overlap_start, overlap_end)

    def get_common_tags(self, other: "Job") -> Set[str]:
        """Get the tags that both jobs share."""
        return self.tags & other.tags

    def __str__(self):
        return f"{self.name}({self.id})"


@planning_solution
@dataclass
class MaintenanceSchedule:
    """
    The planning solution containing the schedule to be optimized.
    """
    work_calendar: Annotated[WorkCalendar, ProblemFactProperty]
    crews: Annotated[List[Crew], ProblemFactCollectionProperty, ValueRangeProvider]
    jobs: Annotated[List[Job], PlanningEntityCollectionProperty]
    start_date_range: Annotated[
        List[date], ProblemFactCollectionProperty, ValueRangeProvider
    ] = field(default_factory=list)
    score: Annotated[Optional[HardSoftScore], PlanningScore] = None
    solver_status: SolverStatus = SolverStatus.NOT_SOLVING

    def __post_init__(self):
        """Initialize start_date_range if not provided."""
        if not self.start_date_range and self.work_calendar:
            self.start_date_range = create_start_date_range(
                self.work_calendar.from_date,
                self.work_calendar.to_date
            )


# ************************************************************************
# Pydantic REST models for API
# ************************************************************************


class WorkCalendarModel(JsonDomainBase):
    id: str
    from_date: str = Field(..., alias="fromDate")  # ISO date string
    to_date: str = Field(..., alias="toDate")


class CrewModel(JsonDomainBase):
    id: str
    name: str


class JobModel(JsonDomainBase):
    id: str
    name: str
    duration_in_days: int = Field(..., alias="durationInDays")
    min_start_date: str = Field(..., alias="minStartDate")
    max_end_date: str = Field(..., alias="maxEndDate")
    ideal_end_date: str = Field(..., alias="idealEndDate")
    tags: List[str] = Field(default_factory=list)
    crew: Optional[Union[str, CrewModel]] = None
    start_date: Optional[str] = Field(None, alias="startDate")
    end_date: Optional[str] = Field(None, alias="endDate")


class MaintenanceScheduleModel(JsonDomainBase):
    work_calendar: WorkCalendarModel = Field(..., alias="workCalendar")
    crews: List[CrewModel]
    jobs: List[JobModel]
    start_date_range: List[str] = Field(default_factory=list, alias="startDateRange")
    score: Optional[str] = None
    solver_status: Optional[str] = Field(None, alias="solverStatus")