Coverage for o2/models/timetable/task_resource_distribution.py: 95%
59 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-05-16 11:18 +0000
« prev ^ index » next coverage.py v7.6.12, created at 2025-05-16 11:18 +0000
1from dataclasses import dataclass, replace
2from functools import cached_property
3from itertools import groupby
4from operator import itemgetter
5from typing import TYPE_CHECKING, Optional
7from dataclass_wizard import JSONWizard
9from o2.models.timetable.distribution_parameter import DistributionParameter
10from o2.models.timetable.distribution_type import DISTRIBUTION_TYPE
11from o2.models.timetable.time_period import TimePeriod
12from o2.util.bit_mask_helper import find_most_frequent_overlap
14if TYPE_CHECKING:
15 from o2.models.timetable.timetable_type import TimetableType
18@dataclass(frozen=True)
19class ArrivalTimeDistribution(JSONWizard):
20 """Distribution for arrival times of cases."""
22 distribution_name: DISTRIBUTION_TYPE
23 distribution_params: list[DistributionParameter]
26@dataclass(frozen=True)
27class TaskResourceDistribution(JSONWizard):
28 """Distribution parameters for a specific resource assigned to a task."""
30 resource_id: str
31 distribution_name: str
32 distribution_params: list[DistributionParameter]
35@dataclass(frozen=True)
36class TaskResourceDistributions(JSONWizard):
37 """Collection of resource distributions for a specific task."""
39 task_id: str
40 resources: list[TaskResourceDistribution]
42 def remove_resource(self, resource_id: str) -> "TaskResourceDistributions":
43 """Remove a resource from the distribution."""
44 return replace(
45 self,
46 resources=[resource for resource in self.resources if resource.resource_id != resource_id],
47 )
49 def add_resource(self, distribution: TaskResourceDistribution) -> "TaskResourceDistributions":
50 """Add a resource to the distribution."""
51 return replace(
52 self,
53 resources=self.resources + [distribution],
54 )
56 def add_resource_based_on_original(
57 self, original_resource_id: str, new_resource_id: str
58 ) -> "TaskResourceDistributions":
59 """Add a resource based on an original resource.
61 If the original resource is not found, the distribution is not changed.
62 """
63 original_distribution = next(
64 (resource for resource in self.resources if resource.resource_id == original_resource_id),
65 None,
66 )
67 if original_distribution is None:
68 return self
70 new_resource = replace(original_distribution, resource_id=new_resource_id)
71 return self.add_resource(new_resource)
73 def get_highest_availability_time_period(
74 self, timetable: "TimetableType", min_hours: int
75 ) -> Optional[TimePeriod]:
76 """Get the highest availability time period for the task.
78 The highest availability time period is the time period with the highest
79 frequency of availability.
80 If no overlapping time periods is found, it will return the longest non-overlapping
81 time period.
82 """
83 resources_assigned_to_task = timetable.get_resources_assigned_to_task(self.task_id)
84 calendars = [timetable.get_calendar_for_resource(resource) for resource in resources_assigned_to_task]
85 bitmasks_by_day = [
86 bitmask for calendar in calendars if calendar is not None for bitmask in calendar.bitmasks_by_day
87 ]
88 if len(bitmasks_by_day) == 0:
89 return None
91 bitmasks_sorted_by_day = sorted(bitmasks_by_day, key=itemgetter(0))
92 bitmasks_grouped_by_day = groupby(bitmasks_sorted_by_day, key=itemgetter(0))
93 max_frequency = 0
94 max_size = 0
95 result = None
96 for day, bitmask_pairs in bitmasks_grouped_by_day:
97 bitmasks = [pair[1] for pair in bitmask_pairs]
98 overlap = find_most_frequent_overlap(bitmasks, min_size=min_hours)
99 if overlap is None:
100 continue
101 frequency, start, stop = overlap
103 size = stop - start
104 if frequency > max_frequency and size > max_size:
105 max_frequency = frequency
106 max_size = size
107 result = TimePeriod.from_start_end(start, stop, day)
109 return result
111 @cached_property
112 def resource_ids(self) -> list[str]:
113 """Get the resource ids in the distribution."""
114 return [resource.resource_id for resource in self.resources]