Coverage for o2/models/timetable/task_resource_distribution.py: 95%

1from dataclasses import dataclass, replace 

2from functools import cached_property 

3from itertools import groupby 

4from operator import itemgetter 

5from typing import TYPE_CHECKING, Optional 

6 

7from dataclass_wizard import JSONWizard 

8 

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 

13 

14if TYPE_CHECKING: 

15 from o2.models.timetable.timetable_type import TimetableType 

16 

17 

18@dataclass(frozen=True) 

19class ArrivalTimeDistribution(JSONWizard): 

20 """Distribution for arrival times of cases.""" 

21 

22 distribution_name: DISTRIBUTION_TYPE 

23 distribution_params: list[DistributionParameter] 

24 

25 

26@dataclass(frozen=True) 

27class TaskResourceDistribution(JSONWizard): 

28 """Distribution parameters for a specific resource assigned to a task.""" 

29 

30 resource_id: str 

31 distribution_name: str 

32 distribution_params: list[DistributionParameter] 

33 

34 

35@dataclass(frozen=True) 

36class TaskResourceDistributions(JSONWizard): 

37 """Collection of resource distributions for a specific task.""" 

38 

39 task_id: str 

40 resources: list[TaskResourceDistribution] 

41 

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 ) 

48 

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 ) 

55 

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. 

60 

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 

69 

70 new_resource = replace(original_distribution, resource_id=new_resource_id) 

71 return self.add_resource(new_resource) 

72 

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. 

77 

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 

90 

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 

102 

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) 

108 

109 return result 

110 

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]