scheduler/backend/scheduler.py

from __future__ import annotations
import os

from dataclasses import dataclass
from datetime import date, timedelta
from test_window import get_shift_sequence_with_capacity, get_shift_capacity_for_date, is_off_day, next_window_start_date
from test_bundle import Test, TestBundle, build_test_bundles, update_mirrored_bundle_priorities

DUT = (os.getenv("DUT") or "DUT").strip()
REF = (os.getenv("REF") or "REF").strip()

@dataclass(frozen=True)
class ScheduleEntry:
    test_id: str
    device: str
    scheduled_date: str
    shift_index: int
    sequence_in_shift: int


@dataclass
class ScheduleWindow:
    index: int
    shifts: list[tuple[date, int]]
    capacity_minutes: int
    remaining_minutes: int
    assigned_device: str | None = None
    assigned_config: tuple[str, ...] | None = None


class Scheduler:
    def __init__(
        self,
        tests: list[Test],
        top_priority_tests: set[tuple[str, str]],
        start_date: str | None = None,
        holiday_dates: set[str] = set(),
        daytime_testing_today: bool = False,
        dual_device_weekend_start_enabled: bool = False,
        dual_device_window_start_dates: set[str] | None = None,
        priority_weight: int = 1000,
    ):
        self.top_priority_tests = top_priority_tests
        self.tests = tests
        self.active_dut: dict[str, Test] = {test.test_id: test for test in tests if test.device == DUT}
        self.active_ref: dict[str, Test] = {test.test_id: test for test in tests if test.device == REF}
        self.start_date = date.fromisoformat(start_date) if start_date else date.today()
        self.holiday_dates = holiday_dates
        self.daytime_testing_today = daytime_testing_today
        self.dual_device_weekend_start_enabled = dual_device_weekend_start_enabled
        self.dual_device_window_start_dates = dual_device_window_start_dates or set()
        self.priority_weight = priority_weight
        self.schedule: list[ScheduleEntry] = []
        self.scheduled_bundle_keys: set[tuple[int, str]] = set() 
        self.scheduled_test_ids: set[str] = set()
    

    def compile_schedule(self) -> str | None:
        bundles = build_test_bundles(self.active_dut, self.active_ref, self.top_priority_tests)
        if not bundles:
            return f"Error: No test bundles could be created from the provided tests."
        
        # Print all test bundles for debugging
        print("Compiled Test Bundles:")
        for bundle in bundles:
            print(f"Bundle Index: {bundle.index}, Device: {bundle.device}, Priority: {bundle.priority}, Config: {bundle.config}, Total Minutes: {bundle.total_minutes}, Tests: {bundle.tests}")
                
        # Group by individual TC values (from config tuples)
        all_tcs = set()
        for bundle in bundles:
            all_tcs.update(bundle.config if bundle.config else [None])
        all_bundles_by_tc = self.create_tc_dict(bundles, all_tcs)


        window_index = 0
        cursor_date = self.start_date

        tc_order = self.get_tc_order(self.tests)

        for tc in tc_order:
            tc_bundles = all_bundles_by_tc[tc]
            if tc_bundles is None or len(tc_bundles) == 0:
                continue
            
            window_device = DUT
            pending_dut: list[TestBundle] = [] 
            pending_ref: list[TestBundle] = []

            
            for bundle in tc_bundles:
                if bundle.device == DUT:
                    pending_dut.append(bundle)

                elif bundle.device == REF:
                    pending_ref.append(bundle)

            while pending_dut or pending_ref:
                active_pending = pending_dut if window_device == DUT else pending_ref
                if len(active_pending) == 0:
                    # If no unscheduled bundles for the current device, switch to the other device
                    window_device = REF if window_device == DUT else DUT
                    active_pending = pending_dut if window_device == DUT else pending_ref

                shifts, capacity = get_shift_sequence_with_capacity(cursor_date, self.holiday_dates, self.daytime_testing_today)
                cursor_date_key = cursor_date.isoformat()
                dual_device_window = cursor_date_key in self.dual_device_window_start_dates
                if not dual_device_window and self.dual_device_weekend_start_enabled:
                    dual_device_window = self._is_weekend_start_day(cursor_date)

                selected_bundles = []
                selected_bundles = self._knapsack_select(capacity, active_pending)
                remaining_time = capacity - sum(bundle.total_minutes for bundle in selected_bundles)

                primary_device = window_device

                if dual_device_window and remaining_time > 0:
                    secondary_device = REF if primary_device == DUT else DUT
                    secondary_pending = pending_dut if secondary_device == DUT else pending_ref
                    secondary_selected_bundles = self._knapsack_select(remaining_time, secondary_pending)
                    selected_bundles.extend(secondary_selected_bundles)

                # Create a mirror of the other device for the next window by updating the priorities
                pending_dut, pending_ref = update_mirrored_bundle_priorities(selected_bundles, pending_dut, pending_ref)

                window = ScheduleWindow(
                    index=window_index,
                    shifts=shifts,
                    capacity_minutes=capacity,
                    remaining_minutes=remaining_time,
                    assigned_config=tc,
                    assigned_device=primary_device,
                )

                self._place_bundles_in_window(window, selected_bundles)
                
                # Mark bundles as scheduled and remove from ALL TC buckets
                selected_keys = {(b.index, b.device) for b in selected_bundles}

                # Remove scheduled bundles from ALL TC buckets globally
                for all_tc in all_bundles_by_tc:
                    all_bundles_by_tc[all_tc] = [b for b in all_bundles_by_tc[all_tc] if (b.index, b.device) not in selected_keys]

                # Remove from pending queues 
                pending_dut = [b for b in pending_dut if (b.index, b.device) not in selected_keys]
                pending_ref = [b for b in pending_ref if (b.index, b.device) not in selected_keys]

                window_index += 1
                window_device = REF if window_device == DUT else DUT
                cursor_date = next_window_start_date(shifts)

        return None

    def _is_weekend_start_day(self, current_date: date) -> bool:
        if is_off_day(current_date, self.holiday_dates):
            return False
        return is_off_day(current_date + timedelta(days=1), self.holiday_dates)

    def _select_bundles(
        self,
        device: str,
        capacity: int,
        unscheduled: list[TestBundle],
        tc,
        mutate: bool,
    ) -> tuple[list[TestBundle], list[TestBundle], int]:
        # Consume pending mirrored bundles only while there is room.
        pending = self.pending_dut if device == DUT else self.pending_ref
        mirrored_bundles: list[TestBundle] = []
        still_pending: list[TestBundle] = [] # Bundles that couldn't fit in the remaining capacity
        remaining_capacity = capacity

        for bundle in pending:
            bundle_key = (bundle.index, bundle.device)
            if bundle_key in self.scheduled_bundle_keys:
                continue
            if bundle.device != device or (tc is not None and tc not in bundle.config) or (tc is None and bundle.config):
                still_pending.append(bundle)
                continue
            if bundle.total_minutes <= remaining_capacity:
                mirrored_bundles.append(bundle)
                remaining_capacity -= bundle.total_minutes
            else:
                still_pending.append(bundle)

        if mutate:
            pending[:] = still_pending

        # Run knapsack selection for remaining capacity
        candidates = [
            b
            for b in unscheduled
            if b not in mirrored_bundles
            and b.device == device
            and (b.index, b.device) not in self.scheduled_bundle_keys
        ]
        knapsack_bundles = self._knapsack_select(remaining_capacity, candidates)
        remaining_capacity -= sum(bundle.total_minutes for bundle in knapsack_bundles)

        print(f"Selected {len(mirrored_bundles)} mirrored bundles and {len(knapsack_bundles)} knapsack bundles for device {device} with remaining capacity {remaining_capacity} minutes.")
        return mirrored_bundles, knapsack_bundles, remaining_capacity
        
    def get_schedule(self) -> list[ScheduleEntry]:
        return self.schedule

    def _knapsack_select(
        self,
        capacity: int,
        candidates: list[TestBundle],
    ) -> list[TestBundle]:
        if capacity <= 0 or not candidates:
            return []

        # Weights are total minutes of each bundle
        weights = [bundle.total_minutes for bundle in candidates]
        # Values strongly favor higher priority tiers without hard-forcing them.
        max_priority = max(bundle.priority for bundle in candidates)
        priority_bias = 2.5
        values = [int(self.priority_weight * (priority_bias ** (max_priority - bundle.priority))) for bundle in candidates]

        # Implement dynamic programming knapsack algorithm to select bundles
        n = len(candidates)
        dp = [[0] * (capacity + 1) for _ in range(n + 1)]
        for i in range(1, n + 1):
            for w in range(capacity + 1):
                if weights[i - 1] <= w:
                    dp[i][w] = max(dp[i - 1][w], dp[i - 1][w - weights[i - 1]] + values[i - 1])
                else:
                    dp[i][w] = dp[i - 1][w]

        # Traceback to find selected bundles
        w = capacity
        selected_indices = []
        for i in range(n, 0, -1):
            if w <= 0:
                break
            if dp[i][w] != dp[i - 1][w]:
                selected_indices.append(i - 1)
                w -= weights[i - 1]

        # Return the selected bundles in the order they were added
        return [candidates[i] for i in reversed(selected_indices)]

    def _place_bundles_in_window(
        self,
        window: ScheduleWindow,
        bundles: list[TestBundle],
    ) -> None:
        # Build per-shift remaining capacity
        shift_remaining = [
            get_shift_capacity_for_date(day, self.holiday_dates, self.daytime_testing_today).get(shift_idx, 0)
            for day, shift_idx in window.shifts
        ]
        current_shift_pos = 0
        sequence_in_shift = [0] * len(window.shifts)

        for bundle in bundles:
            bundle_key = (bundle.index, bundle.device)
            if bundle_key in self.scheduled_bundle_keys:
                continue
            self.scheduled_bundle_keys.add(bundle_key)

            for test in bundle.tests:
                schedule_key = f"{test}:{bundle.device}"
                if schedule_key in self.scheduled_test_ids:
                    continue
                self.scheduled_test_ids.add(schedule_key)

                test_minutes = self._get_test_minutes(test, bundle.device)

                while current_shift_pos < len(window.shifts) and shift_remaining[current_shift_pos] <= 0:
                    current_shift_pos += 1

                if current_shift_pos >= len(window.shifts):
                    break

                if sum(shift_remaining[current_shift_pos:]) < test_minutes:
                    break

                start_shift_pos = current_shift_pos
                shift_day, shift_idx = window.shifts[start_shift_pos]
                sequence_in_shift[start_shift_pos] += 1

                remaining_test_minutes = test_minutes
                consume_shift_pos = start_shift_pos
                while remaining_test_minutes > 0 and consume_shift_pos < len(window.shifts):
                    consumed_minutes = min(shift_remaining[consume_shift_pos], remaining_test_minutes)
                    shift_remaining[consume_shift_pos] -= consumed_minutes
                    remaining_test_minutes -= consumed_minutes
                    consume_shift_pos += 1

                window.remaining_minutes -= test_minutes

                while current_shift_pos < len(window.shifts) and shift_remaining[current_shift_pos] <= 0:
                    current_shift_pos += 1

                self.schedule.append(
                    ScheduleEntry(
                        test_id=test,
                        device=bundle.device,
                        scheduled_date=str(shift_day),
                        shift_index=shift_idx,
                        sequence_in_shift=sequence_in_shift[start_shift_pos],
                    )
                )
        

    def _get_test_minutes(self, test_id: str, device: str) -> int:
        if device == DUT and test_id in self.active_dut:
            return self.active_dut[test_id].estimated_minutes
        if device == REF and test_id in self.active_ref:
            return self.active_ref[test_id].estimated_minutes
        return 0

    def create_tc_dict(self, bundles: list[TestBundle], all_tcs: set[str | None]) -> dict[str | None, list[TestBundle]]:
        # Create dict: TC -> bundles supporting that TC
        all_bundles_by_tc: dict[str | None, list[TestBundle]] = {}
        # Sort with None last
        sorted_tcs = sorted([tc for tc in all_tcs if tc is not None]) + ([None] if None in all_tcs else [])
        for tc in sorted_tcs:
            if tc is None:
                all_bundles_by_tc[tc] = [b for b in bundles if not b.config]
            else:
                all_bundles_by_tc[tc] = [b for b in bundles if tc in b.config]
            # Sort by priority, then index
            all_bundles_by_tc[tc].sort(key=lambda b: (b.priority, b.index))
        return all_bundles_by_tc
    
    def get_tc_order(self, all_tests: list[Test]) -> list[str | None]:
        # Implement the logic for getting the test case order
        all_dut = {test.test_id: test for test in all_tests if test.device == DUT}
        all_ref = {test.test_id: test for test in all_tests if test.device == REF}
        bundles = build_test_bundles(all_dut, all_ref, [])

        all_tcs = set()
        for bundle in bundles:
            all_tcs.update(bundle.config if bundle.config else [None])
        all_bundles_by_tc = self.create_tc_dict(bundles, all_tcs)

        # Sort TCs by the number of bundles available for each TC (most to least)
        sorted_tcs = sorted(all_bundles_by_tc.keys(), key=lambda tc: len(all_bundles_by_tc[tc]), reverse=True)
        return sorted_tcs