import itertools
import json
import os
import time
from collections import defaultdict
from collections.abc import Callable, Generator
from dataclasses import dataclass
from functools import cache
from typing import cast

import z3
from bidict import frozenbidict
from pydantic import BaseModel

from .config import (
    CombinedConfig,
    FacultyConfig,
    OptimizerFlags,
)
from .logging import logger
from .models import (
    Course,
    CourseInstance,
    Day,
    TimeInstance,
    TimePoint,
    TimeSlot,
)
from .time_slot_generator import TimeSlotGenerator


def load_config_from_file[T: BaseModel](
    config_cls: type[T],
    filename: str | os.PathLike[str],
) -> T:
    """
    Load scheduler configuration from a JSON file.

    **Usage:**
    ```python
    from scheduler import CombinedConfig, load_config_from_file
    cfg = load_config_from_file(CombinedConfig, "config.json")
    ```

    **Args:**
    - config_cls: The class of the configuration to load
    - filename: Path to the file (str or pathlib.Path)

    **Returns:**
    The loaded configuration

    **Example:**
    >>> load_config_from_file(CombinedConfig, "config.json")
    >>> load_config_from_file(CombinedConfig, Path("config.json"))
    """
    with open(filename, encoding="utf-8") as f:
        data = json.load(f)
    return config_cls(**data)


def get_faculty_availability(
    faculty_config: FacultyConfig,
) -> list[TimeInstance]:
    """
    Calculate the availability of a faculty as a list of `TimeInstance` objects.

    **Usage:**
    ```python
    slots = get_faculty_availability(faculty_config)
    ```

    **Args:**
    - faculty_config: The configuration of the faculty

    **Returns:**
    The availability of the faculty as a list of `TimeInstance` objects
    """
    days: list[Day] = [Day.MON, Day.TUE, Day.WED, Day.THU, Day.FRI]
    result: list[TimeInstance] = []
    for day in days:
        day_name = day.name
        times = faculty_config.times.get(day_name, [])
        for time_range in times:
            # Parse TimeRange object
            start_str = time_range.start
            end_str = time_range.end
            start_hour, start_minute = map(int, start_str.split(":"))
            end_hour, end_minute = map(int, end_str.split(":"))

            start_time: TimePoint = TimePoint.make_from(start_hour, start_minute)
            end_time: TimePoint = TimePoint.make_from(end_hour, end_minute)
            result.append(
                TimeInstance(
                    day=day,
                    start=start_time,
                    duration=end_time - start_time,
                )
            )
    return result


@dataclass
class _FunctionConstraints:
    """
    Structured data for function constraints and their references.

    **Usage:**
    ```python
    # Built by Scheduler._build_function_constraints
    ```
    """

    constraints: list[z3.BoolRef]
    overlaps: z3.FuncDeclRef
    lab_overlaps: z3.FuncDeclRef
    lecture_next_to: z3.FuncDeclRef
    faculty_available: z3.FuncDeclRef
    lab_next_to: z3.FuncDeclRef


@dataclass
class _Z3SortsAndConstants:
    """
    Structured data for Z3 sorts and their corresponding constant mappings.

    **Usage:**
    ```python
    # Built by Scheduler._create_z3_enumsorts
    ```
    """

    time_slot_sort: z3.SortRef
    time_slot_constants: frozenbidict[TimeSlot, z3.ExprRef]
    faculty_sort: z3.SortRef
    faculty_constants: frozenbidict[str, z3.ExprRef]
    room_sort: z3.SortRef
    room_constants: frozenbidict[str, z3.ExprRef]
    lab_sort: z3.SortRef
    lab_constants: frozenbidict[str, z3.ExprRef]


class Scheduler:
    """
    Scheduler class for generating schedules.

    **Usage:**
    ```python
    sched = Scheduler(full_config)
    next(sched.get_models())
    ```
    """

    def _initialize_faculty_data(self, config) -> None:
        """
        Initialize faculty-related data structures and preferences.

        **Usage:**
        ```python
        self._initialize_faculty_data(config)
        ```
        """
        for faculty_data in config.faculty:
            faculty_name = faculty_data.name
            self._faculty.add(faculty_name)
            self._faculty_maximum_credits[faculty_name] = faculty_data.maximum_credits
            self._faculty_maximum_days[faculty_name] = faculty_data.maximum_days
            self._faculty_minimum_credits[faculty_name] = faculty_data.minimum_credits
            self._faculty_unique_course_limits[faculty_name] = faculty_data.unique_course_limit
            self._faculty_course_preferences[faculty_name] = faculty_data.course_preferences
            self._faculty_room_preferences[faculty_name] = faculty_data.room_preferences
            self._faculty_lab_preferences[faculty_name] = faculty_data.lab_preferences
            self._faculty_mandatory_days[faculty_name] = {
                day if isinstance(day, Day) else Day[day] for day in faculty_data.mandatory_days
            }
            self._faculty_availability[faculty_name] = get_faculty_availability(faculty_data)

    def _initialize_courses(self, config) -> tuple[list[Course], set[int]]:
        """
        Initialize courses and return them along with required credits.

        **Usage:**
        ```python
        courses, credits = self._initialize_courses(config)
        ```
        """
        courses: list[Course] = []
        required_credits = set()
        course_counts: dict[str, int] = defaultdict(int)

        for c in config.courses:
            course_counts[c.course_id] += 1
            required_credits.add(c.credits)
            course_faculty = c.faculty
            if not course_faculty:
                for faculty_data in config.faculty:
                    if c.course_id in faculty_data.course_preferences:
                        course_faculty.append(faculty_data.name)

            course = Course(
                credits=c.credits,
                course_id=c.course_id,
                section=course_counts[c.course_id],
                labs=c.lab,
                rooms=c.room,
                conflicts=c.conflicts,
                faculties=course_faculty,
            )
            courses.append(course)

        return courses, required_credits

    def _initialize_time_slots(self, time_slot_config, required_credits: set[int]) -> None:
        """
        Initialize time slots and create ranges for different credit levels.

        **Usage:**
        ```python
        self._initialize_time_slots(time_slot_config, required_credits)
        ```
        """
        self._time_slot_generator = TimeSlotGenerator(time_slot_config)
        self._ranges: dict[int, tuple[int, int]] = {}
        self._slots: list[TimeSlot] = []

        for creds in sorted(required_credits):
            low = len(self._slots)
            self._slots.extend(self._time_slot_generator.time_slots(creds))
            self._ranges[creds] = (low, len(self._slots) - 1)

    def _create_z3_enumsorts(self) -> _Z3SortsAndConstants:
        """
        Create Z3 EnumSorts for time slots, faculty, rooms, and labs.

        **Usage:**
        ```python
        z3_data = self._create_z3_enumsorts()
        ```
        """

        def sanitize(name):
            return name.replace(" ", "_")

        # Create TimeSlot EnumSort
        time_slot_names = [sanitize(str(slot)) for slot in self._slots]
        time_slot_sort, time_slot_constants = z3.EnumSort("TimeSlot", time_slot_names, ctx=self._ctx)
        time_slot_constants_dict = frozenbidict(
            {time_slot: time_slot_constants[i] for i, time_slot in enumerate(self._slots)}
        )

        # Create Faculty EnumSort
        faculty_names = [sanitize(faculty) for faculty in self._faculty]
        faculty_sort, faculty_constants = z3.EnumSort("Faculty", faculty_names, ctx=self._ctx)
        faculty_constants_dict = frozenbidict(
            {faculty: faculty_constants[i] for i, faculty in enumerate(self._faculty)},
        )

        # Create Room EnumSort
        room_names = [sanitize(room) for room in self._rooms]
        room_sort, room_constants = z3.EnumSort("Room", room_names, ctx=self._ctx)
        room_constants_dict = frozenbidict(
            {room: room_constants[i] for i, room in enumerate(self._rooms)},
        )

        # Create Lab EnumSort
        lab_names = [sanitize(lab) for lab in self._labs]
        lab_sort, lab_constants = z3.EnumSort("Lab", lab_names, ctx=self._ctx)
        lab_constants_dict = frozenbidict(
            {lab: lab_constants[i] for i, lab in enumerate(self._labs)},
        )

        return _Z3SortsAndConstants(
            time_slot_sort=time_slot_sort,
            time_slot_constants=time_slot_constants_dict,
            faculty_sort=faculty_sort,
            faculty_constants=faculty_constants_dict,
            room_sort=room_sort,
            room_constants=room_constants_dict,
            lab_sort=lab_sort,
            lab_constants=lab_constants_dict,
        )

    def _create_course_variables(self, z3_data: _Z3SortsAndConstants) -> None:
        """
        Create Z3 variables for each course.

        **Usage:**
        ```python
        self._create_course_variables(z3_data)
        ```
        """
        for course in self._courses:
            course.time = z3.Const(f"{str(course)}_time", z3_data.time_slot_sort)
            course.faculty = z3.Const(f"{str(course)}_faculty", z3_data.faculty_sort)
            course.room = z3.Const(f"{str(course)}_room", z3_data.room_sort)
            course.lab = z3.Const(f"{str(course)}_lab", z3_data.lab_sort)

    def __init__(self, full_config: CombinedConfig):
        """
        Initializes the scheduler with all the necessary constraints and variables.

        **Usage:**
        ```python
        Scheduler(full_config)
        ```

        **Args:**
        - full_config: `CombinedConfig` object containing all the configuration
                       settings including the scheduler config, time slot config,
                       limit, and optimizer flags

        **Raises:**
        - ValueError: If the optimizer flags are invalid
        """
        # Extract configuration
        config = full_config.config
        time_slot_config = full_config.time_slot_config
        self._optimizer_flags = full_config.optimizer_flags
        self._limit = full_config.limit

        # Initialize Z3 context
        self._ctx = z3.Context()

        # Initialize data structures
        self._faculty: set[str] = set()
        self._faculty_maximum_credits: dict[str, int] = {}
        self._faculty_maximum_days: dict[str, int] = {}
        self._faculty_minimum_credits: dict[str, int] = {}
        self._faculty_unique_course_limits: dict[str, int] = {}
        self._faculty_course_preferences: dict[str, dict[str, int]] = {}
        self._faculty_room_preferences: dict[str, dict[str, int]] = {}
        self._faculty_lab_preferences: dict[str, dict[str, int]] = {}
        self._faculty_mandatory_days: dict[str, set[Day]] = {}
        self._faculty_availability: dict[str, list[TimeInstance]] = {}
        self._initialize_faculty_data(config)

        # Initialize courses and time slots
        self._rooms = set(config.rooms)
        self._labs = set(config.labs)
        self._courses, required_credits = self._initialize_courses(config)
        self._initialize_time_slots(time_slot_config, required_credits)

        # Create Z3 structures
        z3_data = self._create_z3_enumsorts()
        self._create_course_variables(z3_data)

        # Build function constraints and get the function references
        function_data = self._build_function_constraints(z3_data)

        # Build faculty constraints
        faculty_constraints = self._build_faculty_constraints(z3_data)

        # Build course constraints
        course_constraints = self._build_course_constraints(
            function_data.overlaps,
            function_data.faculty_available,
            z3_data,
        )

        # Build resource constraints
        resource_constraints = self._build_resource_constraints(
            function_data.overlaps,
            function_data.lab_overlaps,
            function_data.lecture_next_to,
            function_data.lab_next_to,
            z3_data,
        )

        # Aggregate all constraints
        self._constraints = self._aggregate_constraints(
            function_data.constraints, faculty_constraints, course_constraints, resource_constraints
        )

        self._function_data = function_data
        self._z3_data = z3_data

    @cache
    def _simplify(self, x: z3.ExprRef) -> z3.BoolRef:
        """
        Cached simplification to avoid redundant computation

        **Usage:**
        ```python
        self._simplify(z3_expr)
        ```
        """
        return cast(z3.BoolRef, z3.simplify(x, cache_all=True, local_ctx=True))

    @cache
    def _cached_slot_relationship(self, fn_name: str, slot_i: TimeSlot, slot_j: TimeSlot) -> bool:
        if fn_name == "overlaps":
            return slot_i.overlaps(slot_j)
        elif fn_name == "lab_overlaps":
            return slot_i.lab_overlaps(slot_j)
        elif fn_name == "lecture_next_to":
            return slot_i.lecture_next_to(slot_j)
        elif fn_name == "lab_next_to":
            return slot_i.lab_next_to(slot_j)
        else:
            raise ValueError(f"Unknown relationship function: {fn_name}")

    def _z3ify_time_constraint(
        self, z3_data: _Z3SortsAndConstants, name: str, *, ctx: z3.Context | None = None
    ) -> tuple[z3.FuncDeclRef, list[z3.BoolRef]]:
        z3fn = z3.Function(
            name,
            z3_data.time_slot_sort,
            z3_data.time_slot_sort,
            z3.BoolSort(ctx=ctx),
        )

        true: list[tuple[z3.BoolRef, z3.BoolRef]] = []
        false: list[tuple[z3.BoolRef, z3.BoolRef]] = []

        for slot_i, slot_j in itertools.combinations_with_replacement(self._slots, 2):
            c_i = cast(z3.BoolRef, z3_data.time_slot_constants[slot_i])
            c_j = cast(z3.BoolRef, z3_data.time_slot_constants[slot_j])
            if self._cached_slot_relationship(name, slot_i, slot_j):
                true.append((c_i, c_j))
                true.append((c_j, c_i))
            else:
                false.append((c_i, c_j))
                false.append((c_j, c_i))

        constraints: list[z3.BoolRef] = []
        if true:
            constraints.append(cast(z3.BoolRef, z3.And([z3fn(ts_i, ts_j) for ts_i, ts_j in true])))
        if false:
            constraints.append(
                cast(
                    z3.BoolRef,
                    z3.And([z3.Not(z3fn(ts_i, ts_j)) for ts_i, ts_j in false]),
                )
            )

        return z3fn, constraints

    def _z3ify_time_slot_fn(
        self,
        z3_data: _Z3SortsAndConstants,
        name: str,
        fn: Callable[[TimeSlot], bool],
        *,
        ctx: z3.Context | None = None,
    ) -> tuple[z3.FuncDeclRef, list[z3.BoolRef]]:
        z3fn = z3.Function(name, z3_data.time_slot_sort, z3.BoolSort(ctx=ctx))

        true: list[z3.BoolRef] = []
        false: list[z3.BoolRef] = []
        for slot in self._slots:
            c = cast(z3.BoolRef, z3_data.time_slot_constants[slot])
            if fn(slot):
                true.append(c)
            else:
                false.append(c)
        constraints: list[z3.BoolRef] = []
        if true:
            constraints.append(cast(z3.BoolRef, z3.And([z3fn(ts) for ts in true])))
        if false:
            constraints.append(cast(z3.BoolRef, z3.And([z3.Not(z3fn(ts)) for ts in false])))
        return z3fn, constraints

    def _z3ify_faculty_time_constraint(
        self, z3_data: _Z3SortsAndConstants, name: str, *, ctx: z3.Context | None = None
    ) -> tuple[z3.FuncDeclRef, list[z3.BoolRef]]:
        z3fn = z3.Function(
            name,
            z3_data.faculty_sort,
            z3_data.time_slot_sort,
            z3.BoolSort(ctx=ctx),
        )

        constraints: list[z3.BoolRef] = []
        for faculty in self._faculty:
            true: list[tuple[z3.BoolRef, z3.BoolRef]] = []
            false: list[tuple[z3.BoolRef, z3.BoolRef]] = []
            faculty_times = self._faculty_availability[faculty]
            faculty_constant = cast(z3.BoolRef, z3_data.faculty_constants[faculty])
            for slot in self._slots:
                slot_constant = cast(z3.BoolRef, z3_data.time_slot_constants[slot])
                if slot.in_time_ranges(faculty_times):
                    true.append((faculty_constant, slot_constant))
                else:
                    false.append((faculty_constant, slot_constant))
            if true:
                constraints.append(
                    cast(
                        z3.BoolRef,
                        z3.And([z3fn(faculty, slot) for faculty, slot in true]),
                    )
                )
            if false:
                constraints.append(
                    cast(
                        z3.BoolRef,
                        z3.And([z3.Not(z3fn(faculty, slot)) for faculty, slot in false]),
                    )
                )

        return z3fn, constraints

    def _build_function_constraints(self, z3_data: _Z3SortsAndConstants) -> _FunctionConstraints:
        """
        Create Z3 function definitions and their constraints.

        **Usage:**
        ```python
        self._build_function_constraints(z3_data)
        ```

        **Args:**
        - z3_data: `_Z3SortsAndConstants` object containing the Z3 sorts and constants

        **Returns:**
        - `_FunctionConstraints` object containing the Z3 function definitions and their constraints
        """
        # abstract function constraints
        overlaps, overlaps_C = self._z3ify_time_constraint(z3_data, "overlaps", ctx=self._ctx)
        lab_overlaps, lab_overlaps_C = self._z3ify_time_constraint(z3_data, "lab_overlaps", ctx=self._ctx)
        lecture_next_to, lecture_next_to_C = self._z3ify_time_constraint(z3_data, "lecture_next_to", ctx=self._ctx)
        faculty_available, faculty_available_C = self._z3ify_faculty_time_constraint(
            z3_data, "faculty_available", ctx=self._ctx
        )
        lab_next_to, lab_next_to_C = self._z3ify_time_constraint(z3_data, "lab_next_to", ctx=self._ctx)

        function_constraints: list[z3.BoolRef] = []
        function_constraints.extend(overlaps_C)
        function_constraints.extend(lab_overlaps_C)
        function_constraints.extend(lecture_next_to_C)
        function_constraints.extend(lab_next_to_C)
        function_constraints.extend(faculty_available_C)

        return _FunctionConstraints(
            constraints=function_constraints,
            overlaps=overlaps,
            lab_overlaps=lab_overlaps,
            lecture_next_to=lecture_next_to,
            faculty_available=faculty_available,
            lab_next_to=lab_next_to,
        )

    def _build_faculty_constraints(self, z3_data: _Z3SortsAndConstants) -> list[z3.BoolRef]:
        """
        Create constraints for faculty credit limits and unique course limits.

        **Usage:**
        ```python
        self._build_faculty_constraints(z3_data)
        ```

        **Args:**
        - z3_data: `_Z3SortsAndConstants` object containing the Z3 sorts and constants

        **Returns:**
        - `list[z3.BoolRef]` containing the faculty constraints
        """
        # Pre-compute course groupings to reduce repeated calculations
        faculty_course_map: defaultdict[str, list[Course]] = defaultdict(list)
        for c in self._courses:
            for faculty in c.faculties:
                faculty_course_map[faculty].append(c)

        # Pre-compute time slot constants per day for reuse
        day_slot_map: defaultdict[Day, set[z3.ExprRef]] = defaultdict(set)
        for slot in self._slots:
            slot_constant = z3_data.time_slot_constants[slot]
            for time_instance in slot.times:
                day_slot_map[time_instance.day].add(slot_constant)
        day_to_slot_constants: dict[Day, tuple[z3.ExprRef, ...]] = {
            day: tuple(slot_constants) for day, slot_constants in day_slot_map.items()
        }

        # Add faculty credit and unique course limits - batch generation
        faculty_constraints: list[z3.BoolRef] = []
        for faculty in self._faculty:
            faculty_courses = faculty_course_map.get(faculty, [])
            faculty_constant = z3_data.faculty_constants[faculty]
            max_days = self._faculty_maximum_days[faculty]
            mandatory_days = self._faculty_mandatory_days[faculty]
            if faculty_courses:
                min_credits = self._faculty_minimum_credits[faculty]
                max_credits = self._faculty_maximum_credits[faculty]
                credit_sum = z3.Sum([z3.If(c.faculty == faculty_constant, c.credits, 0) for c in faculty_courses])
                # ensure that each faculty is assigned between min and max credits
                faculty_constraints.append(
                    cast(
                        z3.BoolRef,
                        z3.And(credit_sum >= min_credits, credit_sum <= max_credits),
                    )
                )

                # Unique course limit constraint - only generate if needed
                unique_limit = self._faculty_unique_course_limits[faculty]

                # Group courses by their unique identifier (subject + number)
                unique_courses: defaultdict[str, list[Course]] = defaultdict(list)
                for c in faculty_courses:
                    unique_courses[c.course_id].append(c)

                # Only create constraint if there are multiple unique courses
                if len(unique_courses) > unique_limit:
                    teaches_course: list[z3.BoolRef] = []
                    for course_group in unique_courses.values():
                        teaches_course.append(
                            cast(
                                z3.BoolRef,
                                z3.Or([c.faculty == faculty_constant for c in course_group]),
                            )
                        )
                    # ensure that each faculty is assigned <= unique course limit
                    limit = self._simplify(z3.Sum([z3.If(tc, 1, 0) for tc in teaches_course]) <= unique_limit)
                    faculty_constraints.append(limit)

            # Track whether the faculty teaches on a given day
            day_indicator_map: dict[Day, z3.BoolRef] = {}
            for day in Day:
                slot_constants = day_to_slot_constants.get(day, ())
                course_day_assignments: list[z3.BoolRef] = []
                if slot_constants and faculty_courses:
                    for course in faculty_courses:
                        slot_matches = [course.time == slot_const for slot_const in slot_constants]
                        if slot_matches:
                            course_day_assignments.append(
                                self._simplify(
                                    z3.And(
                                        course.faculty == faculty_constant,
                                        z3.Or(slot_matches),
                                    )
                                )
                            )
                if course_day_assignments:
                    day_indicator_map[day] = self._simplify(z3.Or(course_day_assignments))
                else:
                    day_indicator_map[day] = z3.BoolVal(False, ctx=self._ctx)

            # Maximum-day constraint
            day_sum_terms = [z3.If(indicator, 1, 0) for indicator in day_indicator_map.values()]
            day_sum = z3.Sum(day_sum_terms) if day_sum_terms else z3.IntVal(0, ctx=self._ctx)
            faculty_constraints.append(self._simplify(day_sum <= max_days))

            # Mandatory-day constraints
            for mandatory_day in mandatory_days:
                indicator = day_indicator_map.get(mandatory_day, z3.BoolVal(False, ctx=self._ctx))
                faculty_constraints.append(self._simplify(indicator))

        return faculty_constraints

    def _build_course_constraints(
        self,
        overlaps: z3.FuncDeclRef,
        faculty_available: z3.FuncDeclRef,
        z3_data: _Z3SortsAndConstants,
    ) -> list[z3.BoolRef]:
        """
        Create individual course constraints.

        **Usage:**
        ```python
        self._build_course_constraints(overlaps, faculty_available, z3_data)
        ```

        **Args:**
        - overlaps: `z3.FuncDeclRef` function for checking time overlaps
        - faculty_available: `z3.FuncDeclRef` function for checking faculty availability
        - z3_data: `_Z3SortsAndConstants` object containing the Z3 sorts and constants

        **Returns:**
        - `list[z3.BoolRef]` containing the course constraints
        """
        # Course constraints with optimized conflict checking - batch generation
        course_constraints: list[z3.BoolRef] = []
        for c in self._courses:
            # conflict constraints
            conflict_constraints: list[z3.BoolRef] = []
            for d in self._courses:
                if d != c and d.course_id in c.conflicts:
                    conflict_constraints.append(cast(z3.BoolRef, z3.Not(overlaps(c.time, d.time))))

            # faculty availability constraint
            course_constraint_list: list[z3.BoolRef] = [
                cast(z3.BoolRef, faculty_available(c.faculty, c.time)),
            ]

            # Get valid time slots for this credit level
            start, stop = self._ranges[c.credits]
            valid_time_slots = {i for i, _ in enumerate(self._slots) if start <= i <= stop}
            if valid_time_slots:
                # Constrain time to valid slots for this credit level
                course_constraint_list.append(
                    cast(
                        z3.BoolRef,
                        z3.Or([c.time == z3_data.time_slot_constants[self._slots[i]] for i in valid_time_slots]),
                    )
                )

            if c.labs:
                # we must assign to a lab when we have options
                course_constraint_list.append(
                    cast(
                        z3.BoolRef,
                        z3.Or([c.lab == z3_data.lab_constants[lab] for lab in self._labs if lab in c.labs]),
                    )
                )
            if c.rooms:
                # we must assign to a room when we have options
                course_constraint_list.append(
                    cast(
                        z3.BoolRef,
                        z3.Or([c.room == z3_data.room_constants[room] for room in self._rooms if room in c.rooms]),
                    )
                )
            if c.faculties:
                # we must assign to a faculty from the candidates
                course_constraint_list.append(
                    cast(
                        z3.BoolRef,
                        z3.Or([c.faculty == z3_data.faculty_constants[faculty] for faculty in c.faculties]),
                    )
                )
            if conflict_constraints:
                course_constraint_list.append(cast(z3.BoolRef, z3.And(conflict_constraints)))

            course_constraints.append(cast(z3.BoolRef, z3.And(course_constraint_list)))

        return course_constraints

    def _build_resource_constraints(
        self,
        overlaps: z3.FuncDeclRef,
        lab_overlaps: z3.FuncDeclRef,
        lecture_next_to: z3.FuncDeclRef,
        lab_next_to: z3.FuncDeclRef,
        z3_data: _Z3SortsAndConstants,
    ) -> list[z3.BoolRef]:
        """
        Create resource sharing and faculty scheduling constraints.

        **Usage:**
        ```python
        self._build_resource_constraints(overlaps, lab_overlaps, lecture_next_to, lab_next_to, z3_data)
        ```

        **Args:**
        - overlaps: `z3.FuncDeclRef` function for checking time overlaps
        - lab_overlaps: `z3.FuncDeclRef` function for checking lab overlaps
        - lecture_next_to: `z3.FuncDeclRef` function for checking lecture next to each other
        - lab_next_to: `z3.FuncDeclRef` function for checking lab next to each other
        - z3_data: `_Z3SortsAndConstants` object containing the Z3 sorts and constants

        **Returns:**
        - `list[z3.BoolRef]` containing the resource constraints
        """
        resource_constraints: list[z3.BoolRef] = []

        for i, j in itertools.combinations(self._courses, 2):
            resource: list[z3.BoolRef] = []
            constraint_parts: list[z3.BoolRef] = []

            # Enforce same room usage when both courses can use the same rooms
            if set(i.rooms) & set(j.rooms):
                resource.append(
                    cast(
                        z3.BoolRef,
                        z3.Implies(
                            i.room == j.room,
                            z3.Not(overlaps(i.time, j.time)),
                        ),
                    )
                )
                if i.course_id == j.course_id:
                    # when a faculty teaches two sections of the same course,
                    # they must use the same room
                    constraint_parts.append(cast(z3.BoolRef, i.room == j.room))

            # Enforce same lab usage when both courses have labs and can use the same labs
            if set(i.labs) & set(j.labs):
                resource.append(
                    cast(
                        z3.BoolRef,
                        z3.Implies(
                            i.lab == j.lab,
                            z3.Not(lab_overlaps(i.time, j.time)),
                        ),
                    )
                )
                if i.course_id == j.course_id:
                    # when a faculty teaches two sections of the same course,
                    # they must use the same lab
                    constraint_parts.append(cast(z3.BoolRef, i.lab == j.lab))

            # Prevent time overlap for courses taught by same faculty
            constraint_parts.append(cast(z3.BoolRef, z3.Not(overlaps(i.time, j.time))))
            if i.course_id == j.course_id:
                # when a faculty teaches two sections of the same course,
                # they must be next to each other
                same_course_constraints = [lecture_next_to(i.time, j.time)]
                # Only require lab_next_to if the course has labs
                if i.labs:
                    same_course_constraints.append(lab_next_to(i.time, j.time))
                constraint_parts.append(cast(z3.BoolRef, z3.And(same_course_constraints)))
            else:
                # when a faculty teaches two sections of different courses,
                # they must not be next to each other
                diff_course_constraints = [z3.Not(lecture_next_to(i.time, j.time))]
                # Only require lab_next_to constraint if both courses have labs
                if i.labs and j.labs:
                    diff_course_constraints.append(z3.Not(lab_next_to(i.time, j.time)))
                constraint_parts.append(cast(z3.BoolRef, z3.And(diff_course_constraints)))

            if i.course_id == j.course_id:
                # prevent overlapping times for different sections of the same course
                resource_constraints.append(cast(z3.BoolRef, z3.Not(overlaps(i.time, j.time))))

            if resource:
                # add all resource constraints (room, lab, etc.)
                resource_constraints.append(cast(z3.BoolRef, z3.And(resource)))
            # add all course constraints when faculty is the same
            resource_constraints.append(
                cast(
                    z3.BoolRef,
                    z3.Implies(i.faculty == j.faculty, z3.And(constraint_parts)),
                )
            )

        return resource_constraints

    def _aggregate_constraints(
        self,
        function_constraints: list[z3.BoolRef],
        faculty_constraints: list[z3.BoolRef],
        course_constraints: list[z3.BoolRef],
        resource_constraints: list[z3.BoolRef],
    ) -> list[z3.BoolRef]:
        """
        Combine all constraints and apply simplification.

        **Usage:**
        ```python
        self._aggregate_constraints(function_c, faculty_c, course_c, resource_c)
        ```

        **Args:**
        - function_constraints: `list[z3.BoolRef]` containing the function constraints
        - faculty_constraints: `list[z3.BoolRef]` containing the faculty constraints
        - course_constraints: `list[z3.BoolRef]` containing the course constraints
        - resource_constraints: `list[z3.BoolRef]` containing the resource constraints

        **Returns:**
        - `list[z3.BoolRef]` containing the aggregated constraints
        """
        all_constraints: list[z3.BoolRef] = []

        for c in itertools.chain(
            function_constraints,
            faculty_constraints,
            course_constraints,
            resource_constraints,
        ):
            all_constraints.append(self._simplify(c))

        logger.debug(f"Added {len(function_constraints)} function constraints")
        logger.debug(f"Added {len(faculty_constraints)} faculty constraints")
        logger.debug(f"Added {len(course_constraints)} course constraints")
        logger.debug(f"Added {len(resource_constraints)} resource constraints")

        return all_constraints

    def _get_schedule(self, model: z3.ModelRef) -> list["CourseInstance"]:
        """
        Internal method to convert a Z3 model to a schedule of `CourseInstance` objects.

        **Usage:**
        ```python
        schedule = self._get_schedule(model)
        ```

        **Args:**
        - model: The Z3 model containing assignments

        **Returns:**
        - `list[CourseInstance]` representing the schedule
        """

        schedule = []
        for course in self._courses:
            slot = model.eval(course.time)
            time = self._z3_data.time_slot_constants.inverse[slot]
            faculty = self._z3_data.faculty_constants.inverse.get(model.eval(course.faculty), None)
            room = self._z3_data.room_constants.inverse.get(model.eval(course.room), None)
            lab = self._z3_data.lab_constants.inverse.get(model.eval(course.lab), None)

            if time is None or faculty is None:
                raise ValueError(f"Invalid model: {model}")

            # Create CourseInstance
            course_instance = CourseInstance(
                course=course,
                time=time,
                faculty=faculty,
                room=room,
                lab=lab,
            )
            schedule.append(course_instance)

        return schedule

    def _update(self, s: z3.Optimize):
        """
        Update the Z3 solver with the new constraints.

        **Usage:**
        ```python
        self._update(optimize_solver)
        ```

        **Args:**
        - s: `z3.Optimize` object containing the Z3 solver

        **Returns:**
        - `None`
        """
        m: z3.ModelRef = s.model()
        rearranged = []
        per_course = []
        # group courses by faculty first
        for _, group_iter in itertools.groupby(self._courses, key=lambda x: m[x.faculty]):
            group = list(group_iter)
            for _, cs_iter in itertools.groupby(group, key=lambda x: x.course_id):
                cs = list(cs_iter)
                if len(cs) > 1:
                    rearranged.append(
                        z3.And(
                            [z3.And(i.time != m[j.time], j.time != m[i.time]) for i, j in itertools.combinations(cs, 2)]
                        )
                    )
                for c in cs:
                    per_instance = []
                    per_instance.append(c.time == m[c.time])
                    if c.rooms:
                        per_instance.append(c.room == m[c.room])
                    if c.labs:
                        per_instance.append(c.lab == m[c.lab])
                    per_course.append(z3.Not(z3.And(per_instance)))

        if rearranged:
            logger.debug(f"Adding 1 course rearrangement constraint with {len(rearranged)} predicates")
            s.add(z3.And(rearranged))
        if per_course:
            logger.debug(f"Adding 1 per-course constraint with {len(per_course)} predicates")
            s.add(z3.Or(per_course))

    def get_models(self) -> Generator[list[CourseInstance], None, None]:
        """
        Generate schedules one-at-a-time using the Z3 solver.

        **Usage:**
        ```python
        for schedule in sched.get_models():
            ...
        ```

        **Returns:**
        Generator of lists of `CourseInstance` objects representing complete schedules

        **Example:**
        >>> full_config = load_config_from_file(CombinedConfig, "config.json")
        >>> scheduler = Scheduler(full_config)
        >>> for model in scheduler.get_models():
        ...     for course in model:
        ...         print(course.as_csv())
        """
        s = z3.Optimize(ctx=self._ctx)

        # Optimized solver configuration for EnumSort-based problems
        # Core optimization settings
        s.set("maxres.maximize_assignment", True)
        s.set("maxsat_engine", "maxres")
        s.set("optsmt_engine", "symba")
        s.set("enable_lns", True)
        s.set("maxres.max_core_size", 100)
        s.set("maxres.wmax", True)
        s.set("pb.compile_equality", True)
        s.set("priority", "pareto")

        for c in self._constraints:
            s.add(c)

        # Add faculty preferences as optimization goals with improved caching - only if requested
        if OptimizerFlags.FACULTY_COURSE in self._optimizer_flags:
            course_preference_terms = []
            for faculty_name, preferences in self._faculty_course_preferences.items():
                if not preferences:  # Skip faculty with no preferences
                    continue

                faculty_constant = self._z3_data.faculty_constants[faculty_name]
                for course in self._courses:
                    if course.course_id in preferences:
                        # Use preference value directly
                        # (1-5 scale where 5 is strongly prefer, 1 is weakest)
                        preference_value = preferences[course.course_id]
                        if preference_value == 0:
                            continue
                        term = z3.If(
                            course.faculty == faculty_constant,
                            preference_value,
                            0,
                        )
                        course_preference_terms.append(term)

            if course_preference_terms:
                n = len(course_preference_terms)
                logger.debug(
                    f"Adding {n} faculty course preference optimization goals",
                )
                s.maximize(z3.Sum(course_preference_terms))

        if OptimizerFlags.FACULTY_ROOM in self._optimizer_flags:
            room_preference_terms = []
            for faculty_name, preferences in self._faculty_room_preferences.items():
                if not preferences:  # Skip faculty with no preferences
                    continue

                faculty_constant = self._z3_data.faculty_constants[faculty_name]
                for course in self._courses:
                    for room in course.rooms:
                        room_constant = self._z3_data.room_constants[room]
                        if room in preferences:
                            preference_value = preferences[room]
                            if preference_value == 0:
                                continue
                            term = z3.If(
                                z3.And(
                                    course.faculty == faculty_constant,
                                    course.room == room_constant,
                                ),
                                preference_value,
                                0,
                            )
                            room_preference_terms.append(term)

            if room_preference_terms:
                n = len(room_preference_terms)
                logger.debug(
                    f"Adding {n} faculty room preference optimization goals",
                )
                s.maximize(z3.Sum(room_preference_terms))

        if OptimizerFlags.FACULTY_LAB in self._optimizer_flags:
            lab_preference_terms = []
            for faculty_name, preferences in self._faculty_lab_preferences.items():
                if not preferences:  # Skip faculty with no preferences
                    continue

                faculty_constant = self._z3_data.faculty_constants[faculty_name]
                for course in self._courses:
                    for lab in course.labs:
                        if lab in preferences:
                            preference_value = preferences[lab]
                            if preference_value == 0:
                                continue
                            term = z3.If(
                                z3.And(
                                    course.faculty == faculty_constant,
                                    course.lab == self._z3_data.lab_constants[lab],
                                ),
                                preference_value,
                                0,
                            )
                            lab_preference_terms.append(term)

            if lab_preference_terms:
                logger.debug(
                    f"Adding {len(lab_preference_terms)} faculty lab preference optimization goals",
                )
                s.maximize(z3.Sum(lab_preference_terms))

        same_rooms = []
        same_labs = []
        packing_rooms = []
        packing_labs = []
        for i, j in itertools.combinations(self._courses, 2):
            if set(i.rooms) & set(j.rooms):
                same_rooms.append(
                    z3.If(
                        z3.And(i.faculty == j.faculty, i.room == j.room),
                        1,
                        0,
                    )
                )
                if i.course_id != j.course_id:
                    packing_rooms.append(
                        z3.If(
                            z3.And(
                                i.room == j.room,
                                self._function_data.lecture_next_to(i.time, j.time),
                            ),
                            1,
                            0,
                        )
                    )
            if set(i.labs) & set(j.labs):
                same_labs.append(z3.If(z3.And(i.faculty == j.faculty, i.lab == j.lab), 1, 0))
                if i.course_id != j.course_id:
                    packing_labs.append(
                        z3.If(
                            z3.And(
                                i.lab == j.lab,
                                self._function_data.lab_next_to(i.time, j.time),
                            ),
                            1,
                            0,
                        )
                    )

        if same_rooms and OptimizerFlags.SAME_ROOM in self._optimizer_flags:
            logger.debug(f"Adding {len(same_rooms)} same room optimization goals")
            s.maximize(z3.Sum(same_rooms))
        if same_labs and OptimizerFlags.SAME_LAB in self._optimizer_flags:
            logger.debug(f"Adding {len(same_labs)} same lab optimization goals")
            s.maximize(z3.Sum(same_labs))
        if packing_rooms and OptimizerFlags.PACK_ROOMS in self._optimizer_flags:
            logger.debug(f"Adding {len(packing_rooms)} room packing optimization goals")
            s.maximize(z3.Sum(packing_rooms))
        if packing_labs and OptimizerFlags.PACK_LABS in self._optimizer_flags:
            logger.debug(f"Adding {len(packing_labs)} lab packing optimization goals")
            s.maximize(z3.Sum(packing_labs))

        if len(self._optimizer_flags) > 0:
            logger.debug("Created all optimization goals")
        else:
            logger.debug("Skipping optimization goals")

        for i in range(self._limit):
            start_time = time.time()
            if s.check() == z3.sat:
                generation_time = time.time() - start_time
                logger.debug(f"Schedule {i + 1} generation took {generation_time:.2f}s")
                yield self._get_schedule(s.model())
                if i < self._limit - 1:
                    self._update(s)
                    i += 1
            else:
                generation_time = time.time() - start_time
                if i == 0:
                    logger.error("No solution found")
                else:
                    logger.warning("No more solutions found")
                logger.debug(f"Final check took {generation_time:.2f} seconds")
                break
