constraint_engine.py

# constraint_engine.py

from collections import defaultdict
from typing import List, Dict, Any, Tuple
from ortools.sat.python import cp_model
from models import Task, Faculty, Section, Room, SubjectType
import constants as const

TaskVariables = Tuple[cp_model.IntVar, cp_model.IntVar, cp_model.IntervalVar, cp_model.IntVar]

class ConstraintEngine:
    def __init__(
        self,
        model: cp_model.CpModel,
        tasks: List[Task],
        faculties: List[Faculty],
        sections: List[Section],
        rooms: List[Room],
        max_continuous_stretch: int = 3,
        slm_constraints: List[Dict[str, Any]] = None,
        scheduling_rules: List[Dict[str, Any]] = None,
        locked_schedule: Dict[str, Any] = None,
        locked_semesters: List[int] = None
    ):
        self.model = model
        self.tasks = tasks
        self.faculties = faculties
        self.sections = sections
        self.rooms = rooms
        self.max_continuous_stretch = max_continuous_stretch
        self.slm_constraints = slm_constraints or []
        self.scheduling_rules = scheduling_rules or []
        self.locked_schedule = locked_schedule or {}
        self.locked_semesters = set(locked_semesters or [])
        self.room_map = {room.room_id: i for i, room in enumerate(self.rooms)}
        self.task_vars: Dict[str, TaskVariables] = {}

    def apply_all_constraints(self):
        self._create_task_variables()
        self._apply_clash_constraints()
        self._apply_stretch_constraints()
        self._apply_grouping_constraints()
        self._apply_daily_subject_limit()

    def _create_task_variables(self):
        for task in self.tasks:
            is_locked = task.section.semester in self.locked_semesters
            
            if is_locked and task.task_id in self.locked_schedule:
                # Force the task to the exact slot and room from the existing schedule
                locked_info = self.locked_schedule[task.task_id]
                start_slot = locked_info["start_slot"]
                room_id = locked_info["room_id"]
                room_index = self.room_map.get(room_id, 0) # Fallback to 0 if missing

                start_domain = cp_model.Domain.FromValues([start_slot])
                start_var = self.model.NewIntVarFromDomain(start_domain, name=f"{task.task_id}_start")
                
                room_domain = cp_model.Domain.FromValues([room_index])
                room_var = self.model.NewIntVarFromDomain(room_domain, name=f"{task.task_id}_room")
            else:
                allowed_starts = self._get_allowed_start_slots(task)
                start_domain = cp_model.Domain.FromValues(allowed_starts)
                start_var = self.model.NewIntVarFromDomain(start_domain, name=f"{task.task_id}_start")

                allowed_room_indices = self._get_allowed_rooms(task)
                room_domain = cp_model.Domain.FromValues(allowed_room_indices)
                room_var = self.model.NewIntVarFromDomain(room_domain, name=f"{task.task_id}_room")

            end_var = self.model.NewIntVar(0, const.TOTAL_TEACHING_SLOTS_PER_WEEK, name=f"{task.task_id}_end")
            interval_var = self.model.NewIntervalVar(start_var, task.duration, end_var, name=f"{task.task_id}_interval")

            self.task_vars[task.task_id] = (start_var, end_var, interval_var, room_var)

    def _apply_clash_constraints(self):
        # 1. Faculty Clash Prevention (MODIFIED FOR COMBINED ELECTIVES)
        intervals_by_faculty = defaultdict(list)
        processed_groups_per_faculty = defaultdict(set)

        faculty_ids_set = {f.id for f in self.faculties}
        for task in self.tasks:
            parts = task.faculty.id.split('_')
            fids = parts if len(parts) > 1 and all(p in faculty_ids_set for p in parts) else [task.faculty.id]
            gid = task.elective_group_id
            interval = self.task_vars[task.task_id][2]

            for fid in fids:
                if gid:
                    # If this faculty is teaching a group, only add the interval ONCE for that group
                    if gid not in processed_groups_per_faculty[fid]:
                        intervals_by_faculty[fid].append(interval)
                        processed_groups_per_faculty[fid].add(gid)
                else:
                    intervals_by_faculty[fid].append(interval)
        
        for faculty_id in intervals_by_faculty:
            if faculty_id == "DUMMY_STAFF":
                continue
            self.model.AddNoOverlap(intervals_by_faculty[faculty_id])

        # 2. Section Clash Prevention (MODIFIED FOR ELECTIVES AND BATCHES)
        intervals_by_section = defaultdict(list)
        processed_groups_per_section = defaultdict(set)
        
        # We also need to keep track of parent sections dynamically in case they are missing from self.sections
        parent_sections_dynamic = set()

        for task in self.tasks:
            sid = task.section.section_id
            gid = task.elective_group_id
            interval = self.task_vars[task.task_id][2]

            if gid:
                if gid not in processed_groups_per_section[sid]:
                    intervals_by_section[sid].append(interval)
                    processed_groups_per_section[sid].add(gid)
            else:
                intervals_by_section[sid].append(interval)
                
            # Add to parent_sections_dynamic if it's a parent section
            if '-' not in sid:
                parent_sections_dynamic.add(sid)
            else:
                parent_sections_dynamic.add(sid.split('-')[0])

        for section_id in intervals_by_section:
            self.model.AddNoOverlap(intervals_by_section[section_id])

        # Enforce NoOverlap between Parent and ALL its Batches (including between batches)
        for parent_sid in parent_sections_dynamic:
            parent_intervals = list(intervals_by_section.get(parent_sid, []))

            # Parent can't overlap with any batch task (grouped or not)
            for batch_suffix in ['-B1', '-B2', '-B3']:
                batch_id = f"{parent_sid}{batch_suffix}"
                if batch_id in intervals_by_section:
                    self.model.AddNoOverlap(parent_intervals + intervals_by_section[batch_id])

            # Non-grouped batch tasks across different batches can't overlap.
            # (Grouped tasks are allowed to overlap — the grouping constraint
            #  forces them to the same start time intentionally.)
            non_grouped_cross_batch = []
            for task in self.tasks:
                sid = task.section.section_id
                if '-' in sid and sid.split('-')[0] == parent_sid and not task.elective_group_id:
                    non_grouped_cross_batch.append(self.task_vars[task.task_id][2])

            if len(non_grouped_cross_batch) > 1:
                self.model.AddNoOverlap(non_grouped_cross_batch)

        # 3. Room Clash Prevention
        for i, room in enumerate(self.rooms):
            optional_intervals = []
            for task in self.tasks:
                if i in self._get_allowed_rooms(task):
                    interval, room_var = self.task_vars[task.task_id][2], self.task_vars[task.task_id][3]
                    is_in_room = self.model.NewBoolVar(f"{task.task_id}_in_room_{room.room_id}")
                    self.model.Add(room_var == i).OnlyEnforceIf(is_in_room)
                    self.model.Add(room_var != i).OnlyEnforceIf(is_in_room.Not())
                    opt_interval = self.model.NewOptionalIntervalVar(
                        interval.StartExpr(), interval.SizeExpr(), interval.EndExpr(), is_in_room, name=f"opt_{task.task_id}_{room.room_id}"
                    )
                    optional_intervals.append(opt_interval)
            if optional_intervals:
                self.model.AddNoOverlap(optional_intervals)

    def _apply_stretch_constraints(self):
        tasks_by_entity = defaultdict(list)
        for task in self.tasks:
            if task.faculty.id != "DUMMY_STAFF":
                for fid in task.faculty.id.split('_'):  # Handle composite co-teaching faculties
                    tasks_by_entity[f"fac_{fid}"].append(task)
            tasks_by_entity[f"sec_{task.section.section_id}"].append(task)
        for name, tasks in tasks_by_entity.items():
            self._add_stretch_constraint_for_entity(tasks, name)

    def _apply_grouping_constraints(self):
        tasks_by_group = defaultdict(list)
        for task in self.tasks:
            if task.elective_group_id:
                tasks_by_group[task.elective_group_id].append(task)
        for group_tasks in tasks_by_group.values():
            if len(group_tasks) > 1:
                first_start = self.task_vars[group_tasks[0].task_id][0]
                for i in range(1, len(group_tasks)):
                    self.model.Add(self.task_vars[group_tasks[i].task_id][0] == first_start)

    def _apply_daily_subject_limit(self):
        tasks_by_sec_sub = defaultdict(list)
        for task in self.tasks:
            if task.subject.subject_type == SubjectType.THEORY:
                tasks_by_sec_sub[(task.section.section_id, task.subject.subject_code)].append(task)
        for tasks in tasks_by_sec_sub.values():
            for day in range(const.NUM_WORKING_DAYS):
                literals = []
                for task in tasks:
                    start_var = self.task_vars[task.task_id][0]
                    is_on_day = self.model.NewBoolVar(f"{task.task_id}_on_day_{day}")
                    day_idx = self.model.NewIntVar(0, const.NUM_WORKING_DAYS-1, f"day_{task.task_id}_{day}")
                    self.model.AddDivisionEquality(day_idx, start_var, const.NUM_TEACHING_SLOTS_PER_DAY)
                    self.model.Add(day_idx == day).OnlyEnforceIf(is_on_day)
                    self.model.Add(day_idx != day).OnlyEnforceIf(is_on_day.Not())
                    literals.append(is_on_day)
                if literals: self.model.Add(sum(literals) <= 1)

    def _get_allowed_start_slots(self, task: Task) -> List[int]:
        full_range = set(range(const.TOTAL_TEACHING_SLOTS_PER_WEEK))
        if task.duration > 1:
            allowed_lab_slots = set()
            for day in range(const.NUM_WORKING_DAYS):
                day_offset = day * const.NUM_TEACHING_SLOTS_PER_DAY
                for pos in const.ALLOWED_LAB_START_INDICES:
                    allowed_lab_slots.add(day_offset + pos)
            full_range.intersection_update(allowed_lab_slots)

        # ── Configurable scheduling rules (data-driven) ─────────────────
        DAY_MAP = {'MON':0,'TUE':1,'WED':2,'THU':3,'FRI':4,'SAT':5}
        for rule in self.scheduling_rules:
            rtype = rule.get('rule_type', '').upper()
            
            if rtype == 'FACULTY_UNAVAILABLE':
                rule_fid = rule.get('faculty_id')
                if rule_fid and rule_fid != task.faculty.id:
                    continue
                
                days = [DAY_MAP[d.upper()] for d in (rule.get('days') or []) if d.upper() in DAY_MAP]
                period_idx = rule.get('period_index')
                
                for day in days:
                    off = day * const.NUM_TEACHING_SLOTS_PER_DAY
                    if period_idx is not None:
                        for o in range(task.duration):
                            full_range.discard(off + period_idx - o)
                    else:
                        for p in range(const.NUM_TEACHING_SLOTS_PER_DAY):
                            for o in range(task.duration):
                                full_range.discard(off + p - o)
                continue

            rule_subjects = [s.lower() for s in (rule.get('subject_codes') or [])]
            rule_types = [t.upper() for t in (rule.get('subject_types') or [])]
            # Check if this rule applies to the current task
            matches_subject = task.subject.subject_code.lower() in rule_subjects if rule_subjects else False
            matches_type = task.subject.subject_type.name in rule_types if rule_types else False
            if not matches_subject and not matches_type:
                continue

            if rtype == 'FIXED_PERIOD':
                # Force tasks to a specific period index
                period_idx = rule.get('period_index', 0)
                full_range = {s for s in full_range
                              if (s % const.NUM_TEACHING_SLOTS_PER_DAY) == period_idx}

            elif rtype == 'BEFORE_TIME':
                # Tasks must start at or before a max period index
                max_p = rule.get('max_period_index', 4)
                full_range = {s for s in full_range
                              if (s % const.NUM_TEACHING_SLOTS_PER_DAY) <= max_p}

            elif rtype == 'FIXED_DAYS':
                # Tasks must only be on specific days
                allowed_days = {DAY_MAP[d.upper()] for d in (rule.get('days') or []) if d.upper() in DAY_MAP}
                if allowed_days:
                    full_range = {s for s in full_range
                                  if (s // const.NUM_TEACHING_SLOTS_PER_DAY) in allowed_days}

        if self.slm_constraints:
            DAY_INDEX = {'MON':0,'TUE':1,'WED':2,'THU':3,'FRI':4,'SAT':5}
            MORNING_PERIODS = list(range(0, const.NUM_TEACHING_SLOTS_PER_DAY // 2))
            AFTERNOON_PERIODS = list(range(const.NUM_TEACHING_SLOTS_PER_DAY // 2, const.NUM_TEACHING_SLOTS_PER_DAY))
            for c in self.slm_constraints:
                ctype = c.get('type', '').upper()
                if ctype == 'FACULTY_UNAVAILABLE' and c.get('faculty_id') == task.faculty.id:
                    days = [DAY_INDEX[d] for d in (c.get('days') or []) if d in DAY_INDEX]
                    for day in days:
                        off = day * const.NUM_TEACHING_SLOTS_PER_DAY
                        for p in range(const.NUM_TEACHING_SLOTS_PER_DAY):
                            for o in range(task.duration):
                                full_range.discard(off + p - o)
                elif ctype == 'SUBJECT_PREFERRED_TIME':
                    if c.get('subject_code') == task.subject.subject_code:
                        period = str(c.get('period','')).upper()
                        target = MORNING_PERIODS if 'MORNING' in period else AFTERNOON_PERIODS
                        if target:
                            full_range = {s for s in full_range if (s % const.NUM_TEACHING_SLOTS_PER_DAY) in target}
                elif ctype == 'WORKING_DAYS':
                    allowed_days = {DAY_INDEX[d] for d in (c.get('days') or []) if d in DAY_INDEX}
                    if allowed_days:
                        full_range = {s for s in full_range if (s // const.NUM_TEACHING_SLOTS_PER_DAY) in allowed_days}
                elif ctype == 'SECTION_FREE_SLOT' and c.get('section_id') == task.section.section_id:
                    slot_index = c.get('slot')
                    days = [DAY_INDEX[d] for d in (c.get('days') or []) if d in DAY_INDEX]
                    if not days:
                        days = list(range(const.NUM_WORKING_DAYS))
                    if slot_index is not None:
                        for day in days:
                            off = day * const.NUM_TEACHING_SLOTS_PER_DAY
                            for o in range(task.duration):
                                full_range.discard(off + slot_index - 1 - o)
                                
        return sorted(list(full_range))

    def _get_allowed_rooms(self, task: Task) -> List[int]:
        if task.subject.subject_code in ["LIB_HR", "STU_HR", "FAC_HR", "STDY_HR"]:
            return [-1]
            
        allowed = []
        for i, room in enumerate(self.rooms):
            type_match = (task.subject.subject_type == SubjectType.LAB) == room.is_lab
            cap_match = room.capacity >= task.section.student_strength
            if type_match and cap_match: allowed.append(i)
        if not allowed:
            req_type = 'Computer Lab' if task.subject.subject_type == SubjectType.LAB else 'Standard Classroom'
            raise ValueError(f"Shortage of suitable class rooms! Could not find any {req_type} with a minimum capacity of {task.section.student_strength} students for the subject '{task.subject.name}' (Section {task.section.section_id}). Please add a larger room to the template.")
        return allowed

    def _add_stretch_constraint_for_entity(self, tasks: List[Task], entity_name: str):
        if not tasks: return
        for day in range(const.NUM_WORKING_DAYS):
            for i in range(const.NUM_TEACHING_SLOTS_PER_DAY - self.max_continuous_stretch):
                window_start = day * const.NUM_TEACHING_SLOTS_PER_DAY + i
                window_slots = range(window_start, window_start + self.max_continuous_stretch + 1)
                literals = []
                for slot in window_slots:
                    is_active = self.model.NewBoolVar(f"{entity_name}_act_{slot}")
                    at_slot = []
                    for t in tasks:
                        start, dur = self.task_vars[t.task_id][0], t.duration
                        cov = self.model.NewBoolVar(f"{t.task_id}_cov_{slot}")
                        self.model.Add(start <= slot).OnlyEnforceIf(cov)
                        self.model.Add(start + dur > slot).OnlyEnforceIf(cov)
                        at_slot.append(cov)
                    self.model.AddBoolOr(at_slot).OnlyEnforceIf(is_active)
                    for lit in at_slot: self.model.AddImplication(is_active.Not(), lit.Not())
                    literals.append(is_active)
                self.model.Add(sum(literals) <= self.max_continuous_stretch)