from calendar import week
import datetime
import itertools
from typing import Dict, Iterable, List, Sequence, Tuple, Set, Any, Type, Literal, Optional
import time

from pydantic import BaseModel, ConfigDict, Extra, constr, field_validator, model_validator, StrictInt, Field

from pathlib import Path

import uuid

import datetime

from pyomo.environ import *
from pyomo.opt import SolverFactory, TerminationCondition, SolverStatus
import urllib.parse

from rota_generator.workers import Worker
from rota_generator.console import console
from rota_generator.bank_holidays import bank_holiday_map

import json

import csv

from collections import defaultdict

from io import StringIO
import sys

from rich.pretty import pprint
from rich.progress import track
from rich.panel import Panel
from rich import print

from loguru import logger

logger.add("rota.log", rotation="1 MB", level="DEBUG")

import ctypes
import platform

# Optional dependency: prefer setproctitle if available, fall back to libc on macOS
try:
    import setproctitle  # type: ignore
except Exception:
    setproctitle = None


def _get_proc_title() -> Optional[str]:
    try:
        if setproctitle is not None:
            return setproctitle.getproctitle()
        if platform.system() == "Darwin":
            libc = ctypes.CDLL("libc.dylib")
            libc.getprogname.restype = ctypes.c_char_p
            res = libc.getprogname()
            if res:
                return res.decode()
    except Exception:
        logger.debug("Could not read current process title")
    return None


def _set_proc_title(title: str) -> bool:
    try:
        if setproctitle is not None:
            setproctitle.setproctitle(title)
            return True
        if platform.system() == "Darwin":
            libc = ctypes.CDLL("libc.dylib")
            # setprogname expects a C string
            try:
                libc.setprogname(ctypes.c_char_p(title.encode()))
                return True
            except Exception:
                # some platforms may not support setprogname via ctypes
                pass
    except Exception:
        logger.debug("Failed to set process title to %s", title)
    return False


def _safe_constraint_info(obj):
    """Return a tuple (type_name, safe_dump) for logging without calling repr()

    Avoids triggering pydantic __repr__ which can call __getattr__ and raise
    when models come from mixed pydantic versions. Prefer model_dump() when
    available; fall back to __dict__ or the class name. Any failure returns a
    short error string.
    """
    tname = type(obj).__name__
    try:
        # Only call model_dump on real pydantic BaseModel instances. Some
        # objects may expose a model_dump attribute or similar but aren't
        # proper BaseModel instances in this runtime; calling model_dump on
        # those can raise PydanticUserError. Use isinstance check to be safe.
        if isinstance(obj, BaseModel) and hasattr(obj, "model_dump"):
            try:
                return tname, obj.model_dump()
            except Exception:
                # fall through to other strategies
                pass
        if hasattr(obj, "__dict__"):
            return tname, dict(obj.__dict__)
        return tname, str(type(obj))
    except Exception as e:
        return tname, f"<dump-failed: {e!r}>"


def _safe_as_dict(obj):
    """Return a plain dict view of obj without triggering pydantic __getattr__.

    Prefer model_dump(), then __dict__, then if obj is a mapping return it.
    Returns None if we couldn't obtain a dict safely.
    """
    try:
        # Only call model_dump on actual pydantic BaseModel instances. This
        # avoids raising PydanticUserError for objects that merely expose a
        # model_dump attribute but aren't proper BaseModel instances.
        if isinstance(obj, BaseModel) and hasattr(obj, "model_dump"):
            try:
                return obj.model_dump()
            except Exception:
                pass

        if isinstance(obj, dict):
            return obj

        if hasattr(obj, "__dict__"):
            try:
                return dict(obj.__dict__)
            except Exception:
                pass

    except Exception:
        # Defensive: anything going wrong, return None so caller can fall back
        return None
    return None


def _ensure_constraint_weeks(constraint, rota_builder):
    """Ensure constraint has a `weeks` value, computing from start/end if needed.

    This avoids directly reading attributes on pydantic models which may trigger
    problematic __getattr__ behavior in mixed-version environments. If we can
    extract a dict view, use that; otherwise fall back to guarded getattr/setattr.
    """
    # Don't mutate unknown constraint objects (some can be pydantic objects
    # from different versions that break __setattr__). Instead compute the
    # weeks and cache them on the rota_builder to be used by callers.
    try:
        if not hasattr(rota_builder, "_constraint_weeks_cache"):
            rota_builder._constraint_weeks_cache = {}

        key = id(constraint)
        if key in rota_builder._constraint_weeks_cache:
            return

        cdict = _safe_as_dict(constraint)
        if cdict is not None:
            weeks = cdict.get("weeks", None)
            if weeks is None:
                start = cdict.get("start_date", None)
                end = cdict.get("end_date", None)
                computed = rota_builder.get_weeks_by_date_range(start, end)
                rota_builder._constraint_weeks_cache[key] = computed
            else:
                rota_builder._constraint_weeks_cache[key] = weeks
            return

        # Fallback: try getattr but don't attempt to setattr on the object
        try:
            weeks = getattr(constraint, "weeks")
            rota_builder._constraint_weeks_cache[key] = weeks
            return
        except Exception:
            start = getattr(constraint, "start_date", None)
            end = getattr(constraint, "end_date", None)
            computed = rota_builder.get_weeks_by_date_range(start, end)
            rota_builder._constraint_weeks_cache[key] = computed
            return
    except Exception:
        tname, dump = _safe_constraint_info(constraint)
        logger.exception("Error accessing/setting 'weeks' on constraint type=%s dump=%s", tname, dump)
        raise


def _get_constraint_weeks(constraint, rota_builder):
    """Return the weeks list for a constraint using the cache or safe accessors.

    Never calls repr()/__repr__ on the constraint.
    """
    if hasattr(rota_builder, "_constraint_weeks_cache"):
        wk = rota_builder._constraint_weeks_cache.get(id(constraint), None)
        if wk is not None:
            return wk

    cdict = _safe_as_dict(constraint)
    if cdict is not None:
        return cdict.get("weeks", None)

    try:
        return getattr(constraint, "weeks")
    except Exception:
        return None

ShiftName = str
DayStr = str
WeekInt = int

days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]

full_weekend_count = 2


SHIFT_BOUNDS = {
    "bank_holiday": (0, 9),
    "shift_count": (0, 400),
    "night_shift_count": (0, 100),
    "weekend_count": (0, 60),
}

#VALID_SHIFT_CONSTRAINTS = Literal[
#    "night",
#    "pre",
#    "post",
#    "balance_across_groups",
#    "limit_grade_number",
#    "minimum_grade_number",
#    "max_shifts_per_week",
#    "max_shifts_per_week_block",
#    "require_remote_site_presence",
#    "require_remote_site_presence_week",
#]

class BaseShiftConstraint(BaseModel):
    weeks: Optional[List[int]] = Field(
        None,
        description=(
            "List of week numbers when this constraint is active. "
            "These are selectors that limit WHEN the constraint applies; they must not be used together with start_date/end_date."
        ),
    )
    start_date: Optional[datetime.date] = Field(
        None,
        description="Inclusive start date for the period when this constraint is active (cannot be used with `weeks`).",
    )
    end_date: Optional[datetime.date] = Field(
        None,
        description="Inclusive end date for the period when this constraint is active (cannot be used with `weeks`).",
    )

    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")

    @field_validator("weeks")
    def weeks_none_if_dates_set(cls, weeks, values):
        data = values.data
        if (data.get("start_date") is not None or data.get("end_date") is not None) and weeks is not None:
            raise ValueError("If start_date or end_date is set, weeks must not be set")
        return weeks


class ShiftConstraint(BaseShiftConstraint):
    """Compatibility wrapper for legacy constraint dicts of the form
    {'name': <constraint_name>, 'options': {...}}.

    The DB historically stored constraints as name/options pairs. When we
    encounter that shape we convert into this wrapper so validation passes
    and the builder code can still inspect `name`/`options` if needed.
    """
    name: str
    options: Any = None

class NightConstraint(BaseShiftConstraint):
    options: Any = None

class PreShiftConstraint(BaseShiftConstraint):
    """Pre-shift constraint.

    Semantics:
    - `weeks`, `start_date`, `end_date` (inherited) are selectors that determine WHEN the constraint is active (which weeks or date window).
    - `days` is a DURATION: the number of days the constraint applies for (e.g. 3 means apply the constraint for 3 days relative to the shift). This field is required.
    - `ignore_shifts`, `exclude_days`, `exclude_dates` further refine which shifts/dates are considered.
    """
    days: int = Field(
        ...,
        description=(
            "Duration in days: the number of days the constraint applies for. "
            "This is different from selector fields (weeks/start_date/end_date) which select WHEN the constraint is active."
        ),
    )
    ignore_shifts: List[str] = Field(default_factory=list, description="List of shift names to ignore when applying this constraint")
    exclude_days: List[str] = Field(default_factory=list, description="Weekday names to exclude (e.g. ['Mon','Tue'])")
    exclude_dates: List[datetime.date] = Field(default_factory=list, description="Specific dates to exclude from the constraint")
    allow_self: bool = Field(True, description="Allow the same worker to satisfy this constraint (default: True)")

class PostShiftConstraint(BaseShiftConstraint):
    """Post-shift constraint.

    See PreShiftConstraint docs: selectors (weeks/start_date/end_date) choose WHEN the constraint is active; `days` is a duration and is required.
    """
    days: int = Field(
        ...,
        description=(
            "Duration in days: the number of days the constraint applies for after the shift. "
            "Different from selector fields which choose which weeks/dates the constraint is active on."
        ),
    )
    ignore_shifts: List[str] = Field(default_factory=list, description="List of shift names to ignore when applying this constraint")
    exclude_days: List[str] = Field(default_factory=list, description="Weekday names to exclude (e.g. ['Mon','Tue'])")
    exclude_dates: List[datetime.date] = Field(default_factory=list, description="Specific dates to exclude from the constraint")
    allow_self: bool = Field(True, description="Allow the same worker to satisfy this constraint (default: True)")

class BalanceAcrossGroupsConstraint(BaseShiftConstraint):
    """"""


class MinSummedGradeByShiftsPerDayConstraint(BaseShiftConstraint):
    shifts: List[str] = Field(..., description="List of shift names to sum grades over")
    min_grade_sum: int = Field(..., description="Minimum total grade required")
    days: Optional[List[str]] = None

class LimitGradeNumberConstraint(BaseShiftConstraint):
    grade: int 
    max_number: int

class MinimumGradeNumberConstraint(BaseShiftConstraint):
    grade: int
    min_number: int

class RequireRemoteSitePresenceConstraint(BaseShiftConstraint):
    site: str
    required_number: int

class MaxShiftsPerWeekConstraint(BaseShiftConstraint):
    max_shifts: int | None = None
    days: Optional[List[str]] = None

class MaxShiftsPerWeekBlockConstraint(BaseShiftConstraint):
    week_block : int
    max_shifts: int | None = None

def _pydantic_to_dict(obj):
    # Recursively convert pydantic models, lists and dicts to JSON-serializable structures
    if obj is None:
        return None
    if isinstance(obj, (list, set, tuple)):
        return [_pydantic_to_dict(x) for x in obj]
    if isinstance(obj, dict):
        return {k: _pydantic_to_dict(v) for k, v in obj.items()}
    # datetime/date -> ISO string
    if isinstance(obj, (datetime.date, datetime.datetime)):
        return obj.isoformat()
    if hasattr(obj, "model_dump"):
        return _pydantic_to_dict(obj.model_dump())
    return obj


class WorkerRequirement(BaseModel):
    """
    start date is inclusive, end date is not
    """

    # NOTE: if number is changed mid week block allocations may not work?
    number: int = 1
    start_date: datetime.date | None = None
    end_date: datetime.date | None = None

    @model_validator(mode="after")
    def check_dates(self):
        if self.start_date is not None and self.end_date is not None:
            if self.start_date > self.end_date:
                raise ValueError("Start date must be before end date")
        return self


class SingleShift(BaseModel):
    """Class to hold all details for a shift

    start_date: The day that the shift starts on
    end_date: The day that the shift ends on (the last day it should be rota'd for)

    Due to week block assigments these (probably) need to be week ba

    Valid constraints

    balance_across_groups (requires block assignment)

    post / pre clear

    require_remote_site_presence_week:
        options: (site, required_number)

    night

    limit_grade_number:
        options: {grade: max_number, grade2: max_number2}

    minimum_grade_number:
        options: (grade, min_number)





    """

    sites: List[str]
    name: str
    length: float
    days: str | list[str]
    balance_offset: float | None = None
    balance_weighting: float = 1
    workers_required: StrictInt | list[WorkerRequirement] = 1
    rota_on_nwds: bool = False
    assign_as_block: bool = False
    force_as_block: bool = False
    force_as_block_unless_nwd: bool | List[List[str]] = False
    hard_constrain_shift: bool = True
    bank_holidays_only: bool = False
    #constraint: list[ShiftConstraint] = []
    # NOTE: avoid typing this as List[BaseModel] — instructing pydantic to
    # validate items as `BaseModel` can cause it to attempt to instantiate
    # `BaseModel` directly which raises a PydanticUserError. Use the
    # concrete BaseShiftConstraint base type so pydantic will validate into
    # actual constraint subclasses instead of BaseModel itself.
    constraints: List["BaseShiftConstraint"] = Field(default_factory=list)

    start_date: datetime.date | None = None
    end_date: datetime.date | None = None
    pair_proxy: bool = False
    display_char: str | None = None  # Add this line
    force_assign_with: list[str] = []  # List of shift names to force assign with this shift on the same day
    include_groups: set[str] = set()
    exclude_groups: set[str] = set()

    @model_validator(mode="after")
    def check_groups_non_overlapping(self) -> "SingleShift":
        intersection = self.include_groups.intersection(self.exclude_groups)
        if intersection:
            raise ValueError(
                f"include_groups and exclude_groups cannot overlap. Intersection: {intersection}"
            )
        return self

    model_config = ConfigDict(
        extra="allow",
        # orm_mode = True
    )

    def __init__(self, **data: Any):
        # Backwards-compat: accept legacy `constraints` list entries in the
        # form {'name': ..., 'options': {...}} and convert them to concrete
        # pydantic constraint models before validation. This avoids errors
        # where the DB stored legacy name/options pairs.
        if "constraints" in data and isinstance(data["constraints"], (list, tuple)):
            converted = []
            _name_map = {
                "pre": PreShiftConstraint,
                "post": PostShiftConstraint,
                "night": NightConstraint,
                "balance_across_groups": BalanceAcrossGroupsConstraint,
                "min_summed_grade_by_shifts_per_day": MinSummedGradeByShiftsPerDayConstraint,
                "limit_grade_number": LimitGradeNumberConstraint,
                "minimum_grade_number": MinimumGradeNumberConstraint,
                "require_remote_site_presence": RequireRemoteSitePresenceConstraint,
                "max_shifts_per_week": MaxShiftsPerWeekConstraint,
                "max_shifts_per_week_block": MaxShiftsPerWeekBlockConstraint,
            }
            for item in data["constraints"]:
                if isinstance(item, dict) and "name" in item:
                    cname = item.get("name")
                    opts = item.get("options") or {}
                    cls = _name_map.get(cname)
                    if cls is not None:
                        try:
                            inst = cls.model_validate(opts) if hasattr(cls, "model_validate") else cls(**(opts or {}))
                        except Exception:
                            # fallback to generic wrapper
                            inst = ShiftConstraint(name=cname, options=opts)
                    else:
                        inst = ShiftConstraint(name=cname, options=opts)
                    converted.append(inst)
                else:
                    converted.append(item)
            data["constraints"] = converted

        if "constraint" in data:
            raise ValueError(
                "The 'constraint' field on SingleShift has been deprecated and removed. "
                "Please use the 'constraints' field with Pydantic constraint models instead."
            )
        super().__init__(**data)
        # self.site = sites
        # self.name = name
        # self.length = length
        # self.days = days
        # self.balance_by_site = balance_by_site
        # balance_offset defines the max difference in allocated shifts
        # versus target shifts (hard constraint)
        # if 0 exactly equal shifts must be assigend (unlikely to be possible)

        if self.display_char is not None and len(self.display_char) > 1:
            raise ValueError("display_char must be a single character or None")

        if self.balance_offset is None:
            self.balance_offset = len(self.days)

        if self.force_as_block_unless_nwd:
            self.assign_as_block = True

        # self.assign_as_block = assign_as_block

        #self.constraint_options = {}
        #self.constraints = []
        #for c in self.constraint:
        #    self.constraints.append(c.name)

        #    if c.options:
        #        self.constraint_options[c.name] = c.options

            # constraint = c.constraint
        #            match c:
        #                case (constraint, options):
        #                    self.constraint_options[constraint] = options
        #                case constraint:
        #                    self.constraints.append(constraint)

        if isinstance(self.days, str):
            self.days = [self.days]

    def __str__(self):
        return (
            f"name: {self.name}, "
            f"sites: {self.sites}, "
            f"length: {self.length}, "
            f"balance_weighting: {self.balance_weighting}, "
            f"balance_offset: {self.balance_offset}, "
            f"hard_constrain_shift: {self.hard_constrain_shift}, "
            f"workers_required: {self.workers_required}, "
            f"rota_on_nwds: {self.rota_on_nwds}, "
            f"assign_as_block: {self.assign_as_block}, "
            f"force_as_block: {self.force_as_block}, "
            f"constraints: {self.constraints}"
        )

    def has_constraint(self, constraint: Type) -> bool:
        for c in self.constraints:
            if isinstance(c, constraint):
                return True
        return False

    def get_constraint(self, constraint: Type) -> BaseShiftConstraint | None:
        for c in self.constraints:
            if isinstance(c, constraint):
                return c
        return None

    def get_constraints(self, constraint: Type) -> list[BaseShiftConstraint]:
        return [c for c in self.constraints if isinstance(c, constraint)]

    def get_shift_summary(self):
        """
        Returns a text summary of the shift
        """

        return f"{self.name}: {self.workers_required} workers for sites ({', '.join(self.site)}) on days ({', '.join(self.days)})"

    def get_shift_number(self):
        return len(self.days)

    def get_worker_requirement_by_date(self, date: datetime.date):
        """Must be called after build_shifts

        Returns the number of workers required for a given date

        If none are required 0 is returned"""
        if isinstance(self.workers_required, int):
            return self.workers_required

        for wr in self.workers_required:
            # print(date, wr)
            if wr.start_date <= date < wr.end_date:
                return wr.number

        return 0

    def get_display_char(self):
        """Returns the display character for the shift"""
        if self.display_char is not None:
            return self.display_char
        return self.name[0]


class RotaConstraintOptions(BaseModel):
    """Pydantic model holding rota-level constraint options with types, defaults and descriptions.

    Validate/coerce raw JSON (e.g. from DB) via `RotaConstraintOptions.model_validate(raw)`.
    Use `model_dump()` to produce a JSON-serializable dict for storage.
    """

    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)

    balance_nights: bool = Field(True, description="Balance night shifts across workers")
    constrain_time_off_after_nights: bool = Field(False, description="Apply time-off constraints after night shifts")
    balance_nights_across_sites: bool = Field(True, description="Balance night shifts across sites")
    balance_bank_holidays: bool = Field(True, description="Balance bank-holiday shifts across staff")
    balance_blocks: bool = Field(True, description="Balance shifts across week-blocks")
    balance_shifts: bool = Field(True, description="Balance shifts across workers")
    balance_shifts_true_quadratic: bool = Field(False, description="Use true quadratic balancing for shifts")
    balance_shifts_quadratic: bool = Field(False, description="Use quadratic penalty to discourage spread across shifts")
    balance_shifts_over_workers: bool = Field(True, description="Balance shift distribution over workers")
    minimise_shift_diffs: bool = Field(False, description="Use simpler minimisation to reduce shift differences")
    maximum_allowed_shift_diff: Optional[int] = Field(None, description="Maximum allowed per-worker shift difference (None = no limit)")
    balance_weekends: bool = Field(True, description="Balance weekend allocations")
    max_weekends: int = Field(100, description="Maximum number of weekends considered")
    # If true, try to balance number of weekend days worked per worker separately
    # for Saturday and Sunday (i.e. equivalent workers should have similar Sat counts
    # and similar Sun counts).
    balance_weekend_days: bool = Field(False, description="Balance weekend days (Sat/Sun) separately")
    balance_weekend_days_weight: int = Field(5, description="Objective weight for balancing weekend days per day")
    balance_weekend_days_quadratic: bool = Field(False, description="Use quadratic penalty to balance weekend days per day")
    balance_weekend_days_use_previous_shifts: bool = Field(False, description="When true, adjust per-day weekend (Sat/Sun) targets using worker.previous_shifts (allocated-worked) distributed per shift day")
    weekend_day_hard_max_deviation: Optional[float] = Field(
        None,
        description=(
            "Hard maximum allowed absolute deviation (number) per-worker from the per-day "
            "weekend target (Sat/Sun). If set, adds constraints enforcing |assigned - target| <= value."
        ),
    )
    max_weekend_frequency: int = Field(1, description="Don't assign multiple shifts every N weeks (requires balance_weekends)")
    max_night_frequency: int = Field(1, description="Don't assign multiple night shifts every N weeks")
    max_night_frequency_week_exclusions: List[int] = Field(default_factory=list, description="Weeks to exclude from night frequency checks")
    max_shifts_per_week: int = Field(7, description="Max shifts per worker per week")
    max_shifts_per_month: int = Field(40, description="Max shifts per worker per month")
    max_days_per_week_block: List[tuple] = Field(default_factory=list, description="Optional list of (max_days, week_block) tuples to limit days in blocks")
    prevent_monday_after_full_weekends: List[int] = Field(default_factory=list, description="Prevent assignment on Monday after full weekends")
    prevent_monday_and_tuesday_after_full_weekends: List[int] = Field(default_factory=list, description="Prevent Monday/Tuesday after full weekends")
    prevent_fridays_before_full_weekends: List[int] = Field(default_factory=list, description="Prevent Fridays before full weekends")
    prevent_thursdays_before_full_weekends: List[int] = Field(default_factory=list, description="Prevent Thursdays before full weekends")
    hard_constrain_pair_separation: bool = Field(False, description="Enforce hard pair separation constraint")
    # These are lists of dicts with keys like {"grades": [...], "weeks": [...], "shifts": [...]}
    avoid_shifts_by_grades: List[dict] = Field(default_factory=list, description="List of rules to avoid assigning grades to shifts")
    avoid_shifts_by_worker_names: List[dict] = Field(default_factory=list, description="List of rules to avoid assigning specific workers to shifts")
    # Rules to avoid assigning specific workers to particular shifts on specific dates.
    # Each entry is a dict with keys: {"names": [...], "shifts": [...], "start_date": date|None, "end_date": date|None, "dates": [date,...]|None}
    avoid_shifts_by_worker_dates: List[dict] = Field(default_factory=list, description="List of rules to avoid assigning specific workers to shifts on particular dates or date ranges")
    allocate_locum_shifts: bool = Field(True, description="Allow allocation of locum shifts")
    distribute_locum_shifts: bool = Field(True, description="Distribute locum shifts among available workers")
    balance_locum_shifts: bool = Field(False, description="Balance locum shifts similarly to normal shifts")
    maximum_allowed_locum_shifts_per_worker: Optional[int] = Field(None, description="Maximum locum shifts allowed per worker (None = no limit)")
    minimum_allowed_locum_shifts_per_worker: Optional[int] = Field(None, description="Minimum locum shifts required per worker (None = no minimum)")
    min_summed_grade_by_shifts_per_day: List[MinSummedGradeByShiftsPerDayConstraint] = Field(default_factory=list, description="List of minimum-summed-grade constraints per day")

    # (Removed legacy dict-like convenience methods to enforce
    # attribute access on the typed `RotaConstraintOptions` model.)

class RotaBuilder(object):
    """Class to hold and manipulate shifts"""

    def __init__(
        self,
        start_date: datetime.date = datetime.date.today(),
        weeks_to_rota: int = 26,
        balance_offset_modifier: float = 1.0,
        ltft_balance_offset: int = 1,
        use_previous_shifts: bool = False,
        use_shift_balance_extra: bool = False,
        use_bank_holiday_extra: bool = False,
        allow_force_assignment_with_leave_conflict: bool = False,
        SHIFT_BOUNDS=SHIFT_BOUNDS,
        name: str = "",
        set_process_title: bool = True,
        bank_holidays: dict[datetime.date, str] = bank_holiday_map,
        constraint_options: RotaConstraintOptions | dict | None = None,
    ):
        self.name = name
        self.set_process_title = set_process_title
        # record original process title so we can restore it after the run
        try:
            self._original_proc_title = _get_proc_title()
        except Exception:
            self._original_proc_title = None

        console.print(
            Panel(
                f"""[white]
{locals()}
        """,
                style="green",
                title=f"Generating Rota: {name}",
            ),
            justify="left",
        )

        self.SHIFT_BOUNDS = SHIFT_BOUNDS
        self.shifts: List[SingleShift] = []  # type List[SingleShift]
        # self.allowed_multi_shift_sets: list[set[str]] = []

        self.use_previous_shifts = use_previous_shifts
        self.use_shift_balance_extra = use_shift_balance_extra
        self.use_bank_holiday_extra = use_bank_holiday_extra

        self.workers: list[Worker] = []

        # self.night_blocks = ["weekday", "weekend", "none"]

        self.sites: set = set()

        self.balance_offset_modifier = balance_offset_modifier
        self.ltft_balance_offset = ltft_balance_offset

        # Initialize constraint options with typed Pydantic model. Accept either
        # None (use defaults), a dict (raw values), or an existing model instance.
        # Store the model instance in `self.constraint_options_model` and also
        # expose it as `self.constraint_options` for minimal code changes.
        if constraint_options is None:
            opts_model = RotaConstraintOptions()
        elif isinstance(constraint_options, RotaConstraintOptions):
            opts_model = constraint_options
        elif isinstance(constraint_options, dict):
            opts_model = RotaConstraintOptions.model_validate(constraint_options)
        else:
            opts_model = RotaConstraintOptions()

        self.constraint_options_model: RotaConstraintOptions = opts_model

        logger.debug(f"Rota constraint options: {self.constraint_options_model.model_dump()}")

        self.terminate_on_warning = [
            "Worker/duplicate id",
            "Worker/duplicate name",
            "Worker/no valid shifts",
            "Shift/invalid start date",
            "Shift/invalid end date",
            "Locum/no locum availability",
            #"Force assignment/leave conflict", # Currently will always be infeasible
        ]

        self.results = None

        self.warnings: list[tuple[str, str]] = []

        self.ignore_valid_shifts = False

        self.set_rota_dates(start_date, weeks_to_rota)

        self.run_start_time: datetime.datetime | None = None
        self.run_end_time: datetime.datetime | None = None

        self.allow_force_assignment_with_leave_conflict = allow_force_assignment_with_leave_conflict

        self.bank_holidays = bank_holidays

        self.paired_shifts = []

        self.exported_rota_file = None

        self.enable_unavailabilities = True

        self.MAX_SHIFTS_PER_DAY = 2

    def set_rota_constraint(self, constraint: str, value: Any):
        # Ensure the model has the attribute and set it on the typed model
        if not hasattr(self.constraint_options_model, constraint):
            raise ValueError(f"Constraint {constraint} not valid")

        setattr(self.constraint_options_model, constraint, value)

    @property
    def constraint_options(self):
        """Backward-compat dict-like proxy for the typed
        `constraint_options_model`. This keeps older code/tests working
        while the canonical storage is `constraint_options_model`.
        """

        class _Proxy:
            def __init__(self, model):
                self._m = model

            def __getitem__(self, key):
                return getattr(self._m, key)

            def __setitem__(self, key, value):
                setattr(self._m, key, value)

            def get(self, key, default=None):
                return getattr(self._m, key, default)

            def items(self):
                return self._m.model_dump().items()

            def __contains__(self, key):
                return hasattr(self._m, key)

        return _Proxy(self.constraint_options_model)

    def set_rota_dates(self, start_date: datetime.date, weeks_to_rota: int):
        self.weeks_to_rota = weeks_to_rota

        self.days = days

        self.weeks = [i for i in range(1, weeks_to_rota + 1)]

        self.weeks_days_product = list(itertools.product(self.weeks, days))

        if start_date.weekday() != 0:
            raise ValueError(
                f"Start date {start_date.isoformat()} must be a Mon (not a {days[start_date.weekday()]})"
            )

        self.start_date = start_date
        self.rota_days_length = len(self.weeks) * 7
        self.rota_end_date = self.start_date + datetime.timedelta(self.rota_days_length)

        console.print(
            f"Weeks to rota: {weeks_to_rota} ({self.start_date} - {self.rota_end_date})"
        )

        # Generate a map for week, day combinations to a given date
        self.week_day_date_map = {}
        self.date_week_day_map = {}

        n: int = 0
        for week, day in self.get_week_day_combinations():
            d = self.start_date + datetime.timedelta(n)
            self.week_day_date_map[(week, day)] = d
            self.date_week_day_map[d] = (week, day)
            n = n + 1

        self.unavailable_to_work = set()
        self.unavailable_to_work_reason = {}
        self.pref_not_to_work = {}
        self.pref_not_to_work_reason = {}
        self.work_requests: set[
            Tuple[str | uuid.UUID | int, WeekInt, DayStr, ShiftName]
        ] = set()
        self.work_requests_map = {}
        self.locum_availability: set[
            Tuple[str | uuid.UUID | int, WeekInt, DayStr, ShiftName]
        ] = set()
        self.locum_availability_map = {}

    def solve_model(
        self,
        solver: str = "appsi_highs",
        options: dict = {},
        debug_if_fail: bool = False,
    ):
        console.print("Setting up solver")
        console.print(f"{options=}")

        if solver == "scip":
            if "seconds" in options:
                options["limits/time"] = options.pop("seconds")
            if "ratio" in options:
                options["limits/gap"] = options.pop("ratio")
            if "threads" in options:
                options.pop("threads")
            self.opt = SolverFactory(solver, executable="scip")
        elif solver == "appsi_highs":
            self.opt = SolverFactory(solver)
            if "seconds" in options:
                options["time_limit"] = options.pop("seconds")
            if "ratio" in options:
                options["mip_rel_gap"] = options.pop("ratio")
            if "threads" in options:
                options.pop("threads")
        else:
            self.opt = SolverFactory(solver)

        try:
            console.print("Solving")
            console.print(f"Options: {options}")

            log_file = f"logs/{self.name}_{datetime.datetime.now().strftime('%Y%m%d-%H%M%S')}.log"
            results = self.opt.solve(
                self.model,
                tee=True,
                options=options,
                # options={
                #    "threads": 10,
                # },
                # logfile=log_file,
                keepfiles=True,
                load_solutions=False,
            )
        except KeyboardInterrupt:
            return
            pass

        if (
            not results.solver.termination_condition == TerminationCondition.infeasible
            # and not results.solver.status == SolverStatus.aborted
        ):
            self.model.solutions.load_from(results)

        self.results = results

        console.print(f"Complete - outcome: {results.solver.status}")
        console.print(f"Termination condition: {results.solver.termination_condition}")

        if results.solver.status != "ok" and debug_if_fail:
            console.print(f"Attempting each shift individually")
            self.solve_shifts_individually(options)

        # if not results.solver.status:
        #    sys.exit(0)

    def solve_shifts_by_block(self, options, block_length: int, folder="block_shifts"):
        shifts = self.shifts

        outcomes = {}

        no_blocks = self.weeks_to_rota - block_length

        original_start_date = self.start_date

        for block in range(0, no_blocks):
            start_time = time.time()
            start_date = original_start_date + datetime.timedelta(weeks=block)
            self.set_rota_dates(start_date=start_date, weeks_to_rota=block_length)
            console.print(
                f"Testing block: start_date - {self.start_date} , length {self.weeks_to_rota} weeks"
            )
            self.clear_shifts()
            self.add_shifts(*shifts)
            self.build_and_solve(options)
            self.export_rota_to_html(
                filename=f"{start_date}_{self.weeks_to_rota}", folder=folder
            )
            end_time = time.time()
            time_taken = end_time - start_time

            outcomes[f"{start_date}_{self.weeks}"] = (
                self.results.solver.status,
                self.results.solver.termination_condition,
                time_taken,
            )

        console.print(outcomes)
        sys.exit(0)

    def solve_shifts_individually(self, options, folder="individual_shifts"):
        shifts = self.shifts

        outcomes = {}

        for shift in shifts:
            console.print(f"Testing shifts: {shift.name}")
            self.ignore_valid_shifts = True
            self.clear_shifts()
            self.add_shift(shift)
            self.build_and_solve(options)
            self.export_rota_to_html(filename=shift.name, folder=folder)

            outcomes[shift.name] = (
                self.results.solver.status,
                self.results.solver.termination_condition,
            )

        console.print(outcomes)

    def build_and_solve(
        self,
        options={"ratio": 0.1, "seconds": 1000, "threads": 10},
        solve=True,
        debug_if_fail: bool = False,
        export=False,
        export_with_timestamp=False,
        solver="appsi_highs",
    ):
        # Optionally set an informative process title so OS tools show which rota is running
        self._proc_title_changed = False
        if getattr(self, "set_process_title", False):
            title = f"rota-gen:{self.name}" if self.name else "rota-gen"
            try:
                if _set_proc_title(title):
                    self._proc_title_changed = True
                    logger.debug("Process title set to %s", title)
            except Exception:
                logger.debug("Could not set process title")

        self.run_start_time = datetime.datetime.now()

        self.build_shifts()
        self.build_workers()
        self.build_model()

        console.print("Building model")
        console.print(f"{options=}")

        if solve:
            self.solve_model(
                options=options, debug_if_fail=debug_if_fail, solver=solver
            )

        self.run_end_time = datetime.datetime.now()

        # Restore the original process title if we changed it
        if getattr(self, "_proc_title_changed", False):
            try:
                if getattr(self, "_original_proc_title", None) is not None:
                    _set_proc_title(self._original_proc_title)
                else:
                    # best-effort restore to a generic python name
                    _set_proc_title("python")
            except Exception:
                logger.debug("Could not restore original process title")

        if export:
            self.export_rota_to_html(self.name, timestamp_filename=export_with_timestamp)

    def build_model(self):
        # Initialize model
        self.model = ConcreteModel()

        # binary variables representing if a worker is scheduled somewhere
        # NOTE: this will assign to all possible shift combinations (which we probably don't want)
        self.model.works = Var(
            (
                (worker.id, week, day, shiftname)
                for worker in self.workers
                for week, day, shiftname in self.get_all_shiftname_combinations()
            ),
            within=Binary,
            initialize=0,
        )

        self.model.works_day = Var(
            (
                (worker.id, week, day)
                for worker in self.workers
                for week, day in self.get_week_day_combinations()
            ),
            within=Binary,
            initialize=0,
        )

        self.model.shift_week_worker_assigned = Var(
            (
                (shift, week, worker.id)
                for worker in self.workers
                for week in self.weeks
                for shift in self.get_shift_names()
            ),
            within=Binary,
            initialize=0,
        )

        blocks_worker_keys = []
        for worker, week, shift_name in self.worker_week_shifts_to_assign_as_blocks():
            shift = self.get_shift_by_name(shift_name)
            sub_blocks = self.get_shift_sub_blocks(shift)
            for sb_idx in range(len(sub_blocks)):
                blocks_worker_keys.append((worker.id, week, shift_name, sb_idx))

        self.model.blocks_worker_shift_assigned = Var(
            blocks_worker_keys,
            within=Binary,
            initialize=0,
        )

        blocks_assigned_keys = []
        for week, shift_name in self.week_shifts_to_assign_as_blocks():
            shift = self.get_shift_by_name(shift_name)
            sub_blocks = self.get_shift_sub_blocks(shift)
            for sb_idx in range(len(sub_blocks)):
                blocks_assigned_keys.append((week, shift_name, sb_idx))

        self.model.blocks_assigned = Var(
            blocks_assigned_keys,
            within=NonNegativeIntegers,
            initialize=0,
        )

        # Helper binary variables for per-worker unique-shift block limits.
        # Index is (worker_id, block_length, block_start_week, shift_name).
        unique_shift_block_indices = []
        for worker in self.workers:
            for constraint in worker.max_unique_shifts_per_week_block:
                for week_blocks in self.get_week_block_iterator(constraint.week_block):
                    if not week_blocks:
                        continue
                    block_start_week = week_blocks[0]
                    for shift_name in self.get_shift_names():
                        if any(
                            shift_name in self.get_shift_names_by_week_day(
                                week,
                                day,
                                return_empty_if_week_day_not_found=True,
                            )
                            for week in week_blocks
                            for day in days
                        ):
                            unique_shift_block_indices.append(
                                (worker.id, constraint.week_block, block_start_week, shift_name)
                            )

        self.model.worker_shift_used_in_week_block = Var(
            unique_shift_block_indices,
            within=Binary,
            initialize=0,
        )

        if self.constraint_options_model.balance_bank_holidays:
            # Bank holidays
            self.model.bank_holiday_count = Var(
                ((worker.id) for worker in self.workers),
                within=NonNegativeIntegers,
                initialize=0,
                bounds=self.SHIFT_BOUNDS["bank_holiday"],
            )
            self.model.bank_holiday_count_w = Var(
                ((worker.id) for worker in self.workers),
                within=NonNegativeIntegers,
                initialize=0,
            )

        self.model.multi_shift_together = Var(
            (
                (worker.id, week, day, idx)
                for worker in self.workers
                for week, day in self.get_week_day_combinations()
                for idx, allowed_set in enumerate(
                    getattr(worker, "allowed_multi_shift_sets", [])
                )
                if allowed_set.issubset(self.get_shift_names_by_week_day(week, day))
            ),
            domain=Binary,
            initialize=0,
        )

        # Used to limit number of workers on night shift per site
        # if we force a binary it will in effect hard constrain to < 2
        # from a single site
        if self.constraint_options_model.balance_nights_across_sites:
            self.model.night_per_site = Var(
                (
                    (week, shift.name, site)
                    for site in self.sites
                    for week in self.weeks
                    for shift in self.get_shifts_with_constraint(NightConstraint)
                ),
                # within=Binary,
                within=NonNegativeIntegers,
                initialize=0,
            )

            self.model.night_per_site_t1 = Var(
                (
                    (week, shift.name, site)
                    for site in self.sites
                    for week in self.weeks
                    for shift in self.get_shifts_with_constraint(NightConstraint)
                ),
                domain=NonNegativeIntegers,
                initialize=0,
            )
            self.model.night_per_site_t2 = Var(
                (
                    (week, shift.name, site)
                    for site in self.sites
                    for week in self.weeks
                    for shift in self.get_shifts_with_constraint(NightConstraint)
                ),
                domain=NonNegativeIntegers,
                initialize=0,
            )

        # Replacement for balance_nights_across_sites
        self.model.shift_per_site = Var(
            (
                (week, shift.name, site)
                for site in self.sites
                for week in self.weeks
                for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint)
            ),
            # within=Binary,
            within=NonNegativeIntegers,
            initialize=0,
        )

        self.model.shift_per_site_t1 = Var(
            (
                (week, shift.name, site)
                for site in self.sites
                for week in self.weeks
                for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint)
            ),
            domain=NonNegativeIntegers,
            initialize=0,
        )
        self.model.shift_per_site_t2 = Var(
            (
                (week, shift.name, site)
                for site in self.sites
                for week in self.weeks
                for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint)
            ),
            domain=NonNegativeIntegers,
            initialize=0,
        )

        self.model.shift_count = Var(
            (
                (worker.id, shift)
                for worker in self.workers
                for shift in self.get_shift_names()
            ),
            domain=NonNegativeReals,
            initialize=0,
            bounds=self.SHIFT_BOUNDS["shift_count"],
        )

        self.model.shift_count_diff = Var(
            (
                (worker.id, shift)
                for worker in self.workers
                for shift in self.get_shift_names()
            ),
            within=Reals,
            initialize=0,
        )

        self.model.shift_count_diff_summed = Var(
            (worker.id for worker in self.workers),
            within=Reals,
            initialize=0,
        )

        if self.constraint_options_model.balance_shifts_over_workers:
            self.model.worker_shift_count_t1 = Var(
                ((worker.id) for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
            )

            self.model.worker_shift_count_t2 = Var(
                ((worker.id) for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
            )

        if (
            self.constraint_options_model.balance_shifts
            or self.constraint_options_model.balance_shifts_quadratic
        ):
            self.model.shift_count_t1 = Var(
                (
                    (worker.id, shift.name)
                    for worker in self.workers
                    for shift in self.get_shifts()
                ),
                domain=NonNegativeReals,
                initialize=0,
            )

            self.model.shift_count_t2 = Var(
                (
                    (worker.id, shift.name)
                    for worker in self.workers
                    for shift in self.get_shifts()
                ),
                domain=NonNegativeReals,
                initialize=0,
            )

        if self.constraint_options_model.balance_shifts_quadratic:
            self.model.shift_count_w = Var(
                (
                    (worker.id, shift.name)
                    for worker in self.workers
                    for shift in self.get_shifts()
                ),
                domain=NonNegativeReals,
                initialize=0,
            )

        if self.constraint_options_model.balance_nights:
            # We also try to even out the night shifts seperately
            self.model.night_shift_count = Var(
                ((worker.id) for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
                bounds=self.SHIFT_BOUNDS["night_shift_count"],
            )

            self.model.night_shift_count_t1 = Var(
                ((worker.id) for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
            )

            self.model.night_shift_count_t2 = Var(
                ((worker.id) for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
            )

            self.model.night_shift_count_w = Var(
                ((worker.id) for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
            )

        # self.model.worker_weekend_assigned = Var(
        #     ((worker.id, week) for worker in self.workers
        #      for week in self.weeks),
        #     domain=NonNegativeReals,
        #     initialize=0,
        # )

        # self.model.works_weekend_count = Var(
        #     ((worker.id, week) for worker in self.workers
        #      for week in self.weeks),
        #     domain=NonNegativeIntegers,
        #     initialize=0,
        # )
        # self.model.works_saturday = Var(
        #     ((worker.id, week) for worker in self.workers
        #      for week in self.weeks),
        #     domain=Binary,
        #     initialize=0,
        # )
        # self.model.works_sunday = Var(
        #     ((worker.id, week) for worker in self.workers
        #      for week in self.weeks),
        #     domain=Binary,
        #     initialize=0,
        # )

        # self.model.works_whole_weekend = Var(
        #     ((worker.id, week) for worker in self.workers
        #      for week in self.weeks),
        #     domain=Binary,
        #     initialize=0,
        # )

        self.model.works_weekend = Var(
            ((worker.id, week) for worker in self.workers for week in self.weeks),
            domain=Binary,
            initialize=0,
        )

        self.model.worker_weekend_count = Var(
            ((worker.id) for worker in self.workers),
            domain=NonNegativeIntegers,
            initialize=0,
            bounds=self.SHIFT_BOUNDS["weekend_count"],
        )

        if self.constraint_options_model.balance_weekends:
            self.model.weekend_shift_count_t1 = Var(
                ((worker.id) for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
            )

            self.model.weekend_shift_count_t2 = Var(
                ((worker.id) for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
            )

            self.model.weekend_shift_count_w = Var(
                ((worker.id) for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
            )

        # Per-day weekend balancing (separate Sat / Sun counts)
        # Also create the bookkeeping variables when a hard deviation constraint is requested
        logger.debug(f"Weekend day balancing options: {self.constraint_options_model.balance_weekend_days=}, {self.constraint_options_model.balance_weekend_days_quadratic=}, {self.constraint_options_model.weekend_day_hard_max_deviation=}")
        if (
            self.constraint_options_model.balance_weekend_days
            or self.constraint_options_model.balance_weekend_days_quadratic
            or self.constraint_options_model.weekend_day_hard_max_deviation is not None
        ):
            logger.debug("Creating per-day weekend balancing variables")
            self.model.weekend_day_shift_count = Var(
                ((worker.id, d) for worker in self.workers for d in ("Sat", "Sun")),
                domain=NonNegativeReals,
                initialize=0,
            )

            self.model.weekend_day_shift_count_t1 = Var(
                ((worker.id, d) for worker in self.workers for d in ("Sat", "Sun")),
                domain=NonNegativeReals,
                initialize=0,
            )

            self.model.weekend_day_shift_count_t2 = Var(
                ((worker.id, d) for worker in self.workers for d in ("Sat", "Sun")),
                domain=NonNegativeReals,
                initialize=0,
            )

            self.model.weekend_day_shift_count_w = Var(
                ((worker.id, d) for worker in self.workers for d in ("Sat", "Sun")),
                domain=NonNegativeReals,
                initialize=0,
            )

        # self.model.weekend_count_t1 = Var(
        #     ((worker.id) for worker in self.workers),
        #     domain=NonNegativeReals,
        #     initialize=0,
        # )

        # self.model.weekend_count_t2 = Var(
        #     ((worker.id) for worker in self.workers),
        #     domain=NonNegativeReals,
        #     initialize=0,
        # )

        # self.model.unavailable = Var(((worker, week, day) for worker in self.workers for week in weeks for day in days),
        #                  within=Binary, initialize=0)


        for worker in self.workers:
            for week, day, shift_name in getattr(worker, "forced_assignments", []):
                # Only add if this shift is valid for this week/day
                if shift_name in self.get_shift_names_by_week_day(week, day, return_empty_if_week_day_not_found=True):
                    # Check for leave conflict
                    if hasattr(self, "unavailable_to_work"):
                        if (worker.id, week, day) in self.unavailable_to_work:
                            if self.allow_force_assignment_with_leave_conflict:
                                self.unavailable_to_work.remove((worker.id, week, day))
                                self.add_warning(
                                    "Force assignment/leave conflict",
                                    f"Worker '{worker.name}' has a forced assignment for shift '{shift_name}' on week {week}, day {day} (date: {self.get_date_by_week_day(week, day)}), which conflicts with a leave request. The leave request has been overridden.",
                                )
                            else:
                                self.add_warning(
                                    "Force assignment/leave conflict",
                                    f"Worker '{worker.name}' has a forced assignment for shift '{shift_name}' on week {week}, day {day} (date: {self.get_date_by_week_day(week, day)}), which conflicts with a leave request.",
                                )

        if not self.enable_unavailabilities:
            self.unavailable_to_work = set()
            self.unavailable_to_work_reason = {}

        def availability_init(model, wid, week, day):
            # This needs to be higher than the max number of possible shifts per day
            if (wid, week, day) in self.unavailable_to_work:
                # print((wid, week, day))
                return 0
            return self.MAX_SHIFTS_PER_DAY + 1

        self.model.available = Param(
            (
                (worker.id, week, day)
                for worker in self.workers
                for week, day in self.get_week_day_combinations()
            ),
            initialize=availability_init,
        )

        locum_request_sets = set()
        for locum_request in self.locum_availability:
            if locum_request[3] == "*":
                pass

            else:
                locum_request_sets.add(locum_request[3])

        # Validate locum shifts are valid
        invalid_shifts = locum_request_sets - set(self.get_shift_names())
        if invalid_shifts and not self.ignore_valid_shifts:
            raise InvalidShift(
                f"Invalid locum shift request [shift(s): ${invalid_shifts}]"
            )

        def locum_request_init(model, wid, week, day, shift):
            if (wid, week, day, shift) in self.locum_availability:
                self.locum_availability_map[(wid, week, day)] = shift
                return 1
            return 0

        self.model.locum_availability = Param(
            (
                (worker.id, week, day, shift)
                for worker in self.workers
                if not worker.locum
                for week, day, shift in self.get_all_shiftname_combinations()
            ),
            initialize=locum_request_init,
        )

        work_request_sets = set()
        for work_request in self.work_requests:
            if work_request[3] == "*":
                pass

            else:
                work_request_sets.add(work_request[3])

        # Validate shifts are valid
        invalid_shifts = work_request_sets - set(self.get_shift_names())
        if invalid_shifts and not self.ignore_valid_shifts:
            raise InvalidShift(
                f"Invalid worker shift request [shift(s): ${invalid_shifts}]"
            )

        def work_request_init(model, wid, week, day, shift):
            if (wid, week, day, shift) in self.work_requests:
                self.work_requests_map[(wid, week, day)] = shift
                return 1
            return 0

        self.model.work_requests = Param(
            (
                (worker.id, week, day, shift)
                for worker in self.workers
                for week, day, shift in self.get_all_shiftname_combinations()
            ),
            initialize=work_request_init,
        )

        def pref_init(model, wid, week, day):
            if (wid, week, day) in self.pref_not_to_work:
                return self.pref_not_to_work[(wid, week, day)]
            return 0

        # It would also be possible to assign sequential days as blocks
        # which would allow more weight to be given if the whole block
        # is not allocated
        self.model.pref_not_to_work = Param(
            (
                (worker.id, week, day)
                for worker in self.workers
                for week, day in self.get_week_day_combinations()
            ),
            initialize=pref_init,
        )

        # binary variables representing if a worker is necessary
        # self.model.needed = Var([worker.id for worker in self.workers], within=Binary, initialize=0)

        # binary variables representing if a worker worked on sunday but not on saturday (avoid if possible)
        # self.model.no_pref = Var([worker.id for worker in self.workers], within=Binary, initialize=0)

        if self.get_workers_who_require_locums():
            # We use this as a time check that their is some valid locum ability
            if not self.get_all_locum_availability():
                self.add_warning(
                    "Locum/no locum availability", "No locum availability set"
                )

            self.model.locum_required = Var(
                (
                    (week, day, shiftname)
                    for week, day, shiftname in self.get_all_shiftname_combinations()
                ),
                within=Binary,
                initialize=0,
            )

            self.model.locum_works = Var(
                (
                    (worker.id, week, day, shiftname)
                    for worker in self.workers
                    for week, day, shiftname in self.get_all_shiftname_combinations()
                ),
                within=Binary,
                initialize=0,
            )

            self.model.locum_shifts_by_worker = Var(
                (worker.id for worker in self.workers),
                within=Integers,
                initialize=0,
            )

            self.model.locum_shifts_t1 = Var(
                (worker.id for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
            )

            self.model.locum_shifts_t2 = Var(
                (worker.id for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
            )

        self.build_model_constraints()

    def build_model_constraints(self):
        self.model.constraints = ConstraintList()  # Create a set of constraints


        for constraint in self.constraint_options_model.min_summed_grade_by_shifts_per_day:
            # Safely extract weeks/days without invoking pydantic __getattr__
            cdict = _safe_as_dict(constraint)
            if cdict is not None:
                weeks = cdict.get("weeks") if cdict.get("weeks") is not None else self.weeks
                days_list = cdict.get("days") if cdict.get("days") is not None else self.days
            else:
                try:
                    weeks = constraint.weeks if constraint.weeks is not None else self.weeks
                    days_list = constraint.days if constraint.days is not None else self.days
                except Exception:
                    tname, dump = _safe_constraint_info(constraint)
                    logger.exception("Error accessing 'weeks' or 'days' on constraint type=%s dump=%s", tname, dump)
                    raise
            for week in weeks:
                for day in days_list:
                    # Only consider if all shifts are present that day
                    if all(shift in self.get_shift_names_by_week_day(week, day) for shift in constraint.shifts):
                        expr_terms = [
                            worker.grade * self.model.works[worker.id, week, day, shift]
                            for shift in constraint.shifts
                            for worker in self.get_workers_for_shift(self.get_shift_by_name(shift))
                        ]
                        if expr_terms:  # Only add constraint if there are eligible workers
                            self.model.constraints.add(
                                sum(expr_terms) >= constraint.min_grade_sum
                            )



        # Ensure each shift has the requisit number of workers assigned
        for (
            week,
            day,
            shift,
            workers_required,
            site_required,
        ) in self.get_required_workers_and_site_combinations():
            shift_obj = self.get_shift_by_name(shift)
            eligible_workers = [
                worker for worker in self.workers
                if self.is_worker_eligible_for_shift(worker, shift_obj)
            ]
            self.model.constraints.add(
                workers_required
                == sum(
                    self.model.works[worker.id, week, day, shift]
                    for worker in eligible_workers
                )
            )

            if self.get_workers_who_require_locums():
                self.model.constraints.add(
                    self.model.locum_required[week, day, shift]
                    == sum(
                        self.model.works[worker.id, week, day, shift]
                        for worker in eligible_workers
                        if worker.locum
                    )
                )

                self.model.constraints.add(
                    self.model.locum_required[week, day, shift]
                    == sum(
                        self.model.locum_works[worker.id, week, day, shift]
                        for worker in self.workers
                        # if not worker.locum
                    )
                )

                self.model.constraints.add(
                    0
                    == sum(
                        self.model.locum_works[worker.id, week, day, shift]
                        for worker in self.workers
                        if worker.locum
                    )
                )

            # And it is not assigned if the worker is ineligible
            ineligible_workers = [
                worker for worker in self.workers
                if not self.is_worker_eligible_for_shift(worker, shift_obj)
            ]
            if ineligible_workers:
                self.model.constraints.add(
                    0
                    == sum(
                        self.model.works[worker.id, week, day, shift]
                        for worker in ineligible_workers
                    )
                )
            # # Ensure shifts are only assigned by workers from the allowed sites
            # self.model.constraints.add(0 == sum(
            #     self.model.works[worker.id, week, day, shift]
            #     for worker in self.workers if worker.site not in site_required))

        # Constraint: total hours worked hours worked
        # for worker in self.workers:
        #    self.model.constraints.add(
        #        1200 >= sum(shift_lengths[shift] * self.model.works[worker, week, day, shift] for week in weeks for day in days for shift in days_shifts[day])
        #    )

        # for week in weeks:
        #    for worker in self.workers:
        #        self.model.constraints.add(
        #            48 >= sum(shift_lengths[shift] * self.model.works[worker, week, day, shift] for day in days for shift in days_shifts[day])
        #        )

        # Define a maximum number hours that can be worked per week
        # NOTE: this works on a calander week not a 7 day period

        # TODO: consider fixing
        # def maxHoursPerWeekRule(model, wid, week):
        #    return 72 >= sum(
        #        shift.length * model.works[wid, week, day, shift.name]
        #        for day, shift in self.get_day_shiftclass_combinations())

        # self.model.max_hours_per_week_constraint = Constraint(
        #    [worker.id for worker in self.workers],
        #    self.weeks,
        #    rule=maxHoursPerWeekRule)

        # Define an upper limit for hours worked over the rota period
        # def maxHoursRule(model, wid):
        #    return 8000 >= sum(
        #        shift.length * model.works[wid, week, day, shift.name]
        #        for week, day, shift in self.get_all_shiftclass_combinations())

        # self.model.max_hours_constraint = Constraint(
        #    [worker.id for worker in self.workers], rule=maxHoursRule)

        # Limit grades
        def limitGradesRule(model, week, shift, grade, limit):
            workers = [w for w in self.workers if w.grade == grade]
            if not workers:
                return Constraint.Skip
            return (
                sum(
                    model.shift_week_worker_assigned[shift, week, w.id] for w in workers
                )
                <= limit
            )

        for shift in self.get_shifts_with_constraint(LimitGradeNumberConstraint):
            for constraint in shift.get_constraints(LimitGradeNumberConstraint):
                grade = constraint.grade
                limit = constraint.max_number
                # self.model.limit_grades_constraint = Constraint(
                setattr(
                    self.model,
                    f"limit_grades_constraint_{shift.name}_{grade}",
                    Constraint(
                        [week for week in self.weeks],
                        # [shift for shift in self.get_shifts_with_constraint("limit_grade_number")],
                        [shift.name],
                        [grade],
                        [limit],
                        rule=limitGradesRule,
                        # name=f"limit_grades_constraint_{shift.name}_{grade}"
                    ),
                )

        def minimumGradesRule(model, week, shift, grade, limit):
            workers = [w for w in self.workers if w.grade >= grade]
            if len(workers) < limit:
                # return Constraint.Skip
                raise ValueError("minimum_grade_number: not enough valid workers")
            return (
                sum(
                    model.shift_week_worker_assigned[shift, week, w.id] for w in workers
                )
                >= limit
            )

        for shift in self.get_shifts_with_constraint(MinimumGradeNumberConstraint):
            constraint = shift.get_constraints(MinimumGradeNumberConstraint)[0]
            grade, min_required = constraint.grade, constraint.min_number

            # self.model.limit_grades_constraint = Constraint(
            setattr(
                self.model,
                f"minimum_grade_number_{shift.name}",
                Constraint(
                    [week for week in self.weeks],
                    # [shift for shift in self.get_shifts_with_constraint("limit_grade_number")],
                    [shift.name],
                    [grade],
                    [min_required],
                    rule=minimumGradesRule,
                    # name=f"limit_grades_constraint_{shift.name}_{grade}"
                ),
            )

        def presenceAtRemoteSiteWeek(
            model, week, shift, required_site, required_number
        ):
            required_site_workers = [
                w for w in self.workers if required_site == w.remote_site
            ]

            return (
                sum(
                    model.shift_week_worker_assigned[shift, week, w.id]
                    for w in required_site_workers
                )
                >= required_number
            )

        for shift in self.get_shifts_with_constraint(
            RequireRemoteSitePresenceConstraint
        ):
            constraint = shift.get_constraints(RequireRemoteSitePresenceConstraint)[0]
            site, required_number = constraint.site, constraint.required_number

            setattr(
                self.model,
                f"require_remote_site_presence_week_{shift.name}",
                Constraint(
                    [week for week in self.weeks],
                    [shift.name],
                    [site],
                    [required_number],
                    rule=presenceAtRemoteSiteWeek,
                ),
            )

        #def presenceAtRemoteSite(
        #    model, day, week, shift, required_site, required_number
        #):
        #    required_site_workers = [
        #        w for w in self.workers if required_site == w.remote_site
        #    ]

        #    for w in required_site_workers:
        #        if (w.id, week, day, shift) not in model.works:
        #            return Constraint.Skip

        #    return (
        #        sum(model.works[w.id, week, day, shift] for w in required_site_workers)
        #        >= required_number
        #    )

        #for shift in self.get_shifts_with_constraint("require_remote_site_presence"):
        #    site, required_number = shift.constraint_options[
        #        "require_remote_site_presence"
        #    ]
        #    # self.model.require_presence_at_site_overnight_rule = Constraint(
        #    setattr(
        #        self.model,
        #        f"require_remote_site_presence_{shift.name}",
        #        Constraint(
        #            [day for day in self.days],
        #            [week for week in self.weeks],
        #            [shift.name],
        #            [site],
        #            [required_number],
        #            rule=presenceAtRemoteSite,
        #        ),
        #    )

        # # Count the number of workers from each site on each night shift
        # # As 1 or 0 is optimum we can simply subtract 1 from the number
        # if self.constraint_options["balance_nights_across_sites"]:
        #    for site in self.sites:
        #        for shift in self.get_shifts_with_constraint("night"):
        #            block = shift.name
        #            for week in self.weeks:
        #                self.model.constraints.add(
        #                    self.model.night_per_site[week, block, site]
        #                    == sum(
        #                        self.model.shift_week_worker_assigned[
        #                            block, week, worker.id
        #                        ]
        #                        for worker in self.workers
        #                        if worker.site == site
        #                    )
        #                    # 1 >= sum(self.model.shift_week_worker_assigned[shift.name, week, worker.id] for worker in self.workers if worker.site == site)
        #                )
        #                self.model.constraints.add(
        #                    self.model.night_per_site_t1[week, block, site]
        #                    - self.model.night_per_site_t2[week, block, site]
        #                    == self.model.night_per_site[week, block, site] - 1
        #                )
        for site in track(
            self.sites, description="Generating site balance constraints..."
        ):
            for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint):
                if site in shift.sites:
                    block = shift.name
                    for week in self.weeks:
                        self.model.constraints.add(
                            self.model.shift_per_site[week, block, site]
                            == sum(
                                self.model.shift_week_worker_assigned[
                                    block, week, worker.id
                                ]
                                for worker in self.workers
                                if worker.site == site
                            )
                            # 1 >= sum(self.model.shift_week_worker_assigned[shift.name, week, worker.id] for worker in self.workers if worker.site == site)
                        )
                        self.model.constraints.add(
                            self.model.shift_per_site_t1[week, block, site]
                            - self.model.shift_per_site_t2[week, block, site]
                            == self.model.shift_per_site[week, block, site] - 1
                        )

        weeks_days = self.get_week_day_combinations()

        #for week in track(self.weeks, description="Generating week constraints..."):
        #    for shift_name in self.shifts_to_assign_or_force_as_blocks():

        for week, shift_name in self.week_shifts_to_assign_as_blocks():
                shift = self.get_shift_by_name(shift_name)
                sub_blocks = self.get_shift_sub_blocks(shift)

                for sb_idx, sb in enumerate(sub_blocks):
                    for worker in self.get_workers_for_shift(shift):
                        if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned:
                            try:
                                self.model.constraints.add(
                                    8
                                    * self.model.blocks_worker_shift_assigned[
                                        worker.id, week, shift_name, sb_idx
                                    ]
                                    >= sum(
                                        self.model.works[worker.id, week, day, shift_name]
                                        for day in sb
                                    )
                                )
                            except KeyError:
                                pass

                    self.model.constraints.add(
                        self.model.blocks_assigned[week, shift_name, sb_idx]
                        == sum(
                            self.model.blocks_worker_shift_assigned[
                                worker.id, week, shift_name, sb_idx
                            ]
                            for worker in self.workers
                            if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
                        )
                    )

                    workers_required = shift.get_worker_requirement_by_date(
                        self.get_week_start_date(week)
                    )

                    if shift.force_as_block_unless_nwd:
                        workers = self.get_workers_for_shift(shift)
                        # Get workers who have a nwd on the sub-block
                        nwd_workers = []
                        full_workers = []

                        for w in workers:
                            if w.non_working_day_list:
                                l = []
                                for (
                                    nwd,
                                    start_nwd_date,
                                    end_nwd_date,
                                ) in w.non_working_day_list:
                                    if nwd in sb:
                                        if start_nwd_date > self.get_week_start_date(week):
                                            continue
                                        if end_nwd_date < self.get_week_start_date(week):
                                            continue
                                        l.append(w)

                                if l:
                                    nwd_workers.append(w)
                                else:
                                    full_workers.append(w)
                            else:
                                full_workers.append(w)

                        if not nwd_workers:  # Just do the usual
                            self.model.constraints.add(
                                sum(
                                    self.model.blocks_worker_shift_assigned[
                                        worker.id, week, shift.name, sb_idx
                                    ]
                                    for worker in self.workers
                                    if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
                                )
                                <= workers_required
                            )
                        else:
                            self.model.constraints.add(
                                sum(
                                    self.model.blocks_worker_shift_assigned[
                                        worker.id, week, shift.name, sb_idx
                                    ]
                                    for worker in full_workers
                                    if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
                                )
                                <= workers_required + 1
                            )

                    elif shift.force_as_block:
                        if workers_required:
                            self.model.constraints.add(
                                sum(
                                    self.model.blocks_worker_shift_assigned[
                                        worker.id, week, shift.name, sb_idx
                                    ]
                                    for worker in self.workers
                                    if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
                                )
                                <= workers_required
                            )

        # Most of our constraints apply per worker
        # Worker constraint loop (worker loop)
        worker: Worker
        for worker in self.workers:
            console.rule(
                f"Generate worker constraints: [bold blue]{worker.name}[/bold blue]"
            )
            logger.debug(f"Generate worker constraints: {worker.name}")

            for week, day in self.get_week_day_combinations():
                shifts_today = self.get_shift_names_by_week_day(week, day)
                total_assigned = sum(
                    self.model.works[worker.id, week, day, shift]
                    for shift in shifts_today
                )
                # If assigned to any shift, works_day must be 1
                self.model.constraints.add(total_assigned >= self.model.works_day[worker.id, week, day])
                # If not assigned to any shift, works_day must be 0
                self.model.constraints.add(total_assigned <= 1000 * self.model.works_day[worker.id, week, day])

            # --- Max days worked per week constraint (global or per-worker) ---
            # Use per-worker value if set, otherwise use global constraint option if set
            max_days_worked_per_week = getattr(worker, "max_days_worked_per_week", None)
            if max_days_worked_per_week is None:
                max_days_worked_per_week = getattr(self.constraint_options_model, "max_days_worked_per_week", None)
            if max_days_worked_per_week is not None:
                for week in self.weeks:
                    # Create a binary variable for each day to indicate if the worker works any shift that day
                    for day in self.days:
                        var_name = f"works_any_{worker.id}_{week}_{day}"
                        if not hasattr(self.model, "works_any_day_vars"):
                            self.model.works_any_day_vars = {}
                        if (worker.id, week, day) not in self.model.works_any_day_vars:
                            self.model.works_any_day_vars[(worker.id, week, day)] = Var(within=Binary)
                            setattr(self.model, var_name, self.model.works_any_day_vars[(worker.id, week, day)])
                        works_any = self.model.works_any_day_vars[(worker.id, week, day)]
                        total_assigned = sum(
                            self.model.works[worker.id, week, day, shift]
                            for shift in self.get_shift_names_by_week_day(week, day)
                        )
                        # If assigned to any shift, works_any must be 1
                        self.model.constraints.add(total_assigned >= works_any)
                        # If not assigned to any shift, works_any must be 0
                        self.model.constraints.add(total_assigned <= 1000 * works_any)
                    # Now sum the binaries for the week
                    self.model.constraints.add(
                        sum(self.model.works_any_day_vars[(worker.id, week, day)] for day in self.days)
                        <= max_days_worked_per_week
                    )

            for week, day, shift_name in getattr(worker, "forced_assignments", []):
                # Only add if this shift is valid for this week/day
                if shift_name in self.get_shift_names_by_week_day(week, day, return_empty_if_week_day_not_found=True):
                    # We have already checked for leave conflicts when building the model
                    self.model.constraints.add(
                        self.model.works[worker.id, week, day, shift_name] == 1
                    )
                else:
                    date = self.get_date_by_week_day(week, day)
                    self.add_warning(
                        "Invalid forced assignment",
                        f"Worker '{worker.name}' has forced assignment for non-existent shift: {shift_name} on week {week}, day {day} (date: {date})",
                    )

            if hasattr(worker, "max_shifts_per_week_by_shift_name"):
                for shift_type, max_per_week in worker.max_shifts_per_week_by_shift_name.items():
                    if shift_type not in self.get_shift_names():
                        raise InvalidShift(
                            f"Worker '{worker.name}' has max_shifts_per_week_by_shift_name referencing non-existent shift: {shift_type}"
                        )
                    for week in self.weeks:
                        self.model.constraints.add(
                            sum(
                                self.model.works[worker.id, week, day, shift_type]
                                for day in days
                                if shift_type in self.get_shift_names_by_week_day(week, day)
                            ) <= max_per_week
                        )

            for week, day, shift in self.get_all_shiftclass_combinations():
                # Only apply if this shift has force_assign_with set
                if hasattr(shift, "force_assign_with") and shift.force_assign_with:
                    for other_shift_name in shift.force_assign_with:
                        # Only add constraint if the other shift is available on this day
                        if other_shift_name in self.get_shift_names_by_week_day(week, day):
                            for worker in self.get_workers_for_shift(shift):
                                self.model.constraints.add(
                                    self.model.works[worker.id, week, day, shift.name]
                                    <= self.model.works[worker.id, week, day, other_shift_name]
                                )

            # Add hard shift dependencies for this worker
            if hasattr(worker, "hard_day_dependencies"):
                for dep in worker.hard_day_dependencies:
                    day_group = dep.days

                    weeks_to_apply = self.weeks
                    if dep.weeks is not None:
                        weeks_to_apply = dep.weeks
                    if dep.start_date is not None or dep.end_date is not None:
                        if dep.start_date is None:
                            dep.start_date = self.start_date
                        if dep.end_date is None:
                            dep.end_date = self.rota_end_date

                        weeks_to_apply = self.get_weeks_by_date_range(
                            dep.start_date, dep.end_date
                        )

                    for week in weeks_to_apply:
                        # compute which days in the day_group should be ignored
                        ignored_days = set()
                        try:
                            for ignore_date in getattr(dep, "ignore_dates", []) or []:
                                try:
                                    ww, dd = self.date_week_day_map.get(ignore_date, (None, None))
                                except Exception:
                                    ww = dd = None
                                if ww == week and dd in day_group:
                                    ignored_days.add(dd)
                        except Exception:
                            ignored_days = set()
                        # If the worker has any forced assignment in the dependency group,
                        # emit a concise warning. For visibility include every day in the
                        # group showing the worker's forced shifts (if any) and other
                        # workers' forced shifts on those days so blocking days are visible.
                        worker_has_forced_in_group = False
                        for (fw, fd, fs) in getattr(worker, "forced_assignments", []):
                            if fw == week and fd in day_group:
                                worker_has_forced_in_group = True
                                break
                        if not worker_has_forced_in_group:
                            try:
                                for (d_date, d_shift) in getattr(worker, "forced_assignments_by_date", []):
                                    ww, dd = self.date_week_day_map.get(d_date, (None, None))
                                    if ww == week and dd in day_group:
                                        worker_has_forced_in_group = True
                                        break
                            except Exception:
                                pass

                        if worker_has_forced_in_group:
                            forced_day_details = []
                            try:
                                date_map_for_week = {self.get_date_by_week_day(week, d): d for d in day_group}
                            except Exception:
                                date_map_for_week = {}

                            for day in day_group:
                                if day in ignored_days:
                                    # skip ignored dates for warning/visibility
                                    continue
                                worker_forced_shifts = set()
                                # week-based forced assignments
                                for (fw, fd, fs) in getattr(worker, "forced_assignments", []):
                                    if fw == week and fd == day:
                                        worker_forced_shifts.add(fs)
                                # date-based forced assignments
                                try:
                                    date = self.get_date_by_week_day(week, day)
                                    for (d_date, d_shift) in getattr(worker, "forced_assignments_by_date", []):
                                        try:
                                            ww, dd = self.date_week_day_map.get(d_date, (None, None))
                                        except Exception:
                                            ww = dd = None
                                        if ww == week and dd == day:
                                            worker_forced_shifts.add(d_shift)
                                        else:
                                            try:
                                                if str(d_date) == str(date):
                                                    worker_forced_shifts.add(d_shift)
                                            except Exception:
                                                pass
                                except Exception:
                                    pass

                                # collect other workers' forced shifts on this day
                                others_forced = {}
                                try:
                                    date_for_day = self.get_date_by_week_day(week, day)
                                except Exception:
                                    date_for_day = None

                                for other in self.workers:
                                    if other is worker:
                                        continue
                                    other_shifts = set()
                                    # week-based forced assignments
                                    for (ofw, ofd, ofs) in getattr(other, "forced_assignments", []):
                                        if ofw == week and ofd == day:
                                            other_shifts.add(ofs)
                                    # date-based forced assignments: try mapping then fallback
                                    for (odate, oshift) in getattr(other, "forced_assignments_by_date", []):
                                        try:
                                            oww, odd = self.date_week_day_map.get(odate, (None, None))
                                        except Exception:
                                            oww = odd = None
                                        if oww == week and odd == day:
                                            other_shifts.add(oshift)
                                            continue
                                        try:
                                            if date_for_day is not None and str(odate) == str(date_for_day):
                                                other_shifts.add(oshift)
                                        except Exception:
                                            pass

                                    if other_shifts:
                                        others_forced[other.name] = sorted(list(other_shifts))

                                forced_day_details.append((day, sorted(list(worker_forced_shifts)), others_forced))

                            parts = []
                            for d, wfs, others in forced_day_details:
                                olist = ", ".join([f"{n}: {', '.join(s)}" for n, s in others.items()]) or "none"
                                wfs_str = ", ".join(wfs) or "none"
                                parts.append(f"{d} - worker forced: {wfs_str}; others forced: {olist}")
                            details = "; ".join(parts)
                            logger.warning(f"HardDayDependency forced-assignment warning for worker={worker.name} week={week}: {details}")
                            self.add_warning(
                                "HardDayDependency forced assignment",
                                f"Worker '{worker.name}' has hard_day_dependency {set(day_group)} in week {week}; forced assignments on dependency days: {details}. If the rota is not infeasible this is unlikely to be an issue.",
                            )
                        # Build constraints only for non-ignored days
                        effective_days = [d for d in day_group if d not in ignored_days]
                        if not effective_days:
                            # nothing to constrain for this dependency in this week
                            continue

                        assigned_any_shift = {}
                        for day in effective_days:
                            shifts_today = self.get_shift_names_by_week_day(week, day)
                            var_name = f"hard_day_dep_{worker.id}_{week}_{day}"
                            if not hasattr(self.model, "hard_day_dep_vars"):
                                self.model.hard_day_dep_vars = {}
                            if (worker.id, week, day) not in self.model.hard_day_dep_vars:
                                self.model.hard_day_dep_vars[(worker.id, week, day)] = Var(within=Binary)
                                setattr(self.model, var_name, self.model.hard_day_dep_vars[(worker.id, week, day)])
                            assigned_any_shift[day] = self.model.hard_day_dep_vars[(worker.id, week, day)]
                            total_assigned = sum(self.model.works[worker.id, week, day, shift] for shift in shifts_today)
                            self.model.constraints.add(total_assigned >= assigned_any_shift[day])
                            self.model.constraints.add(total_assigned <= len(shifts_today) * assigned_any_shift[day])
                        # All-or-none: all effective (non-ignored) days in the group must have the same assignment status
                        vals = list(assigned_any_shift.values())
                        for i in range(1, len(vals)):
                            self.model.constraints.add(vals[i] == vals[0])

            # Add hard day exclusions for this worker
            if hasattr(worker, "hard_day_exclusions"):
                for exclusion in worker.hard_day_exclusions:
                    days_to_exclude = exclusion.days
                    weeks_to_apply = self.weeks
                    if exclusion.weeks is not None:
                        weeks_to_apply = exclusion.weeks
                    if exclusion.start_date is not None or exclusion.end_date is not None:
                        if exclusion.start_date is None:
                            exclusion.start_date = self.start_date
                        if exclusion.end_date is None:
                            exclusion.end_date = self.rota_end_date

                        weeks_to_apply = self.get_weeks_by_date_range(
                            exclusion.start_date, exclusion.end_date
                        )

                    for week in weeks_to_apply:
                        # For each week, prevent any shift being assigned on both days
                        self.model.constraints.add(
                            sum(
                                self.model.works_day[worker.id, week, day]
                                for day in days_to_exclude
                                #for shift in self.get_shift_names_by_week_day(week, day)
                            )
                            <= 1
                        )
                        #self.model.constraints.add(
                        #    sum(
                        #        self.model.works[worker.id, week, day1, shift]
                        #        for shift in self.get_shift_names_by_week_day(week, day1)
                        #    )
                        #    + sum(
                        #        self.model.works[worker.id, week, day2, shift]
                        #        for shift in self.get_shift_names_by_week_day(week, day2)
                        #    )
                        #    <= 1
                        #)


            for week, day in self.get_week_day_combinations():
                shifts_today = self.get_shift_names_by_week_day(week, day)
                worker_allowed_sets = getattr(worker, "allowed_multi_shift_sets", [])
                # for idx, allowed_set in enumerate(worker_allowed_sets):
                #    if allowed_set.issubset(shifts_today):
                #        var = self.model.multi_shift_together[worker.id, week, day, idx]
                #        self.model.constraints.add(
                #            sum(self.model.works[worker.id, week, day, shift] for shift in allowed_set)
                #            == len(allowed_set) * var
                #        )
                for idx, allowed_set in enumerate(worker_allowed_sets):
                    if allowed_set.issubset(shifts_today):
                        var = self.model.multi_shift_together[worker.id, week, day, idx]
                        # var is 1 iff all shifts in allowed_set are assigned
                        for shift in allowed_set:
                            self.model.constraints.add(
                                var <= self.model.works[worker.id, week, day, shift]
                            )
                        self.model.constraints.add(
                            var
                            >= sum(
                                self.model.works[worker.id, week, day, shift]
                                for shift in allowed_set
                            )
                            - len(allowed_set)
                            + 1
                        )
                    shifts_today = self.get_shift_names_by_week_day(week, day)
                for shift_name in shifts_today:
                    if hasattr(worker, "force_assign_with") and shift_name in worker.force_assign_with:
                        for other_shift in worker.force_assign_with[shift_name]:
                            if other_shift in shifts_today:
                                self.model.constraints.add(
                                    self.model.works[worker.id, week, day, shift_name]
                                    <= self.model.works[worker.id, week, day, other_shift]
                                )



            if self.get_workers_who_require_locums():
                for week, day, shift in self.get_all_shiftclass_combinations():
                    if not worker.locum:
                        self.model.constraints.add(
                            self.get_locum_availability(worker, week, day, shift)
                            >= self.model.locum_works[worker.id, week, day, shift.name]
                        )

                        self.model.constraints.add(
                            self.model.locum_shifts_by_worker[worker.id]
                            == sum(
                                self.model.locum_works[worker.id, week, day, shift]
                                for week, day, shift in self.get_all_shiftname_combinations()
                            )
                        )

                self.model.constraints.add(
                    self.model.locum_shifts_t1[worker.id]
                    - self.model.locum_shifts_t2[worker.id]
                    == self.model.locum_shifts_by_worker[worker.id]
                )

                if worker.locum_max_shifts:
                    self.model.constraints.add(
                        worker.locum_max_shifts
                        >= sum(
                            self.model.locum_works[worker.id, week, day, shiftname]
                            for week, day, shiftname in self.get_all_shiftname_combinations()
                        )
                    )

                if worker.locum_max_shifts_per_week:
                    for week_to_check in self.weeks:
                        self.model.constraints.add(
                            worker.locum_max_shifts_per_week
                            >= sum(
                                self.model.locum_works[worker.id, week, day, shiftname]
                                for week, day, shiftname in self.get_all_shiftname_combinations(
                                    week=week_to_check
                                )
                            )
                        )

                if self.constraint_options_model.distribute_locum_shifts:
                    self.model.constraints.add(
                        sum(
                            self.model.locum_required[week, day, shift]
                            for week, day, shift in self.get_all_shiftname_combinations()
                        )
                        / (len(self.get_locum_workers()) - 1)
                        >= self.model.locum_shifts_by_worker[worker.id]
                    )

                if (
                    self.constraint_options_model.maximum_allowed_locum_shifts_per_worker
                    is not None
                ):
                    self.model.constraints.add(
                        self.constraint_options_model.maximum_allowed_locum_shifts_per_worker
                        >= self.model.locum_shifts_by_worker[worker.id]
                    )
                if worker.locum_availability:
                    if (
                        self.constraint_options_model.minimum_allowed_locum_shifts_per_worker
                        is not None
                    ):
                        self.model.constraints.add(
                            self.constraint_options_model.minimum_allowed_locum_shifts_per_worker
                            <= self.model.locum_shifts_by_worker[worker.id]
                        )

            if (
                self.constraint_options_model.minimise_shift_diffs
                or self.constraint_options_model.maximum_allowed_shift_diff is not None
            ):
                self.model.constraints.add(
                    self.model.shift_count_diff_summed[worker.id]
                    == sum(
                        self.model.shift_count_diff[worker.id, shift.name]
                        for shift in self.get_shifts()
                    )
                )

            if self.constraint_options_model.maximum_allowed_shift_diff is not None:
                self.model.constraints.add(
                    self.model.shift_count_diff_summed[worker.id]
                    <= self.constraint_options_model.maximum_allowed_shift_diff
                )
                self.model.constraints.add(
                    self.model.shift_count_diff_summed[worker.id]
                    >= -self.constraint_options_model.maximum_allowed_shift_diff
                )

            for max_days, weeks in self.constraint_options_model.max_days_per_week_block:
                for week_blocks in self.get_week_block_iterator(weeks):
                    self.model.constraints.add(
                        sum(
                            self.model.works_day[worker.id, week, day]
                            for week, day in self.get_week_day_combinations()
                            if week in week_blocks
                        )
                        <= max_days
                    )

            for constraint in worker.max_days_per_week_block:
                for week_blocks in self.get_week_block_iterator(constraint.week_block):
                    # Determine which (week, day) combinations to ignore
                    ignored_week_days = set()
                    for ignore_date in (constraint.ignore_dates or []):
                        try:
                            w, d = self.date_week_day_map.get(ignore_date, (None, None))
                            if w is not None and d is not None:
                                ignored_week_days.add((w, d))
                        except Exception:
                            pass
                    
                    included_days = [
                        (week, day)
                        for week, day in self.get_week_day_combinations()
                        if week in week_blocks and (week, day) not in ignored_week_days
                    ]
                    if included_days:
                        self.model.constraints.add(
                            sum(
                                self.model.works_day[worker.id, week, day]
                                for week, day in included_days
                            )
                            <= constraint.max_days
                        )

            for constraint in worker.max_unique_shifts_per_week_block:
                for week_blocks in self.get_week_block_iterator(constraint.week_block):
                    if not week_blocks:
                        continue

                    # Determine which (week, day) combinations to ignore
                    ignored_week_days = set()
                    for ignore_date in (constraint.ignore_dates or []):
                        try:
                            w, d = self.date_week_day_map.get(ignore_date, (None, None))
                            if w is not None and d is not None:
                                ignored_week_days.add((w, d))
                        except Exception:
                            pass
                    
                    block_start_week = week_blocks[0]
                    used_shift_vars = []
                    for shift_name in self.get_shift_names():
                        shift_assignments = []
                        for week, day in self.get_week_day_combinations():
                            if week in week_blocks and (week, day) not in ignored_week_days:
                                if shift_name in self.get_shift_names_by_week_day(
                                    week,
                                    day,
                                    return_empty_if_week_day_not_found=True,
                                ):
                                    shift_assignments.append(
                                        self.model.works[worker.id, week, day, shift_name]
                                    )
                        if not shift_assignments:
                            continue

                        used_shift_var = self.model.worker_shift_used_in_week_block[
                            worker.id,
                            constraint.week_block,
                            block_start_week,
                            shift_name,
                        ]
                        used_shift_vars.append(used_shift_var)
                        self.model.constraints.add(
                            sum(shift_assignments)
                            <= len(shift_assignments) * used_shift_var
                        )

                    if used_shift_vars:
                        self.model.constraints.add(
                            sum(used_shift_vars) <= constraint.max_unique_shifts
                        )

            for week_blocks in self.get_week_block_iterator(4):
                # Prevent more than n number shifts per 4 weeks
                try:
                    self.model.constraints.add(
                        self.constraint_options_model.max_shifts_per_month
                        >= sum(
                            self.model.works[worker.id, week, day, shiftname]
                            for week, day, shiftname in self.get_all_shiftname_combinations()
                            if week in week_blocks
                        )
                    )
                except ValueError:
                    # Occurs if there are no shifts within a defined block
                    # TODO: test if this breaks (and we should check rathar than except)
                    self.add_warning(
                        "max shifts per month constraint",
                        f"Failed to constrain max_shifts_per_month for worker: {worker.name}",
                    )
                    pass

            # Date-based per-worker per-shift avoidance rules
            for constraint_dict in self.constraint_options_model.avoid_shifts_by_worker_dates:
                # validate shift names
                if invalid_shifts := set(constraint_dict["shifts"]).difference(set(self.get_shift_names())):
                    message = f"avoid shifts by worker/dates constraint contains a non existent shift: {invalid_shifts}"
                    if self.ignore_valid_shifts:
                        self.add_warning("Non existent shift", message)
                        continue
                    else:
                        raise InvalidShift(message)

                if worker.name not in constraint_dict["names"]:
                    continue

                # Build set of allowed (week,day) pairs that fall within the date selector
                selected_weeks_days = set()

                # explicit dates take precedence
                if constraint_dict.get("dates"):
                    date_set = set(constraint_dict["dates"])
                    for wk, d in self.get_week_day_combinations():
                        date = self.get_date_by_week_day(wk, d)
                        if date in date_set:
                            selected_weeks_days.add((wk, d))
                else:
                    sd = constraint_dict.get("start_date")
                    ed = constraint_dict.get("end_date")
                    # If neither provided, apply to whole rota
                    if sd is None and ed is None:
                        for wk, d in self.get_week_day_combinations():
                            selected_weeks_days.add((wk, d))
                    else:
                        if sd is None:
                            sd = self.start_date
                        if ed is None:
                            ed = self.rota_end_date
                        for wk, d in self.get_week_day_combinations():
                            date = self.get_date_by_week_day(wk, d)
                            if sd <= date <= ed:
                                selected_weeks_days.add((wk, d))

                if not selected_weeks_days:
                    # nothing to constrain
                    continue

                # Determine which of the requested shifts actually occur on the selected dates
                present_shifts = set(
                    shift
                    for wk, d, shift in self.get_all_shiftname_combinations()
                    if (wk, d) in selected_weeks_days
                )

                requested_shifts = set(constraint_dict["shifts"])
                used_shifts = requested_shifts.intersection(present_shifts)

                missing_shifts = requested_shifts - present_shifts
                if missing_shifts:
                    self.add_warning(
                        "avoid shifts by worker/dates",
                        (
                            f"Requested avoid shifts {sorted(list(missing_shifts))} for worker '{worker.name}'"
                            " are not present in the rota dates selected and will be ignored"
                        ),
                    )

                if not used_shifts:
                    # Nothing to constrain after filtering — skip adding the constraint
                    continue

                # Add constraint that forbids the given shifts for this worker on the selected dates
                self.model.constraints.add(
                    0
                    == sum(
                        self.model.works[worker.id, wk, d, shift]
                        for wk, d, shift in self.get_all_shiftname_combinations()
                        if (wk, d) in selected_weeks_days and shift in used_shifts
                    )
                )

            for constraint_dict in self.constraint_options_model.avoid_shifts_by_worker_names:
                if invalid_shifts := set(constraint_dict["shifts"]).difference(
                    set(self.get_shift_names())
                ):
                    message = f"avoid shifts by worker constraint contains a non existent shift: {invalid_shifts}"
                    if self.ignore_valid_shifts:
                        self.add_warning("Non existent shift", message)
                        continue  # Skip if we don't want to raise an error
                    else:
                        raise InvalidShift(message)

                if worker.name in constraint_dict["names"]:
                    self.model.constraints.add(
                        0
                        == sum(
                            self.model.works[worker.id, week, day, shift]
                            for week, day, shift in self.get_all_shiftname_combinations()
                            if week in constraint_dict["weeks"]
                            and shift in constraint_dict["shifts"]
                        )
                    )

            for constraint_dict in self.constraint_options_model.avoid_shifts_by_grades:
                if invalid_shifts := set(constraint_dict["shifts"]).difference(
                    set(self.get_shift_names())
                ):
                    message = f"avoid shifts by grade constraint contains a non existent shift: {invalid_shifts}"
                    if self.ignore_valid_shifts:
                        self.add_warning("Non existent shift", message)
                        continue  # Skip if we don't want to raise an error
                    else:
                        raise InvalidShift(message)
                if invalid_grades := set(constraint_dict["grades"]).difference(
                    set(self.get_worker_grades())
                ):
                    message = f"avoid shifts by grade constraint contains a non existent worker grade: {invalid_grades}"
                    self.add_warning("Non existent shift", message)
                    raise ValueError(message)

                if worker.grade in constraint_dict["grades"]:
                    self.model.constraints.add(
                        0
                        == sum(
                            self.model.works[worker.id, week, day, shift]
                            for week, day, shift in self.get_all_shiftname_combinations()
                            if week in constraint_dict["weeks"]
                            and shift in constraint_dict["shifts"]
                        )
                    )

            # Count number of weekends an worker works
            # if self.constraint_options["balance_weekends"]:
            self.model.constraints.add(
                self.model.worker_weekend_count[worker.id]
                == sum(self.model.works_weekend[worker.id, week] for week in self.weeks)
            )

            if self.constraint_options_model.max_weekends > -1:
                self.model.constraints.add(
                    self.constraint_options_model.max_weekends
                    >= self.model.worker_weekend_count[worker.id]
                )

            # Balance shifts within required limits (set by balance_offset)
            # If balance_offset is too restrictive (for the rota length)
            # no solution will be possible
            for shift in track(
                self.get_shifts(), description="Generate shift balance constraints"
            ):
                if (
                    self.is_worker_eligible_for_shift(worker, shift)
                ):  # Each site specfies which sites self.workers can fullfill
                    total_shifts = self.shift_worker_counts[shift.name]

                    # Find workers who have exact shifts for this shift type
                    exact_workers = [
                        w for w in self.workers 
                        if self.is_worker_eligible_for_shift(w, shift) and shift.name in getattr(w, "exact_shifts", {})
                    ]
                    exact_shifts_sum = sum(w.exact_shifts[shift.name] for w in exact_workers)

                    if exact_shifts_sum > total_shifts:
                        warning_msg = f"For shift {shift.name}, total exact shifts ({exact_shifts_sum}) exceeds total required shifts ({total_shifts})."
                        if not any(warn[1] == warning_msg for warn in self.warnings):
                            self.add_warning("Exact shifts exceed total shifts", warning_msg)

                    if worker in exact_workers:
                        target_shifts = worker.exact_shifts[shift.name]
                    else:
                        remaining_shifts = total_shifts - exact_shifts_sum
                        # Sum FTE only for workers who do not have exact shifts for this shift type
                        full_time_equivalent_joined = sum(
                            w.get_fte(shift=shift.name)
                            for w in self.workers
                            if self.is_worker_eligible_for_shift(w, shift) and w not in exact_workers
                        )

                        if not full_time_equivalent_joined:
                            # avoid ZeroDivisionError
                            if remaining_shifts > 0:
                                self.add_warning(
                                    "Zero FTE for sites with remaining shifts",
                                    f"full_time_equivalent sum is zero for shift {shift.name} but remaining_shifts = {remaining_shifts}; setting target_shifts=0 for worker {worker.name}",
                                )
                            target_shifts = 0
                        else:
                            target_shifts = (
                                max(0.0, float(remaining_shifts))
                                / full_time_equivalent_joined
                                * worker.get_fte(shift=shift.name)
                            )

                    if worker not in exact_workers:
                        if self.use_previous_shifts:
                            if shift.name in worker.previous_shifts:
                                entry = worker.previous_shifts[shift.name]
                                # entry may be a tuple (worked, allocated) or a dict of per-day tuples
                                if isinstance(entry, dict):
                                    # prefer overall tuple if present
                                    if "__all__" in entry:
                                        try:
                                            worked = float(entry["__all__"][0])
                                            allocated = float(entry["__all__"][1])
                                            target_shifts = target_shifts + allocated - worked
                                        except Exception:
                                            pass
                                    else:
                                        # sum any per-day entries
                                        try:
                                            total_worked = 0.0
                                            total_alloc = 0.0
                                            for v in entry.values():
                                                total_worked += float(v[0])
                                                total_alloc += float(v[1])
                                            target_shifts = target_shifts + total_alloc - total_worked
                                        except Exception:
                                            pass
                                else:
                                    try:
                                        worked, allocated = entry
                                        target_shifts = (
                                            target_shifts + float(allocated) - float(worked)
                                        )
                                    except Exception:
                                        pass

                        if self.use_shift_balance_extra:
                            if shift.name in worker.shift_balance_extra:
                                extra = worker.shift_balance_extra[shift.name]

                                # TODO look at how this affects allocation (how does it affect fte)
                                match extra:
                                    case "double":
                                        target_shifts = target_shifts * 2
                                    case "half":
                                        target_shifts = target_shifts / 2
                                    case _:
                                        target_shifts = target_shifts + extra

                    # print(worker.name, shift.name, target_shifts)
                    worker.shift_target_number[shift.name] = target_shifts

                    if worker in exact_workers:
                        self.model.constraints.add(
                            self.model.shift_count[worker.id, shift.name] == target_shifts
                        )
                    elif shift.hard_constrain_shift:
                        adjusted_balance_offset = shift.balance_offset
                        if worker.get_fte(shift=shift.name) < 100:
                            adjusted_balance_offset = (
                                adjusted_balance_offset * self.ltft_balance_offset
                            )

                        max_shifts = (
                            target_shifts
                            + adjusted_balance_offset * self.balance_offset_modifier
                        )
                        min_shifts = (
                            target_shifts
                            - adjusted_balance_offset * self.balance_offset_modifier
                        )

                        # --- Store for export ---
                        if not hasattr(worker, "hard_constrain_shift_limits"):
                            worker.hard_constrain_shift_limits = {}
                        worker.hard_constrain_shift_limits[shift.name] = {
                            "min_shifts": min_shifts,
                            "max_shifts": max_shifts,
                            "target_shifts": target_shifts,
                        }

                        if self.get_week_day_combinations_for_shift(shift):
                            self.model.constraints.add(
                                inequality(
                                    min_shifts,
                                    sum(
                                        self.model.works[worker.id, week, day, shift.name]
                                        for week, day in self.get_week_day_combinations_for_shift(
                                            shift
                                        )
                                    ),
                                    max_shifts,
                                )
                            )
                else:
                    # Make sure shifts aren't assigned to those from other sites
                    # Not needed if there are no shifts
                    if self.get_week_day_combinations_for_shift(shift):
                        self.model.constraints.add(
                            0
                            == sum(
                                self.model.works[worker.id, week, day, shift.name]
                                for week, day in self.get_week_day_combinations_for_shift(
                                    shift
                                )
                                # for week, day in self.get_week_day_combinations()
                            )
                        )

                # model.shift_count stores the number of shifts that a worker is assigned to worker by shift
                # this is used (by the objective) to balance the total number of shifts worked
                self.model.constraints.add(
                    self.model.shift_count[worker.id, shift.name]
                    == sum(
                        self.model.works[worker.id, week, day, shift.name]
                        for week, day in self.get_week_day_combinations_for_shift(shift)
                    )
                )

                self.model.constraints.add(
                    self.model.shift_count_diff[worker.id, shift.name]
                    == self.model.shift_count[worker.id, shift.name]
                    - worker.shift_target_number[shift.name]
                )

                if self.constraint_options_model.balance_shifts:
                    self.model.constraints.add(
                        self.model.shift_count_t1[worker.id, shift.name]
                        - self.model.shift_count_t2[worker.id, shift.name]
                        == self.model.shift_count_diff[worker.id, shift.name]
                    )

                if self.constraint_options_model.balance_shifts_quadratic:
                    # This may need to be updated
                    xU = 25
                    xL = 1
                    self.model.constraints.add(
                        inequality(
                            xL,
                            self.model.shift_count_t1[worker.id, shift.name]
                            + self.model.shift_count_t2[worker.id, shift.name]
                            + 1,
                            xU,
                        )
                    )

                    self.model.constraints.add(
                        self.model.shift_count_w[worker.id, shift.name]
                        >= xL
                        * (
                            self.model.shift_count_t1[worker.id, shift.name]
                            + self.model.shift_count_t2[worker.id, shift.name]
                            + 1
                        )
                        * 2
                        - xL * xL
                    )

                    self.model.constraints.add(
                        self.model.shift_count_w[worker.id, shift.name]
                        >= xU
                        * (
                            self.model.shift_count_t1[worker.id, shift.name]
                            + self.model.shift_count_t2[worker.id, shift.name]
                            + 1
                        )
                        * 2
                        - xU * xU
                    )

            # Define worker_shift_count_t1 and worker_shift_count_t2 constraints for the object
            # Thus bypassing the need for a quadratic solver
            # t1-t2 is the target
            # As the objective is to minimise t1+t2 and t1 and t2 are positive reals
            # t1+t2 approximates the absolute target (which otherwise requires a quadratic solver)
            # TODO: quadratic implementation so perfect solutions will be chosen

            # NOTE: as this balances across the sum of the worker's shifts
            # they can cancel each other out if a high weighting is given to multiple shifts
            if self.constraint_options_model.balance_shifts_over_workers:
                self.model.constraints.add(
                    self.model.worker_shift_count_t1[worker.id]
                    - self.model.worker_shift_count_t2[worker.id]
                    == sum(
                        (self.model.shift_count_diff[worker.id, shift.name])
                        * shift.balance_weighting
                        for shift in self.get_shifts()
                    )
                )

            if self.constraint_options_model.balance_bank_holidays:
                extra_bank_holiday = 0
                if self.use_bank_holiday_extra:
                    extra_bank_holiday = worker.bank_holiday_extra

                self.model.constraints.add(
                    self.model.bank_holiday_count[worker.id]
                    == sum(
                        self.model.works[worker.id, week, day, shift]
                        for week, day, shift in self.get_all_shiftname_combinations()
                        if self.week_day_date_map[(week, day)] in self.bank_holidays
                    )
                    + extra_bank_holiday
                )

                xU = len(self.get_bank_holiday_week_days()) + 1
                xL = 1
                self.model.constraints.add(
                    inequality(
                        xL,
                        self.model.bank_holiday_count[worker.id] + 1,
                        xU,
                    )
                )

                self.model.constraints.add(
                    self.model.bank_holiday_count_w[worker.id]
                    >= xL * (self.model.bank_holiday_count[worker.id] + 1) * 2 - xL * xL
                )

                self.model.constraints.add(
                    self.model.bank_holiday_count_w[worker.id]
                    >= xU * (self.model.bank_holiday_count[worker.id] + 1) * 2 - xU * xU
                )

            if self.constraint_options_model.balance_nights:
                self.model.constraints.add(
                    self.model.night_shift_count[worker.id]
                    == sum(
                        self.model.works[worker.id, week, day, shift.name]
                        for week, day, shift in self.get_week_day_shift_combinations_for_constraint(
                            NightConstraint
                        )
                        # for week, day in self.get_week_day_combinations_for_shift(shift)
                        # for shift in self.get_shifts_with_constraint("night")
                    )
                )

                night_shift_target_number = sum(
                    worker.shift_target_number[shift.name]
                    for shift in self.get_shifts_with_constraint(NightConstraint)
                )

                self.model.constraints.add(
                    self.model.night_shift_count_t1[worker.id]
                    - self.model.night_shift_count_t2[worker.id]
                    == self.model.night_shift_count[worker.id]
                    - night_shift_target_number
                )

                # This may need to be updated
                xU = 6
                xL = 1
                self.model.constraints.add(
                    inequality(
                        xL,
                        self.model.night_shift_count_t1[worker.id]
                        + self.model.night_shift_count_t2[worker.id]
                        + 1,
                        xU,
                    )
                )

                self.model.constraints.add(
                    self.model.night_shift_count_w[worker.id]
                    >= xL
                    * (
                        self.model.night_shift_count_t1[worker.id]
                        + self.model.night_shift_count_t2[worker.id]
                        + 1
                    )
                    * 2
                    - xL * xL
                )

                self.model.constraints.add(
                    self.model.night_shift_count_w[worker.id]
                    >= xU
                    * (
                        self.model.night_shift_count_t1[worker.id]
                        + self.model.night_shift_count_t2[worker.id]
                        + 1
                    )
                    * 2
                    - xU * xU
                )

                # self.model.constraints.add(
                #     self.model.night_shift_count_w[worker.id] >= 0)

            # We use a similar method to balance the number of weekends worked
            # This works as long as weekend shifts are assigned as blocks!
            weekend_shift_target_number = sum(
                worker.shift_target_number[shift.name] / len(shift.days)
                for shift in self.get_shifts()
                if "Sat" in shift.days or "Sun" in shift.days
            )

            worker.weekend_shift_target_number = weekend_shift_target_number

            if self.constraint_options_model.balance_weekends:
                if weekend_shift_target_number > 0:
                    self.model.constraints.add(
                        self.model.weekend_shift_count_t1[worker.id]
                        - self.model.weekend_shift_count_t2[worker.id]
                        == self.model.worker_weekend_count[worker.id]
                        - weekend_shift_target_number
                    )

                    xU = self.constraint_options_model.max_weekends

                    xL = 1
                    self.model.constraints.add(
                        inequality(
                            xL,
                            self.model.weekend_shift_count_t1[worker.id]
                            + self.model.weekend_shift_count_t2[worker.id]
                            + 1,
                            xU,
                        )
                    )

                    self.model.constraints.add(
                        self.model.weekend_shift_count_w[worker.id]
                        >= xL
                        * (
                            self.model.weekend_shift_count_t1[worker.id]
                            + self.model.weekend_shift_count_t2[worker.id]
                            + 1
                        )
                        * 2
                        - xL * xL
                    )

                    self.model.constraints.add(
                        self.model.weekend_shift_count_w[worker.id]
                        >= xU
                        * (
                            self.model.weekend_shift_count_t1[worker.id]
                            + self.model.weekend_shift_count_t2[worker.id]
                            + 1
                        )
                        * 2
                        - xU * xU
                    )

            # Per-day weekend balancing (Sat and Sun separately)
            logger.debug(f"Adding per-day weekend balancing constraints for worker {worker.name} ({worker.id})")
            logger.debug(self.constraint_options_model.weekend_day_hard_max_deviation)
            if (
                self.constraint_options_model.balance_weekend_days
                or self.constraint_options_model.balance_weekend_days_quadratic
                or self.constraint_options_model.weekend_day_hard_max_deviation is not None
            ):
                logger.debug(f"Worker {worker.name}: per-day weekend block active (hard_dev={self.constraint_options_model.weekend_day_hard_max_deviation})")
                console.print(f"Worker {worker.name}: per-day weekend block active (hard_dev={self.constraint_options_model.weekend_day_hard_max_deviation})")
                hard_weekend_constraints_added = 0
                for day_name in ("Sat", "Sun"):
                    # count assignments on this specific weekend day
                    self.model.constraints.add(
                        self.model.weekend_day_shift_count[worker.id, day_name]
                        == sum(
                            self.model.works[worker.id, w, d, shift]
                            for w, d, shift in self.get_all_shiftname_combinations()
                            if d == day_name
                        )
                    )

                    # target number for this day (derived from shift targets)
                    weekend_day_shift_target_number = sum(
                        worker.shift_target_number[shift.name] / len(shift.days)
                        for shift in self.get_shifts()
                        if day_name in shift.days
                    )

                    # Optionally adjust per-day weekend targets using previous_shifts
                    if self.use_previous_shifts and self.constraint_options_model.balance_weekend_days_use_previous_shifts:
                        # For each shift that includes this day, if previous_shifts contains an entry
                        # add (allocated - worked) / len(shift.days) to the per-day target
                        prev_adj = 0.0
                        for shift in self.get_shifts():
                            if day_name in shift.days and shift.name in getattr(worker, "previous_shifts", {}):
                                entry = worker.previous_shifts.get(shift.name)
                                # support per-day nested dicts: { 'Sat': (worked, allocated), 'Sun': (...) }
                                if isinstance(entry, dict):
                                    # prefer explicit day entry if present
                                    if day_name in entry:
                                        try:
                                            worked = float(entry[day_name][0])
                                            allocated = float(entry[day_name][1])
                                            prev_adj += (allocated - worked)
                                        except Exception:
                                            continue
                                    # fallback: if an overall tuple was preserved under '__all__', use distributed value
                                    elif "__all__" in entry:
                                        try:
                                            worked = float(entry["__all__"][0])
                                            allocated = float(entry["__all__"][1])
                                            prev_adj += (allocated - worked) / max(1, len(shift.days))
                                        except Exception:
                                            continue
                                else:
                                    try:
                                        worked = float(entry[0])
                                        allocated = float(entry[1])
                                        prev_adj += (allocated - worked) / max(1, len(shift.days))
                                    except Exception:
                                        continue
                        weekend_day_shift_target_number = weekend_day_shift_target_number + prev_adj
                    # store adjusted per-day weekend target on the worker object for later use in objectives
                    if not hasattr(worker, 'weekend_day_target'):
                        worker.weekend_day_target = {}
                    worker.weekend_day_target[day_name] = weekend_day_shift_target_number
                    self.model.constraints.add(
                        self.model.weekend_day_shift_count_t1[worker.id, day_name]
                        - self.model.weekend_day_shift_count_t2[worker.id, day_name]
                        == self.model.weekend_day_shift_count[worker.id, day_name]
                        - weekend_day_shift_target_number
                    )

                    xU_day = self.constraint_options_model.max_weekends
                    xL_day = 1
                    self.model.constraints.add(
                        inequality(
                            xL_day,
                            self.model.weekend_day_shift_count_t1[worker.id, day_name]
                            + self.model.weekend_day_shift_count_t2[worker.id, day_name]
                            + 1,
                            xU_day,
                        )
                    )

                    self.model.constraints.add(
                        self.model.weekend_day_shift_count_w[worker.id, day_name]
                        >= xL_day
                        * (
                            self.model.weekend_day_shift_count_t1[worker.id, day_name]
                            + self.model.weekend_day_shift_count_t2[worker.id, day_name]
                            + 1
                        )
                        * 2
                        - xL_day * xL_day
                    )

                    self.model.constraints.add(
                        self.model.weekend_day_shift_count_w[worker.id, day_name]
                        >= xU_day
                        * (
                            self.model.weekend_day_shift_count_t1[worker.id, day_name]
                            + self.model.weekend_day_shift_count_t2[worker.id, day_name]
                            + 1
                        )
                        * 2
                        - xU_day * xU_day
                    )

                    # Optional hard constraint: limit absolute deviation from per-day target
                    if self.constraint_options_model.weekend_day_hard_max_deviation is not None:
                        d = self.constraint_options_model.weekend_day_hard_max_deviation
                        console.print(f"Adding hard weekend-day deviation constraint: worker={worker.name} ({worker.id}), day={day_name}, target={weekend_day_shift_target_number:.3f}, d={d}")
                        # assigned - target <= d
                        self.model.constraints.add(
                            self.model.weekend_day_shift_count[worker.id, day_name]
                            - weekend_day_shift_target_number
                            <= d
                        )
                        # target - assigned <= d
                        self.model.constraints.add(
                            weekend_day_shift_target_number
                            - self.model.weekend_day_shift_count[worker.id, day_name]
                            <= d
                        )
                        hard_weekend_constraints_added += 2

                    if hard_weekend_constraints_added:
                        console.print(f"Worker {worker.name}: added {hard_weekend_constraints_added} hard weekend-day constraints")

            # Ensure worker is not allocated shifts on non working days
            if worker.non_working_day_list:
                for week, day, shift in self.get_all_shiftclass_combinations():
                    for n, start, end in worker.non_working_day_list:
                        if not shift.rota_on_nwds and day == n:
                            # print(start, self.week_day_date_map[(week, day)], end)
                            # If this shift is force assigned, allow it but add a warning
                            if start <= self.week_day_date_map[(week, day)] < end:
                                if (week, day, shift.name) in getattr(worker, "forced_assignments", []):
                                    self.add_warning(
                                        "Force assignment/NWD conflict",
                                        f"Worker '{worker.name}' has a forced assignment for shift '{shift.name}' on week {week}, day {day} (date: {self.get_date_by_week_day(week, day)}), which is a non-working day."
                                    )
                                else:
                                    self.model.constraints.add(
                                        0
                                        == self.model.works[
                                            worker.id, week, day, shift.name
                                        ]
                                    )

                        if self.get_workers_who_require_locums():
                            if not worker.locum_on_nwds and day == n:
                                if start <= self.week_day_date_map[(week, day)] < end:
                                    self.model.constraints.add(
                                        0
                                        == self.model.locum_works[
                                            worker.id, week, day, shift.name
                                        ]
                                    )

            if self.constraint_options_model.balance_blocks:
                if self.constraint_options_model.max_night_frequency:
                    for week_blocks in self.get_week_block_iterator(
                        self.constraint_options_model.max_night_frequency
                    ):
                        # Ignore excluded weeks
                        if set(week_blocks).intersection(
                            set(
                                self.constraint_options_model.max_night_frequency_week_exclusions
                            )
                        ):
                            continue

                        if self.get_shifts_with_constraint(NightConstraint):
                            # Prevent nights more than once every n weeks
                            self.model.constraints.add(
                                1
                                >= sum(
                                    self.model.shift_week_worker_assigned[
                                        shift.name, week, worker.id
                                    ]
                                    for week in week_blocks
                                    for shift in self.get_shifts_with_constraint(
                                        NightConstraint
                                    )
                                )
                            )

            # if self.constraint_options_model.balance_weekends:
            if self.constraint_options_model.max_weekend_frequency:
                for week_blocks in self.get_week_block_iterator(
                    self.constraint_options_model.max_weekend_frequency
                ):
                    # Prevent weekend shifts more than once every n weeks
                    self.model.constraints.add(
                        1
                        >= sum(
                            self.model.works_weekend[worker.id, week]
                            for week in week_blocks
                        )
                    )

            for constraint_shift in self.get_shifts_with_constraints(
                MaxShiftsPerWeekBlockConstraint
            ):
                constraint = constraint_shift.get_constraint(MaxShiftsPerWeekBlockConstraint)

                # If not supplied, default to the number of days per shift
                if constraint.max_shifts is not None:
                    shift_number = constraint.max_shifts
                else:
                    shift_number = constraint_shift.get_shift_number()

                for week_blocks in self.get_week_block_iterator(
                    constraint.week_block
                ):
                    # Prevent shifts more than once every n weeks
                    self.model.constraints.add(
                        shift_number
                        >= sum(
                            self.model.works[
                                worker.id, week, day, constraint_shift.name
                            ]
                            for week in week_blocks
                            for day in constraint_shift.days
                        )
                    )

            for week in self.weeks:

                weekend_days = self.days[5:]  # ["Sat", "Sun"]
                total_assigned = sum(
                    self.model.works[worker.id, week, day, shiftname]
                    for w, day, shiftname in self.get_all_shiftname_combinations(week=week)
                    if day in weekend_days
                )
                # If any shift is assigned on Sat or Sun, works_weekend is 1
                self.model.constraints.add(total_assigned >= self.model.works_weekend[worker.id, week])
                # If works_weekend is 1, at least one shift must be assigned
                self.model.constraints.add(total_assigned <= 1000 * self.model.works_weekend[worker.id, week])



                for constraint_shift in self.get_shifts_with_constraints(
                    MaxShiftsPerWeekConstraint
                ):

                    # If not supplied, default to the number of days per shift
                    constraint = constraint_shift.get_constraint(MaxShiftsPerWeekConstraint)
                    if constraint.max_shifts is not None:
                        shift_number = constraint.max_shifts
                    else:
                        shift_number = constraint_shift.get_shift_number()
                    try:
                        # If the constraint specifies particular days, only count those days,
                        # otherwise count the shift's own days.
                        constraint_days = (
                            constraint.days if constraint.days is not None else self.days
                        )

                        self.model.constraints.add(
                            shift_number
                            >= sum(
                                self.model.works[worker.id, w, day, constraint_shift.name]
                                for w, day in self.get_week_day_combinations_for_shift(constraint_shift)
                                if w == week and day in constraint_days
                            )
                        )
                    except ValueError:
                        # This happens if a shift if not assigned on the week (should we test for this instead?)
                        pass

                try:
                    self.model.constraints.add(
                        self.constraint_options_model.max_shifts_per_week
                        >= sum(
                            self.model.works[worker.id, w, day, shiftname]
                            # for shiftname in self.get_shift_names_by_week_day(week, day)
                            # for day in self.days
                            for w, day, shiftname in self.get_all_shiftname_combinations(
                                week=week
                            )
                        )
                    )
                except ValueError:
                    pass

                # TODO: consider excluding shifts that span relevant period
                if (
                    worker.site
                    in self.constraint_options_model.prevent_monday_after_full_weekends
                ):
                    # Ignore last week
                    if week + 1 in self.weeks:
                        self.model.constraints.add(
                            full_weekend_count
                            >= sum(
                                self.model.works[worker.id, week, "Sat", shift]
                                for shift in self.get_shift_names_by_week_day(
                                    week, "Sat"
                                )
                            )
                            + sum(
                                self.model.works[worker.id, week, "Sun", shift]
                                for shift in self.get_shift_names_by_week_day(
                                    week, "Sun"
                                )
                            )
                            + sum(
                                self.model.works[worker.id, week + 1, "Mon", shift]
                                for shift in self.get_shift_names_by_week_day(
                                    week + 1, "Mon"
                                )
                            )
                        )
                if (
                    worker.site
                    in self.constraint_options_model.prevent_monday_and_tuesday_after_full_weekends
                ):
                    # Ignore last week
                    if week + 1 in self.weeks:
                        self.model.constraints.add(
                            full_weekend_count
                            >= sum(
                                self.model.works[worker.id, week, "Sat", shift]
                                for shift in self.get_shift_names_by_week_day(
                                    week, "Sat"
                                )
                            )
                            + sum(
                                self.model.works[worker.id, week, "Sun", shift]
                                for shift in self.get_shift_names_by_week_day(
                                    week, "Sun"
                                )
                            )
                            + sum(
                                self.model.works[worker.id, week + 1, "Mon", shift]
                                for shift in self.get_shift_names_by_week_day(
                                    week + 1, "Mon"
                                )
                            )
                            + sum(
                                self.model.works[worker.id, week + 1, "Tue", shift]
                                for shift in self.get_shift_names_by_week_day(
                                    week + 1, "Tue"
                                )
                            )
                        )
                if (
                    worker.site
                    in self.constraint_options_model.prevent_fridays_before_full_weekends
                ):
                    self.model.constraints.add(
                        full_weekend_count
                        >= sum(
                            self.model.works[worker.id, week, "Sat", shift.name]
                            for shift in self.get_shift_names_by_week_day(week, "Sat")
                        )
                        + sum(
                            self.model.works[worker.id, week, "Sun", shift.name]
                            for shift in self.get_shift_names_by_week_day(week, "Sat")
                        )
                        + sum(
                            self.model.works[worker.id, week, "Fri", shift.name]
                            for shift in self.get_shift_names_by_week_day(week, "Sat")
                        )
                    )
                if (
                    worker.site
                    in self.constraint_options_model.prevent_thursdays_before_full_weekends
                ):
                    self.model.constraints.add(
                        full_weekend_count
                        >= sum(
                            self.model.works[worker.id, week, "Sat", shift.name]
                            for shift in self.get_shift_names_by_week_day(week, "Sat")
                        )
                        + sum(
                            self.model.works[worker.id, week, "Sun", shift.name]
                            for shift in self.get_shift_names_by_week_day(week, "Sun")
                        )
                        + sum(
                            self.model.works[worker.id, week, "Thu", shift.name]
                            for shift in self.get_shift_names_by_week_day(week, "Thu")
                        )
                    )

                # # model.weekend_count stores the number of weekend shifts that a worker is assigned to wor
                # # this is used (by the objective) to balance the total number of weekend shifts worked
                # self.model.constraints.add(
                #     # We could use a helper to get required shifts to check
                #     self.model.works_whole_weekend[worker.id, week] >= sum(
                #         self.model.works[worker.id, week, day, shift.name]
                #         for day in self.days[5:]
                #         for shift in self.get_shifts()) - 1)
                # self.model.constraints.add(
                #     self.model.works_whole_weekend[worker.id, week] <= sum(
                #         self.model.works[worker.id, week, day, shift.name]
                #         for day in self.days[5:]
                #         for shift in self.get_shifts()) / 2)
                # if self.constraint_options["balance_weekends"]:


                # Try disabling tihs to allow more than 2 shifts ot be assigned on a weekend
                #self.model.constraints.add(
                #    self.model.works_weekend[worker.id, week]
                #    >= sum(
                #        self.model.works[worker.id, week, day, shiftname]
                #        for w, day, shiftname in self.get_all_shiftname_combinations(
                #            week=week
                #        )
                #        if day in self.days[5:]
                #    )
                #    / 2
                #)
                #self.model.constraints.add(
                #    self.model.works_weekend[worker.id, week]
                #    <= sum(
                #        self.model.works[worker.id, week, day, shiftname]
                #        for w, day, shiftname in self.get_all_shiftname_combinations(
                #            week=week
                #        )
                #        if day in self.days[5:]
                #    )
                #)
                # self.model.constraints.add(
                #     self.model.works_saturday[worker.id, week] == sum(
                #         self.model.works[worker.id, week, "Sat", shift.name]
                #         for shift in self.get_shifts()))
                # self.model.constraints.add(
                #     self.model.works_sunday[worker.id, week] == sum(
                #         self.model.works[worker.id, week, "Sun", shift.name]
                #         for shift in self.get_shifts()))

                # self.model.constraints.add(1 == sum(
                #     self.model.shift_week_worker_assigned[s.name, week, worker.id]
                #     for s in self.get_shifts_with_constraint("night")))

                # for shift in self.get_shifts_with_constraint("single_block_per_week"):
                #    if shift.force_as_block:
                #        # Force nights to be assigned in blocks
                #        # self.model.constraints.add(8* self.model.shift_week_worker_assigned[shift.name, week, worker.id] >= sum(self.model.works[worker.id, week, day, shift.name] for day in self.days)
                #        #     )

                #        # Force night shifts to be assigned in blocks
                #        self.model.constraints.add(
                #            self.model.shift_week_worker_assigned[
                #                shift.name, week, worker.id
                #            ]
                #            * len(shift.days)
                #            == sum(
                #                self.model.works[worker.id, w, d, s]
                #                for w, d, s in self.get_all_shiftname_combinations(week=week)
                #                #for day in self.days
                #            )
                #        )
                #    else:
                #        raise ValueError(f"Requires {shift.name} to have force_as_block")

                # for shift in self.get_shifts_with_constraint("night"):
                for shift in self.get_shifts(week=week):
                    if shift.force_as_block:
                        if shift.get_worker_requirement_by_date(
                            self.get_week_start_date(week)
                        ):
                            # Force shifts to be assigned in blocks
                            self.model.constraints.add(
                                self.model.shift_week_worker_assigned[
                                    shift.name, week, worker.id
                                ]
                                * len(shift.days)
                                == sum(
                                    self.model.works[worker.id, week, day, shift.name]
                                    for day in shift.days
                                )
                            )

                #

            # def get_pre_may(week_days, max_pre)

            for n in track(
                range(len(weeks_days)), description="Generate week/day constraints"
            ):
                week, day = weeks_days[n]

                pre_map = []
                for pre_n in range(1, self.max_pre + 1):
                    try:
                        pweek, pday = weeks_days[n - pre_n]
                        p = 1
                    except:
                        pweek, pday = weeks_days[n]
                        p = 0

                    pre_map.append((p, pweek, pday))

                post_map = []
                for post_n in range(1, self.max_post + 1):
                    try:
                        pweek, pday = weeks_days[n + post_n]
                        p = 1
                    except:
                        pweek, pday = weeks_days[n]
                        p = 0

                    post_map.append((p, pweek, pday))

                # p1 = 1
                # p2 = 1
                # p3 = 1
                # p4 = 1
                # p5 = 1
                # p6 = 1
                # p7 = 1
                # if n > 0:
                #    pweek, pday = weeks_days[n - 1]
                # else:
                #    p1 = 0
                #    pweek, pday = weeks_days[n]

                # try:
                #    p2week, p2day = weeks_days[n - 2]
                # except IndexError:
                #    p2week, p2day = weeks_days[n]
                #    p2 = 0

                # try:
                #    p3week, p3day = weeks_days[n - 3]
                # except IndexError:
                #    p3week, p3day = weeks_days[n]
                #    p3 = 0

                # try:
                #    p4week, p4day = weeks_days[n - 4]
                # except IndexError:
                #    p4week, p4day = weeks_days[n]
                #    p4 = 0

                # try:
                #    p5week, p5day = weeks_days[n - 5]
                # except IndexError:
                #    p5week, p5day = weeks_days[n]
                #    p5 = 0

                # try:
                #    p6week, p6day = weeks_days[n - 6]
                # except IndexError:
                #    p6week, p6day = weeks_days[n]
                #    p6 = 0

                # try:
                #    p7week, p7day = weeks_days[n - 7]
                # except IndexError:
                #    p7week, p7day = weeks_days[n]
                #    p7 = 0

                # pre_map = [
                #    (p1, pweek, pday),
                #    (p2, p2week, p2day),
                #    (p3, p3week, p3day),
                #    (p4, p4week, p4day),
                #    (p5, p5week, p5day),
                #    (p6, p6week, p6day),
                #    (p7, p7week, p7day),
                # ]

                # n1 = 1
                # n2 = 1
                # n3 = 1
                # n4 = 1
                # n5 = 1
                # n6 = 1
                # n7 = 1
                # if n > 0:
                #    nweek, nday = weeks_days[n - 1]
                # else:
                #    n1 = 0
                #    nweek, nday = weeks_days[n]

                # try:
                #    n2week, n2day = weeks_days[n - 2]
                # except IndexError:
                #    n2week, n2day = weeks_days[n]
                #    n2 = 0

                # try:
                #    n3week, n3day = weeks_days[n - 3]
                # except IndexError:
                #    n3week, n3day = weeks_days[n]
                #    n3 = 0

                # try:
                #    n4week, n4day = weeks_days[n - 4]
                # except IndexError:
                #    n4week, n4day = weeks_days[n]
                #    n4 = 0

                # try:
                #    n5week, n5day = weeks_days[n - 5]
                # except IndexError:
                #    n5week, n5day = weeks_days[n]
                #    n5 = 0

                # try:
                #    n6week, n6day = weeks_days[n - 6]
                # except IndexError:
                #    n6week, n6day = weeks_days[n]
                #    n6 = 0

                # try:
                #    n7week, n7day = weeks_days[n - 7]
                # except IndexError:
                #    n7week, n7day = weeks_days[n]
                #    n7 = 0

                # post_map = [
                #    (n1, nweek, nday),
                #    (n2, n2week, n2day),
                #    (n3, n3week, n3day),
                #    (n4, n4week, n4day),
                #    (n5, n5week, n5day),
                #    (n6, n6week, n6day),
                #    (n7, n7week, n7day),
                # ]

                # n1 = 1
                # n2 = 1
                # try:
                #    nweek, nday = weeks_days[n + 1]
                # except IndexError:
                #    nweek, nday = weeks_days[n]
                #    n1 = 0

                # try:
                #    n2week, n2day = weeks_days[n + 2]
                # except IndexError:
                #    n2week, n2day = weeks_days[n]
                #    n2 = 0

                # try:
                #    n3week, n3day = weeks_days[n + 3]
                # except IndexError:
                #    n3week, n3day = weeks_days[n]
                #    n3 = 0

                # IF paired we check the following against both workers
                workers = [worker]
                if self.constraint_options_model.hard_constrain_pair_separation:
                    for worker_pairs in self.worker_pairs:
                        if worker_pairs[0] == worker:
                            workers = worker_pairs

                if self.get_shift_names_by_week_day(week, day):
                    # Unable to work (hard constraint not preference)
                    self.model.constraints.add(
                        self.model.available[worker.id, week, day]
                        >= sum(
                            self.model.works[worker.id, week, day, shift]
                            for shift in self.get_shift_names_by_week_day(week, day)
                        )
                    )
                    if self.get_locum_workers():
                        self.model.constraints.add(
                            self.model.available[worker.id, week, day]
                            >= sum(
                                self.model.locum_works[worker.id, week, day, shift]
                                for shift in self.get_shift_names_by_week_day(week, day)
                            )
                        )

                    # Shift-specific active dates constraint
                    for shift in self.get_shifts():
                        if (worker.id, week, day, shift.name) in self.model.works:
                            calc_start = getattr(worker, "calculated_shift_start_dates", {}).get(shift.name, worker.calculated_start_date)
                            calc_end = getattr(worker, "calculated_shift_end_dates", {}).get(shift.name, worker.calculated_end_date)
                            if calc_start is not None and calc_end is not None:
                                date = self.week_day_date_map[(week, day)]
                                if date < calc_start or date >= calc_end:
                                    self.model.constraints.add(
                                        self.model.works[worker.id, week, day, shift.name] == 0
                                    )
                                    if self.get_locum_workers() and (worker.id, week, day, shift.name) in self.model.locum_works:
                                        self.model.constraints.add(
                                            self.model.locum_works[worker.id, week, day, shift.name] == 0
                                        )

                    # single shift per day (unless multi-shift allowed)
                    # This is signifantly slower so only enable if required
                    shifts_today = self.get_shift_names_by_week_day(week, day)
                    assigned_vars = [
                        self.model.works[worker.id, week, day, shift]
                        for shift in shifts_today
                    ]
                    worker_allowed_sets = getattr(
                        worker, "allowed_multi_shift_sets", []
                    )
                    force_assign_sets = []
                    for shift_name in shifts_today:
                        if hasattr(worker, "force_assign_with") and shift_name in worker.force_assign_with:
                            force_set = set([shift_name] + list(worker.force_assign_with[shift_name]))
                            if force_set.issubset(shifts_today):
                                force_assign_sets.append(force_set)
                        shift_obj = self.get_shift_by_name(shift_name)
                        if getattr(shift_obj, "force_assign_with", []):
                            force_set = set([shift_name] + list(shift_obj.force_assign_with))
                            if force_set.issubset(shifts_today):
                                force_assign_sets.append(force_set)

                    # If no allowed sets, we can only assign one shift
                    allowed_sets_today = [
                        allowed_set
                        for allowed_set in worker_allowed_sets
                        if allowed_set.issubset(shifts_today)
                    ] + force_assign_sets

                    if allowed_sets_today:
                        # Forbid any multi-shift assignment that includes shifts from different allowed sets
                        # Build all possible multi-shift assignments
                        all_multi = [
                            set(c)
                            for i in range(2, len(shifts_today) + 1)
                            for c in itertools.combinations(shifts_today, i)
                        ]
                        # Build all allowed multi-shift patterns (all non-empty subsets of each allowed set)
                        allowed_patterns = [{shift} for shift in shifts_today]
                        for allowed_set in allowed_sets_today:
                            allowed_patterns += [
                                set(s)
                                for i in range(2, len(allowed_set) + 1)
                                for s in itertools.combinations(allowed_set, i)
                            ]
                        # Forbid all other multi-shift combinations
                        forbidden = [s for s in all_multi if s not in allowed_patterns]
                        for forbidden_set in forbidden:
                            self.model.constraints.add(
                                sum(
                                    self.model.works[worker.id, week, day, shift]
                                    for shift in forbidden_set
                                )
                                <= len(forbidden_set) - 1
                            )
                        # Still, at most all shifts per day
                        self.model.constraints.add(
                            sum(
                                self.model.works[worker.id, week, day, shift]
                                for shift in shifts_today
                            )
                            <= len(shifts_today)
                        )
                    else:
                        # Only one shift per day if no allowed sets
                        self.model.constraints.add(sum(assigned_vars) <= 1)

                    # This applies to locums as well
                    if self.get_workers_who_require_locums():
                        self.model.constraints.add(
                            1
                            >= sum(
                                self.model.locum_works[worker.id, week, day, shift]
                                for shift in self.get_shift_names_by_week_day(week, day)
                            )
                        )
                        self.model.constraints.add(
                            1
                            >= sum(
                                self.model.works[worker.id, week, day, shift]
                                + self.model.locum_works[worker.id, week, day, shift]
                                for shift in self.get_shift_names_by_week_day(week, day)
                            )
                        )

                    if self.constraint_options_model.hard_constrain_pair_separation:
                        for worker_pairs in self.worker_pairs:
                            if worker_pairs[0] == worker:
                                self.model.constraints.add(
                                    1
                                    >= sum(
                                        self.model.works[w.id, week, day, shift]
                                        for shift in self.get_shift_names_by_week_day(
                                            week, day
                                        )
                                        for w in worker_pairs
                                    )
                                )

                # NOTE: due to the way that this is implemented, if a
                # shift spans 7 days you may get >7 allocations in a row
                # as it only checks for a different shift allocation
                for constraint_shift in self.get_shifts_with_constraints(
                    PreShiftConstraint
                ):
                    constraint = constraint_shift.get_constraint(PreShiftConstraint)
                    if constraint.allow_self:
                        ignore_shifts = constraint.ignore_shifts + [constraint_shift.name]
                    else:
                        ignore_shifts = constraint.ignore_shifts
                    weeks = _get_constraint_weeks(constraint, self)
                    if weeks is not None and week not in weeks:
                        continue
                    date_main = self.get_date_by_week_day(week, day)
                    for n in range(0, constraint.days):
                        # skip if main day excluded by weekday or specific date
                        if day in getattr(constraint, "exclude_days", []) or date_main in getattr(constraint, "exclude_dates", []):
                            continue
                        if day in constraint_shift.days:
                            # Only apply if worker can work this shift
                            if not self.is_worker_eligible_for_shift(worker, constraint_shift):
                                continue
                            try:
                                works = self.model.works[
                                    worker.id, week, day, constraint_shift.name
                                ]
                            except KeyError:
                                continue
                            # compute the pre date and skip if it's in exclude_dates
                            try:
                                pre_date = self.get_date_by_week_day(pre_map[n][1], pre_map[n][2])
                            except Exception:
                                pre_date = None
                            if pre_date is not None and pre_date in getattr(constraint, "exclude_dates", []):
                                continue
 
                            self.model.constraints.add(
                                1
                                >= works
                                + sum(
                                    pre_map[n][0]
                                    * self.model.works[
                                        w.id, pre_map[n][1], pre_map[n][2], shiftname
                                    ]
                                    for shiftname in self.get_shift_names_by_week_day(
                                        pre_map[n][1], pre_map[n][2]
                                    )
                                    if shiftname not in ignore_shifts
                                    for w in workers
                                    if self.is_worker_eligible_for_shift(w, constraint_shift)  # Only workers who can work this shift
                                )
                            )

                for constraint_shift in self.get_shifts_with_constraints(
                    PostShiftConstraint
                ):
                    constraint = constraint_shift.get_constraint(PostShiftConstraint)
                    if constraint.allow_self:
                        ignore_shifts = constraint.ignore_shifts + [constraint_shift.name]
                    else:
                        ignore_shifts = constraint.ignore_shifts
                    weeks = _get_constraint_weeks(constraint, self)
                    if weeks is not None and week not in weeks:
                        continue
                    date_main = self.get_date_by_week_day(week, day)
                    for n in range(0, constraint.days):
                        # skip if main day excluded by weekday or specific date
                        if day in getattr(constraint, "exclude_days", []) or date_main in getattr(constraint, "exclude_dates", []):
                            continue
                        if day in constraint_shift.days:
                            # Only apply if worker can work this shift
                            if not self.is_worker_eligible_for_shift(worker, constraint_shift):
                                continue
                            try:
                                works = self.model.works[
                                    worker.id, week, day, constraint_shift.name
                                ]
                            except KeyError:
                                continue
                            # compute the post date and skip if it's in exclude_dates
                            try:
                                post_date = self.get_date_by_week_day(post_map[n][1], post_map[n][2])
                            except Exception:
                                post_date = None
                            if post_date is not None and post_date in getattr(constraint, "exclude_dates", []):
                                continue
 
                            self.model.constraints.add(
                                1
                                >= works
                                + sum(
                                    post_map[n][0]
                                    * self.model.works[
                                        w.id, post_map[n][1], post_map[n][2], shiftname
                                    ]
                                    for shiftname in self.get_shift_names_by_week_day(
                                        post_map[n][1], post_map[n][2]
                                    )
                                    if shiftname not in ignore_shifts
                                    for w in workers
                                    if self.is_worker_eligible_for_shift(w, constraint_shift)  # Only workers who can work this shift
                                )
                            )

                # Night constraint means we won't assign a shift the day before
                # an unavailability
                for constraint_shift in self.get_shifts_with_constraint(NightConstraint):
                    if (
                        worker.id,
                        # pweek,
                        pre_map[0][1],
                        # pday,
                        pre_map[0][2],
                        constraint_shift.name,
                    ) in self.model.works:
                        # Ensure night prior to unavalibity is not assigned
                        self.model.constraints.add(
                            self.model.available[worker.id, week, day]
                            >= self.model.works[
                                # worker.id, pweek, pday, constraint_shift.name
                                worker.id,
                                pre_map[0][1],
                                pre_map[0][2],
                                constraint_shift.name,
                            ]
                        )

        ## Ensure each shift has the requisit number of workers assigned
        # for (
        #    week,
        #    day,
        #    shift,
        #    workers_required,
        #    site_required,
        # ) in self.get_required_workers_and_site_combinations():
        #    # print(week, day, shift, workers_required, site_required)
        #    self.model.constraints.add(
        #        workers_required
        #        == sum(
        #            self.model.works[worker.id, week, day, shift]
        #            for worker in self.workers
        #            if worker.locum
        #        )
        #    )

        # for worker in self.workers:
        #    if worker.locum:
        #        for week, day, shift in self.get_all_shiftname_combinations():
        #            self.model.constraints.add(
        #                self.model.locum_works[worker.id, week, day, shift]
        #                == self.model.works[worker.id, week, day, shift]
        #            )

        self.define_objectives()

        print("Building model completed")

    def define_objectives(self):
        # Define an objective function with model as input, to pass later
        def obj_rule(m):
            # c = len(workers)
            prefer_multi_shift_expr = sum(
                getattr(worker, "prefer_multi_shift_together", 0)
                * self.model.multi_shift_together[worker.id, week, day, idx]
                for worker in self.workers
                for week, day in self.get_week_day_combinations()
                for idx, allowed_set in enumerate(
                    getattr(worker, "allowed_multi_shift_sets", [])
                )
                if allowed_set.issubset(self.get_shift_names_by_week_day(week, day))
            )

            balance_modifier_constant = 10
            balance_quadratic_shift_modifier_constant = 3
            block_shift_balancing_constant = 1
            locum_shift_balance_modifier_constant = 40

            locum_shift_balancing = 0
            if self.get_workers_who_require_locums():
                if self.constraint_options_model.balance_locum_shifts:
                    locum_shift_balancing = sum(
                        locum_shift_balance_modifier_constant
                        * (
                            self.model.locum_shifts_t1[(worker.id)]
                            + self.model.locum_shifts_t2[(worker.id)]
                        )
                        for worker in self.workers
                    )

            if self.constraint_options_model.balance_shifts_over_workers:
                worker_shift_balancing = sum(
                    balance_modifier_constant
                    * (
                        self.model.worker_shift_count_t1[(worker.id)]
                        + self.model.worker_shift_count_t2[(worker.id)]
                    )
                    for worker in self.workers
                )
            else:
                worker_shift_balancing = 0

            if self.constraint_options_model.balance_shifts:
                shift_balancing = sum(
                    balance_modifier_constant
                    * (
                        self.model.shift_count_t1[worker.id, shift.name]
                        + self.model.shift_count_t2[worker.id, shift.name]
                    )
                    for worker in self.workers
                    for shift in self.get_shifts()
                )
            else:
                shift_balancing = 0

            if self.constraint_options_model.balance_shifts_quadratic:
                quadratic_shift_balancing = sum(
                    balance_quadratic_shift_modifier_constant
                    * (
                        self.model.shift_count_t1[worker.id, shift.name]
                        + self.model.shift_count_t2[worker.id, shift.name]
                    )
                    for worker in self.workers
                    for shift in self.get_shifts()
                )
            else:
                quadratic_shift_balancing = 0

            true_quadratic_shift_balancing = 0
            if self.constraint_options_model.balance_shifts_true_quadratic:
                shift_diff_modifier_constant = 10
                true_quadratic_shift_balancing = sum(
                    shift_diff_modifier_constant
                    * self.model.shift_count_diff_summed[(worker.id)]
                    for worker in self.workers
                )

            if self.constraint_options_model.minimise_shift_diffs:
                shift_diff_modifier_constant = 10
                shift_diff_balancing = sum(
                    shift_diff_modifier_constant
                    * self.model.shift_count_diff_summed[(worker.id)]
                    for worker in self.workers
                )
            else:
                shift_diff_balancing = 0

            if self.constraint_options_model.balance_nights:
                night_balance_modifier_constant = 10
                night_shift_balancing = sum(
                    night_balance_modifier_constant
                    * self.model.night_shift_count_w[(worker.id)]
                    # (self.model.night_shift_count_t1[(worker.id)] +
                    #  self.model.night_shift_count_t2[(worker.id)])
                    for worker in self.workers
                )
            else:
                night_shift_balancing = 0

            if self.constraint_options_model.balance_bank_holidays:
                bank_holiday_balance_modifier_constant = 1000
                bank_holiday_balancing = sum(
                    bank_holiday_balance_modifier_constant
                    * (self.model.bank_holiday_count_w[(worker.id)] - 1)
                    # (self.model.night_shift_count_t1[(worker.id)] +
                    #  self.model.night_shift_count_t2[(worker.id)])
                    for worker in self.workers
                )
            else:
                bank_holiday_balancing = 0

            if self.constraint_options_model.balance_weekends:
                weekend_balance_modifier_constant = 5
                weekend_shift_balancing = sum(
                    weekend_balance_modifier_constant
                    * self.model.weekend_shift_count_w[(worker.id)]
                    for worker in self.workers
                )
            else:
                weekend_shift_balancing = 0

            # Balance each weekend day (Sat / Sun) separately if requested.
            # Support either a linear term or a stronger quadratic penalty.
            weekend_day_term = 0
            if self.constraint_options_model.balance_weekend_days_quadratic:
                wd_weight = self.constraint_options_model.balance_weekend_days_weight
                # HiGHS/appsi does not accept arbitrary nonlinear objective expressions here.
                # Use the existing t1/t2 linear representation (which approximates absolute
                # deviation) in the objective instead of a true squared term.
                weekend_day_term = sum(
                    wd_weight
                    * (
                        self.model.weekend_day_shift_count_t1[(worker.id, day_name)]
                        + self.model.weekend_day_shift_count_t2[(worker.id, day_name)]
                    )
                    for worker in self.workers
                    for day_name in ("Sat", "Sun")
                )
            elif self.constraint_options_model.balance_weekend_days:
                weekend_day_modifier = self.constraint_options_model.balance_weekend_days_weight
                weekend_day_term = sum(
                    weekend_day_modifier * self.model.weekend_day_shift_count_w[(worker.id, day_name)]
                    for worker in self.workers
                    for day_name in ("Sat", "Sun")
                )

            preference_constant = 10
            # Preferences (not to work)
            preferences = sum(
                self.model.pref_not_to_work[worker.id, week, day]
                * self.model.works[worker.id, week, day, shift]
                * preference_constant
                for worker in self.workers
                for week, day, shift in self.get_all_shiftname_combinations()
            )

            # Work requsets
            work_request_constant = 10000  # Needs to be very big to override loss from bank holiday requests
            work_requests = sum(
                self.model.work_requests[worker.id, week, day, shift]
                * self.model.works[worker.id, week, day, shift]
                * work_request_constant
                for worker in self.workers
                for week, day, shift in self.get_all_shiftname_combinations()
            )

            # # Spread nights
            if self.constraint_options_model.balance_nights_across_sites:
                nights_site_balancing = sum(
                    (
                        self.model.night_per_site_t1[week, shift.name, site]
                        + self.model.night_per_site_t2[week, shift.name, site]
                    )
                    * 20
                    # self.model.night_per_site2[week, block, site]
                    for week in self.weeks
                    for shift in self.get_shifts_with_constraint(NightConstraint)
                    for site in self.sites
                )
            else:
                nights_site_balancing = 0

            if self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint):
                block_site_balancing = sum(
                    (
                        self.model.shift_per_site_t1[week, shift.name, site]
                        + self.model.shift_per_site_t2[week, shift.name, site]
                    )
                    * 2000
                    # self.model.night_per_site2[week, block, site]
                    for week in self.weeks
                    for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint)
                    for site in self.sites
                )
            else:
                block_site_balancing = 0

            if self.constraint_options_model.balance_blocks:
                blocks_balancing = sum(
                    block_shift_balancing_constant
                    * self.model.blocks_assigned[week, shift, sb_idx]
                    for week in self.weeks
                    for shift in self.shifts_to_assign_as_blocks()
                    for sb_idx in range(len(self.get_shift_sub_blocks(self.get_shift_by_name(shift))))
                    if (week, shift, sb_idx) in self.model.blocks_assigned
                )
            else:
                blocks_balancing = 0

            prefer_block_expr = sum(
                worker.assign_as_block_preferences.get(shift.name, 0)
                * self.model.blocks_worker_shift_assigned[worker.id, week, shift.name, sb_idx]
                for worker in self.workers
                for week in self.weeks
                for shift in self.shifts
                if hasattr(worker, "assign_as_block_preferences")
                and shift.name in worker.assign_as_block_preferences
                and self.is_worker_eligible_for_shift(worker, shift)
                for sb_idx in range(len(self.get_shift_sub_blocks(shift)))
                if (worker.id, week, shift.name, sb_idx)
                in self.model.blocks_worker_shift_assigned
            )

            # Quadratic :(
            # worker_pairs_balancing = 0
            # worker_pairs_constant = 100000
            # if self.worker_pairs:
            #    for worker_a, worker_b in self.worker_pairs:
            #        print(worker_a, worker_b)
            #        for week, day in self.get_week_day_combinations():
            #            worker_pairs_balancing = worker_pairs_balancing + sum(self.model.works[worker_a.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day)) * sum(self.model.works[worker_b.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day)) * worker_pairs_constant
            return (
                weekend_shift_balancing
                + quadratic_shift_balancing
                + shift_balancing
                + worker_shift_balancing
                + night_shift_balancing
                + shift_diff_balancing
                + preferences
                + nights_site_balancing
                + block_site_balancing
                + bank_holiday_balancing
                + blocks_balancing
                - work_requests
                + locum_shift_balancing
                + true_quadratic_shift_balancing
                - prefer_multi_shift_expr
                + prefer_block_expr
                + weekend_day_term
            )

        # add objective function to the model. rule (pass function) or expr (pass expression directly)
        self.model.obj = Objective(rule=obj_rule, sense=minimize)

    # def allow_shifts_together(self, *shift_names: str):
    #    """
    #    Allow the given set of shift names to be assigned together on a single day.
    #    """
    #    self.allowed_multi_shift_sets.append(set(shift_names))

    def add_warning(self, warning_type: str, message: str):
        print(f"[bold red]WARNING:[/bold red] {warning_type} - {message}")
        self.warnings.append((warning_type, message))

        if warning_type in self.terminate_on_warning:
            raise WarningTermination(warning_type)

    def get_warnings(self, warning_type: None | str = None):
        if warning_type is None:
            return self.warnings
        else:
            return [warning for warning in self.warnings if warning[0] == warning_type]

    def add_worker(self, worker: Worker) -> None:
        """Add a worker to the rota

        Args:
            worker (Worker):
        """
        self.workers.append(worker)

    def add_workers(self, workers: List) -> None:
        """Add multiple workers to the rota

        Args:
            workers (List(Worker)):
        """
        for worker in workers:
            self.add_worker(worker)
        # self.workers.extend(workers)

    def build_workers(self) -> None:
        """Process loaded shifts and workers

        Must be called prior to attempting to solve
        """
        if not hasattr(self, "shifts_by_name"):
            self.build_shifts()

        if not self.workers:
            raise NoWorkers("Workers must be added prior to calling build_workers")

        self.workers_id_map: dict[str, Worker] = {}
        self.workers_name_map: dict[str, Worker] = {}

        for worker in track(self.workers, description="Building workers"):
            print(worker.name)
            worker.load_rota(self)
            for s_name in getattr(worker, "exact_shifts", {}):
                try:
                    s = self.get_shift_by_name(s_name)
                    if worker.site not in s.sites:
                        self.add_warning(
                            "Invalid exact shift site",
                            f"Worker {worker.name} requested exact shifts for shift {s_name} but their site {worker.site} is not in the shift sites {s.sites}",
                        )
                    elif not self.is_worker_eligible_for_shift(worker, s):
                        self.add_warning(
                            "Invalid exact shift group",
                            f"Worker {worker.name} requested exact shifts for shift {s_name} but is not eligible for it",
                        )
                except KeyError:
                    self.add_warning(
                        "Invalid exact shift",
                        f"Worker {worker.name} requested exact shifts for non-existent shift {s_name}",
                    )

            # Validate shift-dependent start/end dates
            start_dates = getattr(worker, "shift_start_dates", [])
            end_dates = getattr(worker, "shift_end_dates", [])

            start_dict = {}
            if isinstance(start_dates, dict):
                start_dict = start_dates
            elif isinstance(start_dates, list):
                for item in start_dates:
                    if hasattr(item, "shift"):
                        start_dict[item.shift] = item.start_date

            end_dict = {}
            if isinstance(end_dates, dict):
                end_dict = end_dates
            elif isinstance(end_dates, list):
                for item in end_dates:
                    if hasattr(item, "shift"):
                        end_dict[item.shift] = item.end_date

            for s_name in set(list(start_dict.keys()) + list(end_dict.keys())):
                if s_name not in self.shifts_by_name:
                    raise InvalidShift(
                        f"Worker {worker.name} specified shift-dependent dates for non-existent shift {s_name}"
                    )
                s = self.get_shift_by_name(s_name)
                if not self.is_worker_eligible_for_shift(worker, s):
                    self.add_warning(
                        "Worker/ineligible shift date constraint",
                        f"Worker {worker.name} specified shift-dependent dates for shift {s_name} but is not eligible to work it"
                    )
            wid = worker.id
            if wid in self.workers_id_map:
                message = f"Worker with id '{wid}' has been added twice"
                self.add_warning("Worker/duplicate id", message)

            self.workers_id_map[wid] = worker

            if worker.name in self.workers_name_map:
                message = f"Worker with name '{worker.name}' has been added twice"
                self.add_warning("Worker/duplicate name", message)
            self.workers_name_map[worker.name] = worker

            worker_has_valid_shift = False
            for shift in self.shifts:
                if self.is_worker_eligible_for_shift(worker, shift):
                    worker_has_valid_shift = True
                    break

            if not worker_has_valid_shift:
                message = f"Worker with name '{worker.name}' ({worker.id}) has no valid shifts (site: {worker.site})"
                logger.warning(message)
                self.add_warning("Worker/no valid shifts", message)


            for date, shift_name in worker.forced_assignments_by_date:
                try:
                    week, day = self.get_week_day_by_date(date)
                    worker.force_assign_shift(week, day, shift_name)
                except WeekDayNotFound:
                    self.add_warning(
                        "Worker/forced assignment date not found",
                        f"Worker {worker.name} ({worker.id}) has a forced assignment for {date} ({shift_name}) but this date is not in the rota",
                    )
                    continue

        self.workers = sorted(self.workers)

        # Pairwise hard constraints to limit per-day weekend skew among equivalent workers
        # Group workers by (site, grade) and enforce for each pair and each day:
        #   assigned_i_day - assigned_j_day <= d
        #   assigned_j_day - assigned_i_day <= d
        # Use `weekend_day_hard_max_deviation` as `d` when set.
        if self.constraint_options_model.weekend_day_hard_max_deviation is not None:
            d = self.constraint_options_model.weekend_day_hard_max_deviation
            groups = defaultdict(list)
            for w in self.workers:
                groups[(w.site, w.grade)].append(w)

            pairwise_added = 0
            for (site, grade), workers_group in groups.items():
                if len(workers_group) < 2:
                    continue
                # iterate unique unordered pairs
                for i in range(len(workers_group)):
                    for j in range(i + 1, len(workers_group)):
                        wi = workers_group[i]
                        wj = workers_group[j]
                        for day_name in ("Sat", "Sun"):
                            try:
                                # wi - wj <= d
                                self.model.constraints.add(
                                    self.model.weekend_day_shift_count[wi.id, day_name]
                                    - self.model.weekend_day_shift_count[wj.id, day_name]
                                    <= d
                                )
                                # wj - wi <= d
                                self.model.constraints.add(
                                    self.model.weekend_day_shift_count[wj.id, day_name]
                                    - self.model.weekend_day_shift_count[wi.id, day_name]
                                    <= d
                                )
                                pairwise_added += 2
                            except Exception:
                                # If variables aren't present for some worker/day, skip
                                continue

            console.print(f"Added {pairwise_added} pairwise weekend-day hard constraints (d={d})")

        self.full_time_equivalent = sum(
            w.fte_adj for w in self.workers
        )  # TODO: rewrite as per shift
        self.full_time_equivalent_sites = {}
        self.workers_at_sites = {}

        for site in track(self.sites, description="Generating site workers"):
            self.workers_at_sites[site] = [w for w in self.workers if w.site == site]

            self.full_time_equivalent_sites[site] = {}

            for shift in self.get_shifts():
                self.full_time_equivalent_sites[site][shift.name] = sum(
                    w.get_fte(shift=shift.name) for w in self.workers if w.site == site
                )

        self.worker_pairs = []
        pairs = defaultdict(list)
        for w in self.workers:
            if w.pair is not None:
                pairs[w.pair].append(w)

        for p in pairs:
            self.worker_pairs.append(tuple(pairs[p]))


    def add_worker_name_constraint_by_week(
        self, names: Iterable[str], weeks: Iterable[int], shifts
    ):
        self.constraint_options_model.avoid_shifts_by_worker_names.append(
            {"names": names, "weeks": weeks, "shifts": shifts}
        )

    def add_avoid_shifts_for_workers_on_dates(
        self,
        names: Iterable[str],
        shifts: Iterable[str],
        start_date: datetime.date | str | None = None,
        end_date: datetime.date | str | None = None,
        dates: Iterable[datetime.date] | None = None,
        reason: str | None = None,
    ):
        """Add a rule preventing the given workers from being assigned the given shifts
        on the provided date range or explicit dates.

        - `names`: iterable of worker names to which the rule applies.
        - `shifts`: iterable of shift names to avoid.
        - `start_date`/`end_date`: inclusive date range (accepts date or string "%d/%m/%Y" or "%d/%m/%y").
        - `dates`: optional explicit iterable of `datetime.date` objects.
        """
        def _coerce_date(d):
            if d is None:
                return None
            if isinstance(d, datetime.date):
                return d
            if isinstance(d, str):
                for fmt in ("%d/%m/%Y", "%d/%m/%y", "%Y-%m-%d"):
                    try:
                        return datetime.datetime.strptime(d, fmt).date()
                    except ValueError:
                        continue
            raise ValueError(f"Cannot parse date: {d}")

        sd = _coerce_date(start_date)
        ed = _coerce_date(end_date)

        entry = {"names": list(names), "shifts": list(shifts), "start_date": sd, "end_date": ed, "dates": None, "reason": reason}
        if dates is not None:
            entry["dates"] = [d if isinstance(d, datetime.date) else _coerce_date(d) for d in dates]

        self.constraint_options_model.avoid_shifts_by_worker_dates.append(entry)

    def add_grade_constraint_by_week(
        self, grades: Iterable[int], weeks: Iterable[int], shifts
    ):
        self.constraint_options_model.avoid_shifts_by_grades.append(
            {"grades": grades, "weeks": weeks, "shifts": shifts}
        )

    def add_min_summed_grade_by_shifts_per_day_constraint(
        self,
        shifts: list[str],
        min_grade_sum: int,
        weeks: list[int] = None,
        days: list[str] = None,
    ):
        """
        Helper to add a minimum summed grade constraint for a set of shifts per day.
        """
        constraint = MinSummedGradeByShiftsPerDayConstraint(
            shifts=shifts,
            min_grade_sum=min_grade_sum,
            weeks=weeks,
            days=days,
        )
        self.constraint_options_model.min_summed_grade_by_shifts_per_day.append(constraint)

    def add_shift(self, shift: SingleShift) -> None:
        """Add a shift to the collection

        :param SingleShift shift: Shift object

        """
        self.shifts.append(shift)

    def add_shifts(self, *shifts: SingleShift) -> None:
        """Add multiple shifts

        Returns:
            None
        """
        self.shifts.extend(shifts)

    def clear_shifts(self) -> None:
        self.shifts = []

    def get_shift_names_by_week_day(self, week, day: DayStr, return_empty_if_week_day_not_found:bool=False) -> set:
        """Returns the shifts required for a specific day

        Returns:
            set: Set of ShiftName

        """
        if (week, day) not in self.week_day_shifts_dict:
            if return_empty_if_week_day_not_found:
                return set()
            raise WeekDayNotFound(f"Week {week}, Day {day} not found in week_day_shifts_dict")
        return self.week_day_shifts_dict[(week, day)]

    def get_week_day_shift_names_by_week_day(self, week, day: DayStr) -> set:
        """Returns the shifts required for a specific day

        Returns:
            set: Set of ShiftName

        """
        return set(
            (week, day, shift_name)
            for shift_name in self.week_day_shifts_dict[(week, day)]
        )

    def get_shift_by_name(self, name: str) -> SingleShift:
        """

        :param name:

        """
        return self.shifts_by_name[name]

    def get_shift_length_by_name(self, name):
        """Get the length of a shift

        Args:
            name (str): Name of a shift

        Returns:
            int: Length of the shift
        """
        return self.shifts_by_name[name].length

    def get_date_by_week_day(self, week: WeekInt, day: DayStr) -> datetime.date:
        try:
            return self.week_day_date_map[(week, day)]
        except KeyError:
            self.add_warning(
                "Rota/week day date map",
                f"Week {week}, Day {day} not found in week_day_date_map",
            )
            return datetime.date(1900, 1, 1)

    def get_weeks_by_date_range(self, start_date: datetime.date | None, end_date: datetime.date | None) -> list[WeekInt]:
        """Get a list of weeks that fall within a date range.

        Args:
            start_date (datetime.date): The start date of the range.
            end_date (datetime.date): The end date of the range.

        Returns:
            list[WeekInt]: A list of week numbers that fall within the date range.
        """
        # TODO: think about error if not monday start date?
        weeks = []
        if start_date is None:
            start_date = self.start_date
        if end_date is None:
            end_date = self.rota_end_date

        for week in self.weeks:
            week_start = self.get_date_by_week_day(week, "Mon")
            week_end = self.get_date_by_week_day(week, "Sun")
            if week_start >= start_date and week_end <= end_date:
                weeks.append(week)
        return weeks

    def pair_shifts(self, *shifts: SingleShift):
        # NOTE: currently designed to pair two shifts
        # this could be extended...
        if len(shifts) < 2:
            raise ValueError(f"Must pair at least two shifts: {shifts}")
        self.paired_shifts.append([*shifts])

    def get_date_range(
        self, start_date: datetime.date = None, end_date: datetime.date = None
    ):
        """Gets a range of dates

        If either start_date or end_date are not provided defaults to the rota dates"""
        if start_date is None:
            start_date = self.start_date
        if end_date is None:
            end_date = self.rota_end_date

        for n in range(int((end_date - start_date).days) + 1):
            yield start_date + datetime.timedelta(n)

    def allow_shifts_together_for_all_workers(self, *shift_names: str):
        """
        Allow the given set of shift names to be assigned together on a single day for all workers.
        """
        for worker in self.workers:
            worker.allow_shifts_together(*shift_names)

    def build_shifts(self):
        """
        Process the added shifts
        """

        self.shifts_by_name = {}
        self.shift_names = []  # type: List[ShiftName]

        # self.week_day_shifts_dict = {(week, day): set() for week in self.weeks for day in days}
        self.week_day_shifts_dict = {}
        self.sites = set()

        self.week_day_shift_product = []
        self.week_day_shiftclass_product = []

        for s in self.shifts:
            pass

        for s in self.shifts:
            if s.name in self.shift_names:
                raise InvalidShift(f"Duplicate shift: {s.name}")

            if s.start_date is not None:
                if s.start_date > self.rota_end_date or s.start_date < self.start_date:
                    self.add_warning(
                        "Shift/invalid start date",
                        f"Shift '{s.name}' has a start date outside of the rota limits ({s.start_date} vs [{self.start_date}---{self.rota_end_date}])",
                    )
            else:
                s.start_date = self.start_date

            if s.end_date is not None:
                if s.end_date > self.rota_end_date or s.end_date < self.start_date:
                    self.add_warning(
                        "Shift/invalid end date",
                        f"Shift '{s.name}' has an end date outside of the rota limits ({s.end_date} vs [{self.start_date}---{self.rota_end_date}])",
                    )
            else:
                s.end_date = self.rota_end_date


            if hasattr(s, "force_assign_with") and s.force_assign_with:
                missing = set(s.force_assign_with) - set(shift.name for shift in self.shifts)
                if missing:
                    raise InvalidShift(
                        f"Shift '{s.name}' has force_assign_with referencing non-existent shift(s): {missing}"
                    )

            all_shift_names = set(shift.name for shift in self.shifts)
            for worker in self.workers:
                if hasattr(worker, "force_assign_with"):
                    for shift_name, with_shifts in getattr(worker, "force_assign_with", {}).items():
                        missing = set([shift_name] + list(with_shifts)) - all_shift_names
                        if missing:
                            raise InvalidShift(
                                f"Worker '{worker.name}' has force_assign_with referencing non-existent shift(s): {missing}"
                            )

            # Check worker requirements
            if not isinstance(s.workers_required, int):
                # Validate the worker requirements
                for worker_requirement in s.workers_required:
                    if not isinstance(worker_requirement, WorkerRequirement):
                        raise InvalidShift(
                            f"Invalid worker requirement: {worker_requirement}"
                        )

                    if worker_requirement.start_date is None:
                        worker_requirement.start_date = s.start_date
                    elif worker_requirement.start_date > s.end_date:
                        self.add_warning(
                            "Shift/worker requirement start date",
                            f"Worker requirement start date is after shift end date: {worker_requirement.start_date} > {s.end_date}",
                        )

                    if worker_requirement.end_date is None:
                        worker_requirement.end_date = s.end_date
                    elif worker_requirement.end_date < s.start_date:
                        self.add_warning(
                            "Shift/worker requirement end date",
                            f"Worker requirement end date is before shift start date: {worker_requirement.end_date} < {s.start_date}",
                        )

            self.shifts_by_name[s.name] = s
            self.shift_names.append(s.name)

            # for day in s.days:
            #    self.week_day_shifts_dict[(week, day)].add(s.name)

            for site in list(s.sites):
                self.sites.add(site)

            for constraint in s.constraints:
                # Use a safe accessor to avoid triggering pydantic __getattr__ bugs
                _ensure_constraint_weeks(constraint, self)

        self.shift_counts = defaultdict(int)
        self.shift_worker_counts = defaultdict(int)
        for week, day in self.weeks_days_product:
            self.week_day_shifts_dict[(week, day)] = set()

            current_date = self.get_date_by_week_day(week, day)
            for s in self.shifts:
                # Check if the shift has started or ended
                if current_date < s.start_date or current_date > s.end_date:
                    continue

                # Shift not required if there are no worker requirements
                workers_required = s.get_worker_requirement_by_date(current_date)
                if not workers_required:
                    self.add_warning(
                        "Shift/worker requirement is zero",
                        f"Shift '{s.name}' has no worker requirements on {current_date}",
                    )
                    continue

                if s.bank_holidays_only:
                    if current_date in self.bank_holidays:
                        self.week_day_shift_product.append((week, day, s.name))
                        self.week_day_shiftclass_product.append((week, day, s))
                        self.week_day_shifts_dict[(week, day)].add(s.name)
                        self.shift_counts[s.name] = self.shift_counts[s.name] + 1
                        self.shift_worker_counts[s.name] = (
                            self.shift_worker_counts[s.name] + workers_required
                        )
                else:
                    if day in s.days:
                        self.week_day_shift_product.append((week, day, s.name))
                        self.week_day_shiftclass_product.append((week, day, s))
                        self.week_day_shifts_dict[(week, day)].add(s.name)
                        self.shift_counts[s.name] = self.shift_counts[s.name] + 1
                        self.shift_worker_counts[s.name] = (
                            self.shift_worker_counts[s.name] + workers_required
                        )

        self.max_pre = 1
        self.max_post = 1
        for shift in self.get_shifts_with_constraint(PreShiftConstraint):
            self.max_pre = max(shift.get_constraint(PreShiftConstraint).days, self.max_pre)

        for shift in self.get_shifts_with_constraint(PostShiftConstraint):
            self.max_post = max(shift.get_constraint(PostShiftConstraint).days, self.max_post)

        # --- Check allowed_multi_shift_sets validity ---
        all_shift_names = set(s.name for s in self.shifts)
        for worker in self.workers:
            for allowed_set in getattr(worker, "allowed_multi_shift_sets", []):
                missing = set(allowed_set) - all_shift_names
                if missing:
                    raise InvalidShift(
                        f"Worker '{worker.name}' has allowed_multi_shift_set(s) with non-existent shift(s): {missing}"
                    )

        ## For paired shifts we create proxy shifts
        # self.paired_shift_map = {}
        # for shifts in self.paired_shifts:
        #    sites = shifts[0].sites
        #    for shift in shifts:
        #        self.paired_shift_map[shift] = shifts
        #
        #        if shift.sites != sites:
        #            # This will lead to odd/broken behaviour
        #            self.add_warning(
        #                "Paired shifts/site mismatch",
        #                f"Shift '{shift.name}' has different sites to paired shifts",
        #            )

    #        # Todo replace with week_day.....
    #        self.day_shift_product = []
    #        self.day_shiftclass_product = []
    #        for s in self.shifts:
    #            for day in days:
    #                if day in s.days:
    #                    self.day_shift_product.append((day, s.name))
    #                    self.day_shiftclass_product.append((day, s))
    #
    #    def get_day_shiftname_combinations(self):
    #        """Returns a list of all required day / shift combinations
    #
    #        Returns:
    #            list: list of in the following format (str->day, str->shift)
    #        """
    #        return self.day_shift_product
    #
    #    def get_day_shiftclass_combinations(self) -> List[Tuple[DayStr, SingleShift]]:
    #        """Returns a list of all required day / shift combinations
    #
    #        Returns:
    #            list: list of in the following format (str->day, shiftObject->shift)
    #        """
    #        return self.day_shiftclass_product

    def get_all_shiftname_combinations(
        self, week: int = None, day: int = None
    ) -> List[Tuple[WeekInt, DayStr, ShiftName]]:
        """Returns a list of all possible week / day / shift combinations

        Returns:
            list: contains tuple of all possible week / day / shift combinations
        """
        week_day_shifts = self.week_day_shift_product

        if week is not None:
            week_day_shifts = [t for t in week_day_shifts if t[0] == week]

        if day is not None:
            week_day_shifts = [t for t in week_day_shifts if t[1] == day]

        return week_day_shifts

    def get_all_shiftclass_combinations(
        self, week: int = None, day: int = None
    ) -> List[Tuple[WeekInt, DayStr, SingleShift]]:
        """Returns a list of all possible week / day / shift combinations

        Returns:
            list: contains tuple of all possible week / day / shift combinations
        """
        week_day_shifts = self.week_day_shiftclass_product

        if week is not None:
            week_day_shifts = [t for t in week_day_shifts if t[0] == week]

        if day is not None:
            week_day_shifts = [t for t in week_day_shifts if t[1] == day]

        return week_day_shifts

    def get_week_day_combinations(self) -> list:
        """Returns a list of all week / day tuple combinations

        Returns:
            list: list of possible week day combinations
                    [(1, "Mon"), (1, "Tue"), ... (n, "Fri")]
        """
        return self.weeks_days_product

    def get_week_day_by_date(self, date: datetime.date) -> tuple[int, DayStr]:
        """Returns the week and day for a given date

        Args:
            date (datetime.date): Date to get week and day for

        Returns:
            tuple[int, DayStr]: Tuple containing week number and day string
        """
        try:
            week, day = self.date_week_day_map[date]
            return week, day
        except KeyError:
            raise WeekDayNotFound(f"No week/day found for date {date}")

    def get_week_day_combinations_for_shift(self, shift: SingleShift) -> list:
        return [
            (week, day)
            for week, day in self.get_week_day_combinations()
            if shift.name in self.get_shift_names_by_week_day(week, day)
        ]

    def get_week_day_shift_combinations_for_constraint(self, constraint) -> list:
        constraint_shifts = self.get_shifts_with_constraint(constraint)

        week_day_shifts = []

        for tup in self.get_all_shiftclass_combinations():
            if tup[2] in constraint_shifts:
                week_day_shifts.append(tup)

        return week_day_shifts

    def get_shifts(self, week: int | None = None) -> List[SingleShift]:
        """Returns a list of all the registered shifts

        Returns:
            List[SingleShift]: list of registered shifts (as SingleShift)
        """
        if week is not None:
            date = self.get_week_start_date(week)

            shifts = []

            for shift in self.shifts:
                if shift.start_date <= date < shift.end_date:
                    if shift.get_worker_requirement_by_date(date):
                        shifts.append(shift)

            return shifts

        else:
            return self.shifts

    def get_shifts_for_worker_site(self, worker_site):
        shifts = [shift for shift in self.shifts if worker_site in shift.sites]
        return shifts

    def get_shifts_for_worker(self, worker_id):
        worker = self.get_worker_by_id(worker_id)
        return [shift for shift in self.shifts if self.is_worker_eligible_for_shift(worker, shift)]

    def get_shifts_with_constraint(self, constraint) -> List[SingleShift]:
        return [shift for shift in self.shifts if shift.has_constraint(constraint)]

    def get_shifts_with_constraints(
        self, *constraints, week: int | None = None
    ) -> List[SingleShift]:
        shift_names = set()

        for constraint in constraints:
            shift_names.update(
                [
                    shift.name
                    for shift in self.get_shifts(week)
                    if shift.has_constraint(constraint)
                ]
            )

        return [self.get_shift_by_name(s) for s in shift_names]

    def get_shift_names(self) -> List[ShiftName]:
        (
            """Returns a list of all the registered shift names
        
        Returns:
            List[ShiftName]: List of names of all available shifts
        """
            """ """
        )
        return self.shift_names

    def get_week_block_iterator(self, block_length: int):
        """Gets a two dimensional list of week blocks in specified length

        e.g. block_length = 4
            [
                [ 1, 2, 3, 4],
                [ 2, 3, 4, 5],
                ...
            ]

        Args:
            block_length (int): length of blocks to create

        Returns:
            list: two dimensional list containing weeks in blocks
        """
        blocks = []
        for i in range(len(self.weeks) - block_length + 1):
            block = self.weeks[i : i + block_length]
            blocks.append(block)
        return blocks

    def get_required_workers_and_site_combinations(self):
        """Returns a list of all possible shifts, the workers required and site

        Returns:
            list: list (week, day, shift name, workers required, shift.sites)
        """
        l = []
        for week, day, shift in self.week_day_shiftclass_product:
            workers_required = shift.get_worker_requirement_by_date(
                self.get_date_by_week_day(week, day)
            )
            # if day in shift.days:
            l.append((week, day, shift.name, workers_required, shift.sites))

        return l

    def get_not_required_shifts(self):
        """Returns a set of all possible shifts combinations that are not required

        Includes those on days

        Returns:
            list: (week, day, shift)
        """
        l = set()

        for week in self.weeks:
            for day in days:
                for shift in self.shifts:
                    l.add((week, day, shift))

        return l - set(self.get_all_shiftclass_combinations())

    def shifts_to_assign_as_blocks(self) -> List[str]:
        return [shift.name for shift in self.get_shifts() if shift.assign_as_block]

    def shifts_to_force_as_blocks(self) -> List[str]:
        return [
            shift.name
            for shift in self.get_shifts()
            if (shift.force_as_block or shift.force_as_block_unless_nwd)
        ]

    # def shifts_to_assign_or_force_as_blocks(self, week: int | None = None) -> List[str]:
    def shifts_to_assign_or_force_as_blocks(self) -> List[str]:
        s = self.shifts_to_assign_as_blocks()
        s.extend(self.shifts_to_force_as_blocks())
        return s

    def worker_week_shifts_to_assign_as_blocks(self) -> List[Tuple[Worker, int, str]]:
        """Returns a list of tuples containing worker id, week and shift name for shifts to assign as blocks"""
        return [
            (worker, week, shift.name)
            for worker in self.workers
            for week in self.weeks
            for shift in self.get_shifts()
            if (
                shift.assign_as_block
                or shift.force_as_block
                or shift.force_as_block_unless_nwd
                or (
                    hasattr(worker, "assign_as_block_preferences")
                    and shift.name in getattr(worker, "assign_as_block_preferences", {})
                    and getattr(worker, "assign_as_block_preferences")[shift.name] != 0
                )
            )
        ]

    def week_shifts_to_assign_as_blocks(self) -> List[Tuple[int, str]]:
        """Returns a list of tuples containing week and shift name for shifts to assign as blocks"""
        return set([
            (week, shift.name)
            for worker in self.workers
            for week in self.weeks
            for shift in self.get_shifts()
            if (
                shift.assign_as_block
                or shift.force_as_block
                or shift.force_as_block_unless_nwd
                or (
                    hasattr(worker, "assign_as_block_preferences")
                    and shift.name in getattr(worker, "assign_as_block_preferences", {})
                    and getattr(worker, "assign_as_block_preferences")[shift.name] != 0
                )
            )
        ])

    def get_shift_sub_blocks(self, shift: SingleShift) -> List[List[str]]:
        if getattr(shift, "force_as_block_unless_nwd", None):
            val = shift.force_as_block_unless_nwd
            if isinstance(val, (list, tuple)) and all(isinstance(x, (list, tuple, set)) for x in val):
                return [list(x) for x in val]
            return [list(shift.days)]
        
        # Check if it is a block shift globally or per-worker
        is_block = (
            shift.assign_as_block
            or shift.force_as_block
            or any(
                hasattr(worker, "assign_as_block_preferences")
                and getattr(worker, "assign_as_block_preferences", {}).get(shift.name, 0) != 0
                for worker in self.workers
            )
        )
        if is_block:
            return [list(shift.days)]
        return []

    def get_all_locum_availability(self):
        return self.locum_availability_map

    def get_locum_availability(self, worker: Worker, week: int, day: DayStr, shift):
        if (worker.id, week, day) in self.locum_availability_map:
            if shift.name in self.locum_availability_map[(worker.id, week, day)]:
                return 1
        return 0

    def get_all_workers(self) -> List[Worker]:
        return self.workers

    def get_worker_by_id(self, worker_id: str) -> Worker:
        return self.workers_id_map[worker_id]

    def get_worker_by_name(self, name: str) -> Worker:
        return self.workers_name_map[name]

    def get_workers_by_group(self) -> dict[str, list[Worker]]:
        group_workers: dict[str, list[Worker]] = defaultdict(list)
        for worker in self.workers:
            group_workers[worker.site].append(worker)

        return group_workers

    def get_workers_by_remote_group(self) -> dict[str, list[Worker]]:
        group_workers: dict[str, list[Worker]] = defaultdict(list)
        for worker in self.workers:
            group_workers[worker.remote_site].append(worker)

        return group_workers

    def get_workers_by_site(self, site: str) -> list[Worker]:
        """
        Returns a list of workers assigned to a specific site.

        Args:
            site (str): The site name.

        Returns:
            list[Worker]: List of workers at the given site.
        """
        return [worker for worker in self.workers if worker.site == site]

    def get_workers_by_grade(self) -> Dict[int, list[Worker]]:
        group_workers: dict[int, list[Worker]] = defaultdict(list)
        for worker in self.workers:
            group_workers[worker.grade].append(worker)

        return group_workers

    def is_worker_eligible_for_shift(self, worker: Worker, shift: SingleShift) -> bool:
        """Returns True if the worker is eligible for the shift based on site and group options."""
        # 1. Group checks
        worker_groups = getattr(worker, "groups", set())
        if not isinstance(worker_groups, (set, list, tuple)):
            worker_groups = {worker_groups} if worker_groups else set()
        worker_groups_set = set(worker_groups)

        # exclude_groups check: worker MUST NOT belong to any group in exclude_groups
        if getattr(shift, "exclude_groups", None):
            exclude_set = set(shift.exclude_groups)
            if worker_groups_set.intersection(exclude_set):
                return False

        # Additive Site or Include Group check
        has_site = worker.site in shift.sites
        has_include_group = False
        if getattr(shift, "include_groups", None):
            include_set = set(shift.include_groups)
            if worker_groups_set.intersection(include_set):
                has_include_group = True

        if getattr(shift, "include_groups", None):
            if not (has_site or has_include_group):
                return False
        else:
            if not has_site:
                return False

        return True

    def get_workers_in_group(self, group_name: str) -> List[Worker]:
        """Returns a list of all workers belonging to the specified group."""
        return [
            w for w in self.workers
            if group_name in getattr(w, "groups", set())
        ]

    def get_workers_for_shift(self, shift: SingleShift) -> List[Worker]:
        return [worker for worker in self.workers if self.is_worker_eligible_for_shift(worker, shift)]

    def get_workers_total_fte(self) -> float:
        """Does not take into account shift adjusted ftes"""
        return self.full_time_equivalent

    def get_worker_details(self):
        w: dict[str, Worker] = defaultdict(list)
        worker: Worker
        for worker in self.workers:
            w[worker.site].append(worker)

        l = []
        for site in w:
            l.append(f"{site}: {sum(worker.get_fte() for worker in w[site]) / 100}")
        t = "\n".join(l)
        return f"Full time equivalent trainees by site:\n{t}"

    def get_workers(self):
        return self.workers

    def get_worker_grades(self) -> set[int]:
        return set([worker.grade for worker in self.workers])

    def get_workers_who_require_locums(self):
        return [worker for worker in self.get_all_workers() if worker.locum]

    def get_locum_workers(self):
        return [
            worker
            for worker in self.get_all_workers()
            if not worker.locum and worker.locum_availability
        ]

    def get_bank_holiday_week_days(self):
        return [
            (week, day)
            for week, day in self.get_week_day_combinations()
            if self.week_day_date_map[(week, day)] in self.bank_holidays
        ]

    def get_week_start_date(self, week: int):
        """Week start date is 1 indexed"""
        week = week - 1
        return self.start_date + datetime.timedelta(weeks=week)

    def force_assign_shift(self, worker_name: str, week: int, day: str, shift_name: str):
        """
        Force a specific worker to be assigned to a shift on a specific week and day.
        """
        worker = self.get_worker_by_name(worker_name)
        worker.forced_assignments.append((week, day, shift_name))

    def force_assign_shift_by_date(self, worker_name: str, date: datetime.date, shift_name: str):
        """
        Force a specific worker to be assigned to a shift on a specific date.
        """
        worker = self.get_worker_by_name(worker_name)
        worker.forced_assignments_by_date.append((date, shift_name))
        
        
        
        
        
        


    # RESULTS
    def export_rota_to_html(
        self, filename: str = "rota", folder=None, timestamp_filename=False
    ):
        if timestamp_filename:
            filename = f"{filename}_{datetime.datetime.now().strftime('%Y%m%d-%H%M%S')}"

        if folder is not None:
            output_file = Path("output", folder, f"{filename}.html")
        else:
            output_file = Path("output", f"{filename}.html")

        # Make sure path exits
        output_file.parent.mkdir(parents=True, exist_ok=True)
        with open(output_file, "w") as f:
            f.write(self.get_worker_timetable_html(True))
            url = urllib.parse.quote(str(output_file.resolve()))
            print(f"Rota exported [bold blue]('[link url={url}]{url}')[/bold blue]")
            self.exported_rota_file = output_file

    def export_rota_to_csv(self, filename: str = "rota"):
        output_file = Path("output", f"{filename}.csv")
        works = self.model.works
        with open(output_file, "w", newline="") as f:
            wr = csv.writer(f, quoting=csv.QUOTE_ALL)
            l = ["Name"]
            l.extend([worker.name for worker in self.workers])
            wr.writerow(l)

            l2 = ["Site"]
            l2.extend([worker.site for worker in self.workers])
            wr.writerow(l2)

            l3 = ["Grade"]
            l3.extend([worker.grade for worker in self.workers])
            wr.writerow(l3)

            l4 = ["FTE"]
            l4.extend([worker.fte for worker in self.workers])
            wr.writerow(l4)

            l5 = ["Exact Shifts"]
            l5.extend([json.dumps(getattr(worker, "exact_shifts", {})) for worker in self.workers])
            wr.writerow(l5)

            for week, day in self.get_week_day_combinations():
                d = [f"Week {week} Day {day}"]

                for worker in self.workers:
                    i = ""
                    for shift in self.get_shift_names_by_week_day(week, day):
                        if works[worker.id, week, day, shift].value == 1:
                            i = shift

                    d.append(i)
                wr.writerow(d)

    def get_work_table(self):
        """Build a timetable of the week as a dictionary from the model's optimal solution."""
        works = self.model.works
        week_table = {
            week: {day: {shift: [] for shift in self.get_shift_names()} for day in days}
            for week in self.weeks
        }
        for week in self.weeks:
            for worker in self.workers:
                for day, shift in self.get_day_shiftname_combinations():
                    if works[worker.id, week, day, shift].value == 1:
                        week_table[week][day][shift].append(worker.get_details())
        return week_table

    def get_shift_timetable_by_week(self):
        """Build a timetable of the week as a dictionary from the model's optimal solution."""
        works = self.model.works
        shift_timetable = {
            shift.name: {week: {day: [] for day in days} for week in self.weeks}
            for shift in self.shifts
        }
        for shift in self.shifts:
            for week in self.weeks:
                for day in days:
                    assigned_workers = []
                    for worker in self.workers:
                        if works[worker.id, week, day, shift.name].value > 0.5:
                            assigned_workers.append(worker)
                    shift_timetable[shift.name][week][day] = assigned_workers
        return shift_timetable

    def get_worker_timetable(self):
        works = self.model.works
        timetable = {
            worker.name: {week: {day: [] for day in days} for week in self.weeks}
            for worker in self.workers
        }
        for worker in self.workers:
            for week in self.weeks:
                for day in days:
                    assigned_shifts = []
                    for shift in self.get_shift_names_by_week_day(week, day):
                        if works[worker.id, week, day, shift].value > 0.5:
                            assigned_shifts.append(shift)
                    timetable[worker.name][week][day] = assigned_shifts
        return timetable

    def get_assigned_shifts_for_worker_on_day(self, worker_id: int, week: int, day: DayStr) -> list:
        """Return a list of assigned SingleShift objects for the given worker on a specific week/day.

        This is a small compatibility helper used by tests that expect objects rather than shift-name strings.
        """
        works = self.model.works
        assigned = []
        for shift_name in self.get_shift_names_by_week_day(week, day, return_empty_if_week_day_not_found=True):
            try:
                if works[worker_id, week, day, shift_name].value > 0.5:
                    assigned.append(self.get_shift_by_name(shift_name))
            except Exception:
                # If variable missing or access error, skip
                continue
        return assigned

    def get_worker_timetable_brief(
        self, show_prefs=False, marker_every=30, show_unavailable=False
    ):
        model = self.model
        week_string = f"{'-Week-':20}" + "".join(
            [7 * str(f"{w}")[-1:] for w in self.weeks]
        )
        days_string = f"{'-Day-':20}" + "".join("MTWTFSS" * len(self.weeks))
        timetable = []
        for worker in self.workers:
            shifts = []
            w = [f"{worker.name:20}"]
            for week, day in self.get_week_day_combinations():
                shift_display_char = "-"
                shift_names = []
                for shift in self.get_shift_names_by_week_day(week, day):
                    if model.works[worker.id, week, day, shift].value > 0.5:
                        shift_names.append(
                            self.get_shift_by_name(shift).get_display_char()
                        )
                if shift_names:
                    shift_display_char = "".join(shift_names)
                    shifts.extend(shift_names)
                w.append(shift_display_char)
            shift_count = ""
            for s in set(shifts):
                shift_count = shift_count + f"{s}: {shifts.count(s)}, "
            shift_count = (
                shift_count
                + f"#weekends_worked: {model.worker_weekend_count[worker.id].value}\\#"
            )

            # compute assigned Saturday / Sunday counts for this worker
            sat_count = 0
            sun_count = 0
            for w, d in self.get_week_day_combinations():
                if d == "Sat":
                    for s in self.get_shift_names_by_week_day(w, d, return_empty_if_week_day_not_found=True):
                        if model.works[worker.id, w, d, s].value > 0.5:
                            sat_count += 1
                            break
                if d == "Sun":
                    for s in self.get_shift_names_by_week_day(w, d, return_empty_if_week_day_not_found=True):
                        if model.works[worker.id, w, d, s].value > 0.5:
                            sun_count += 1
                            break
            timetable.append(f"{''.join(w)} {shift_count}")

            if show_prefs:
                # prefs
                w = [f"{'Preferences':20}"]
                for week, day in self.get_week_day_combinations():
                    if model.pref_not_to_work[worker.id, week, day] > 0:
                        w.append("Y")
                    else:
                        w.append("N")
                timetable.append("".join(w))

            if show_unavailable:
                # prefs
                w = [f"{'Unavailable':20}"]
                for week, day in self.get_week_day_combinations():
                    if model.available[worker.id, week, day] > 0:
                        w.append("A")
                    else:
                        w.append("U")
                timetable.append("".join(w))

        i = marker_every
        while i < len(timetable):
            timetable.insert(i, week_string)
            timetable.insert(i + 1, days_string)
            i = i + marker_every + 2

        timetable.append(week_string)
        timetable.insert(0, week_string)
        timetable.insert(0, days_string)
        return "\n".join(timetable)

    def get_worker_timetable_html(
        self, include_html_tag=False, table_name="rota-table", html_static_prefix=None
    ):
        model = self.model

        # Build a lookup of per-worker/day avoid-shift rules (worker.id, week, day) -> {shifts: [...], reason: str}
        avoid_lookup: dict[tuple, dict] = {}
        try:
            for entry in getattr(self.constraint_options_model, "avoid_shifts_by_worker_dates", []):
                names = entry.get("names", []) or []
                shifts = entry.get("shifts", []) or []
                reason = entry.get("reason")
                dates = entry.get("dates")
                sd = entry.get("start_date")
                ed = entry.get("end_date")

                date_list = []
                if dates:
                    date_list = dates
                elif sd is not None or ed is not None:
                    try:
                        for d in self.get_date_range(sd, ed):
                            date_list.append(d)
                    except Exception:
                        # fall back to trying to iterate
                        pass

                for name in names:
                    matching_workers = [w for w in self.workers if w.name.strip() == name.strip()]
                    if not matching_workers:
                        continue
                    wid = matching_workers[0].id
                    for d in date_list:
                        try:
                            wk, day = self.get_week_day_by_date(d)
                        except Exception:
                            continue
                        key = (wid, wk, day)
                        existing = avoid_lookup.get(key, {"shifts": [], "reason": None})
                        existing["shifts"].extend(shifts)
                        # preserve reason if provided
                        if reason:
                            existing["reason"] = reason
                        avoid_lookup[key] = existing
        except Exception:
            # If constraint options not present or parsing fails, continue without avoid info
            avoid_lookup = {}

        timetable = []

        # Run times and heading go outside the table; store for later insertion
        run_times_html = ""
        if self.run_start_time is not None and self.run_end_time is not None:
            run_times_html = (
                "<div id='run_times_row' style='display: flex; gap: 2em; margin-bottom: 1em;'>"
                f"<div id='run_start_time_cell'>Start time: <span id='run_start_time'>{self.run_start_time:%Y-%m-%d %H:%M:%S%z}</span></div>"
                f"<div id='run_end_time_cell'>End time: <span id='run_end_time'>{self.run_end_time:%Y-%m-%d %H:%M:%S%z}</span></div>"
                f"<div id='run_elapsed_time_cell'>Elapsed: <span id='run_elapsed_time'>{self.run_end_time - self.run_start_time}</span></div>"
                "</div>"
            )

        heading_html = f"<h2>Rota start date: {self.start_date.isoformat()} ({self.weeks[-1]} weeks)</h2>"

        date_row = ["<th class='worker'></th>"]

        n = 0
        for week, day in self.get_week_day_combinations():
            d = self.start_date + datetime.timedelta(n)

            th_class = "date"
            if d in self.bank_holidays:
                th_class = th_class + " bank-holiday"

            date_row.append(
                f"<th title='{d}' class='{th_class}' data-date='{d}'>Week {week}: {day}</th>"
            )
            n = n + 1
        # Header row (dates) should be placed in a proper thead
        header_row_html = f"<thead><tr class='data-row'>{''.join(date_row)}</tr></thead>"

        current_site = ""

        for worker in self.get_all_workers():
            if worker.site != current_site:
                try:
                    site = worker.site
                    site_count = len(self.workers_at_sites.get(site, []))
                    timetable.append(
                        f"<tr><th class='site-title'>{site} n={site_count}</th></tr>"
                    )
                except KeyError as e:
                    print(e)
                current_site = worker.site

            shifts = []
            locum_shifts = []
            nwds = json.dumps(worker.non_working_day_list, default=str)
            # if worker.nwd:
            #    # TODO: limit to dates
            #    nwds = ", ".join([i[0] for i in worker.nwd])
            # else:
            #    nwds = None

            worker_targets = json.dumps(worker.shift_target_number)

            tds = []

            n = 0

            shift_tds = []
            for week, day in self.get_week_day_combinations():
                d = self.start_date + datetime.timedelta(n)
                n = n + 1
                shift_display_char = "-"
                shift_names = []
                assigned_shift_names = []
                locum = False
                shift_name = ""
                force_assigned = False

                for shift in self.get_shift_names_by_week_day(week, day):
                    shift_obj = self.get_shift_by_name(shift)
                    if model.works[worker.id, week, day, shift].value > 0.8:
                        shift_names.append(shift_obj.get_display_char())
                        assigned_shift_names.append(shift)
                        # Check if this shift was force assigned
                        if (week, day, shift) in getattr(worker, "forced_assignments", []):
                            force_assigned = True
                    elif self.get_workers_who_require_locums():
                        if model.locum_works[worker.id, week, day, shift].value > 0.8:
                            locum_shifts.append(shift)
                            shift_names.append(shift_obj.get_display_char())
                            assigned_shift_names.append(shift)
                            locum = True
                            # Check if this shift was force assigned
                            if (week, day, shift) in getattr(worker, "forced_assignments", []):
                                force_assigned = True
                if shift_names:
                    shift_display_char = "<br/>".join(
                        [
                            "<span class='multi-shift-shift'>" + name + "</span>"
                            for name in shift_names
                        ]
                    )
                    shifts.extend(assigned_shift_names)
                    shift_name = assigned_shift_names[0] if assigned_shift_names else ""

                # Visual indication for multiple shifts
                multi_shift = len(shift_names) > 1
                css_class = day
                if multi_shift:
                    css_class += " multi-shift"

                title = f"{shift_name} ({d})"
                unavailable_reason = ""
                if model.available[worker.id, week, day] > 0:
                    available = True
                else:
                    available = False
                    css_class = "unavailable"
                    try:
                        unavailable_reason = self.unavailable_to_work_reason[
                            (worker.id, week, day)
                        ]
                    except KeyError:
                        print("Error getting reason")
                        print(worker.id, worker.name)
                        print(f"Week {week}, Day {day}")
                        unavailable_reason = "????"
                    title = f"{title} / {unavailable_reason}"


                # --- Highlight force assigned shifts ---
                if force_assigned:
                    css_class += " force-assigned"
                    title = f"{title} [FORCE ASSIGNED]"

                bank_holiday = ""
                if d in self.bank_holidays:
                    title = f"{title} [Bank Holiday]"
                    css_class = " ".join((css_class, "bank-holiday"))
                    bank_holiday = f" data-bank-holiday='{self.bank_holidays[d]}'"

                requests = ""
                if (worker.id, week, day) in self.work_requests_map:
                    css_class = " ".join((css_class, "shift-requested"))
                    title = " ".join((title, "[REQUESTED]"))
                    requests = f" data-shift-request='{self.work_requests_map[worker.id, week, day]}'"

                if (worker.id, week, day) in self.locum_availability_map:
                    css_class = " ".join((css_class, "locum-availability"))
                    title = " ".join((title, "[LOCUM AVAILABILITY]"))
                    requests = f" data-locum-request='{self.locum_availability_map[worker.id, week, day]}'"

                if locum:
                    css_class = " ".join((css_class, "locum-shift"))

                # Add visual marker/data if this day/worker has avoid-shift rules
                avoid_info = avoid_lookup.get((worker.id, week, day))
                avoid_attr = ""
                if avoid_info is not None:
                    # mark the cell as having avoids; include which shifts and the reason
                    avoid_shifts = ",".join(avoid_info.get("shifts", []))
                    avoid_reason = avoid_info.get("reason") or ""
                    css_class = " ".join((css_class, "avoid-shift"))
                    avoid_attr = f" data-avoid-shifts='{avoid_shifts}' data-avoid-reason='{avoid_reason}'"

                remote_site = ""
                if shift_name:
                    shift = self.get_shift_by_name(shift_name)
                    if shift.has_constraint(RequireRemoteSitePresenceConstraint):
                        remote_site = f" data-shift-remote-site='{shift.get_constraint(RequireRemoteSitePresenceConstraint).site}'"
                    if shift.has_constraint(NightConstraint):
                        css_class = " ".join((css_class, "night-shift"))

                title_escaped = title.replace("'", "&apos;").replace('"', "&quot;")
                shift_tds.append(
                    f"<td title='{title_escaped}' class='rota-day {css_class}'"
                    f" data-shift='{','.join(assigned_shift_names)}'"
                    f" data-shift-display='{','.join(shift_names)}'"
                    f" data-available='{available}'"
                    f" data-unavailable_reason='{unavailable_reason}'"
                    f" data-date='{d}' data-week='{week}' data-day='{day}'"
                    f"{remote_site}{requests}{bank_holiday}{avoid_attr}>{shift_display_char}</td>"
                )

            shift_count = ""
            shift_count_dict = {}
            for s in set(shifts):
                c = shifts.count(s)
                shift_count_dict[s] = c
                shift_count = shift_count + f"{s}: {c}, "

            locum_shift_count = ""
            locum_shift_count_dict = {}
            for s in set(locum_shifts):
                c = locum_shifts.count(s)
                locum_shift_count_dict[s] = c
                locum_shift_count = locum_shift_count + f"{s}: {c}, "

            shift_diff_dict = {}
            for shift in self.get_shifts():
                diff = model.shift_count_diff[worker.id, shift.name].value
                shift_diff_dict[shift.name] = diff

            # compute assigned Saturday / Sunday counts for this worker (for HTML data attributes)
            sat_count = 0
            sun_count = 0
            for w, d in self.get_week_day_combinations():
                if d == "Sat":
                    for s in self.get_shift_names_by_week_day(w, d, return_empty_if_week_day_not_found=True):
                        if model.works[worker.id, w, d, s].value > 0.5:
                            sat_count += 1
                            break
                if d == "Sun":
                    for s in self.get_shift_names_by_week_day(w, d, return_empty_if_week_day_not_found=True):
                        if model.works[worker.id, w, d, s].value > 0.5:
                            sun_count += 1
                            break
            # Small HTML summary to be inserted inside the worker td (friendly display)
            # Build a small adjustments summary from prior allocations (if any)
            prior_map = {}
            adjustments_list = []
            for shift in self.get_shifts():
                prev = getattr(worker, "previous_shifts", {}) or {}
                if shift.name in prev:
                    entry = prev[shift.name]
                    # entry may be:
                    # - a tuple/list: (worked, allocated)
                    # - a dict with per-day entries (e.g., {'Sat': (w,a), 'Sun': (w,a)})
                    # - a dict with an '__all__' key or other buckets
                    if isinstance(entry, dict):
                        # If the dict already contains per-day keys (Sat/Sun) or an explicit
                        # '__all__' bucket, preserve the dict as-is so downstream code
                        # can render per-day adjustments. Otherwise, attempt to aggregate
                        # any tuple-like values found in the dict into totals.
                        if any(k in entry for k in ("Sat", "Sun")) or "__all__" in entry:
                            prior_map[shift.name] = entry
                            # also compute an overall adjustment for the adjustments list
                            try:
                                # compute totals by summing tuple/list values where present
                                worked = sum(v[0] for v in entry.values() if isinstance(v, (list, tuple)))
                                allocated = sum(v[1] for v in entry.values() if isinstance(v, (list, tuple)))
                                adj = float(allocated) - float(worked)
                            except Exception:
                                adj = 0
                        else:
                            # fallback: sum any tuple values found in the dict
                            try:
                                worked = sum(v[0] for v in entry.values() if isinstance(v, (list, tuple)))
                                allocated = sum(v[1] for v in entry.values() if isinstance(v, (list, tuple)))
                                adj = float(allocated) - float(worked)
                            except Exception:
                                worked = 0
                                allocated = 0
                                adj = 0
                            prior_map[shift.name] = dict(worked=worked, allocated=allocated, adjustment=adj)
                    else:
                        # tuple/list style entry
                        try:
                            worked, allocated = entry
                            adj = float(allocated) - float(worked)
                        except Exception:
                            worked = 0
                            allocated = 0
                            adj = 0
                        prior_map[shift.name] = dict(worked=worked, allocated=allocated, adjustment=adj)

                    if adj != 0:
                        adjustments_list.append(f"{shift.name}:{adj:+g}")

            # If weekend prior is present as a top-level worked/allocated dict (not per-day),
            # expand it into per-day entries so Sat/Sun adjustments are always available.
            if 'weekend' in prior_map:
                w = prior_map['weekend']
                # detect the simple dict shape produced earlier: {'worked': x, 'allocated': y, 'adjustment': z}
                if isinstance(w, dict) and 'worked' in w and 'allocated' in w and not ('Sat' in w or 'Sun' in w or '__all__' in w):
                    try:
                        worked = float(w.get('worked', 0))
                        allocated = float(w.get('allocated', 0))
                        # split evenly across Sat and Sun
                        sat_worked = worked / 2.0
                        sun_worked = worked - sat_worked
                        sat_alloc = allocated / 2.0
                        sun_alloc = allocated - sat_alloc
                        prior_map['weekend'] = {
                            'Sat': (sat_worked, sat_alloc),
                            'Sun': (sun_worked, sun_alloc),
                        }
                    except Exception:
                        # leave as-is if parsing fails
                        pass

            adjustments_html = ""
            if adjustments_list:
                adjustments_html = (
                    "<div class='worker-adjustments auto-generated' style='margin-top:3px; font-size:0.85em; color:#666;'>"
                    + "Adjust: "
                    + ", ".join(adjustments_list)
                    + "</div>"
                )

            # Compute per-day (Sat/Sun) adjustments for the 'weekend' shift if provided
            sat_adj = None
            sun_adj = None
            prev_all = getattr(worker, "previous_shifts", {}) or {}
            weekend_entry = prev_all.get("weekend")
            if weekend_entry is not None:
                try:
                    if isinstance(weekend_entry, dict):
                        if "Sat" in weekend_entry:
                            w_worked, w_alloc = weekend_entry["Sat"]
                            sat_adj = float(w_alloc) - float(w_worked)
                        if "Sun" in weekend_entry:
                            w_worked, w_alloc = weekend_entry["Sun"]
                            sun_adj = float(w_alloc) - float(w_worked)
                        if sat_adj is None and sun_adj is None and "__all__" in weekend_entry:
                            w_worked, w_alloc = weekend_entry["__all__"]
                            total = float(w_alloc) - float(w_worked)
                            sat_adj = sun_adj = total / 2.0
                        if sat_adj is None and sun_adj is None:
                            # fallback: sum any tuple values and split
                            worked_sum = sum(v[0] for v in weekend_entry.values() if isinstance(v, (list, tuple)))
                            alloc_sum = sum(v[1] for v in weekend_entry.values() if isinstance(v, (list, tuple)))
                            total = float(alloc_sum) - float(worked_sum)
                            sat_adj = sun_adj = total / 2.0
                    else:
                        # tuple-style entry, split evenly across Sat/Sun
                        w_worked, w_alloc = weekend_entry
                        total = float(w_alloc) - float(w_worked)
                        sat_adj = sun_adj = total / 2.0
                except Exception:
                    sat_adj = sun_adj = None

            def _adj_html(adj):
                if adj is None:
                    return ""
                try:
                    adjf = float(adj)
                except Exception:
                    return ""
                sign_text = f"{adjf:+.2f}"
                # colour intensity scaled by magnitude (cap divisor at 3.0 to avoid overpowering)
                mag = min(1.0, abs(adjf) / 3.0)
                # ensure a minimum alpha so small adjustments remain visible
                if mag < 0.12:
                    mag = 0.12
                if adjf > 0:
                    bg = f"rgba(0,128,0,{mag:.2f})"
                    color = "#fff"
                elif adjf < 0:
                    bg = f"rgba(196,0,0,{mag:.2f})"
                    color = "#fff"
                else:
                    bg = "rgba(128,128,128,0.15)"
                    color = "#000"
                return f" <span class='weekend-day-adjust' style='margin-left:6px; padding:1px 6px; border-radius:4px; background:{bg}; color:{color}; font-weight:600; font-size:0.85em;'>({sign_text})</span>"

            sat_adj_html = _adj_html(sat_adj)
            sun_adj_html = _adj_html(sun_adj)

            worker_summary_html = (
                f"<div class='worker-summary auto-generated' style='margin-top:4px; font-size:0.9em; color:#333;'>"
                f"<span class='weekend-sat'>Sat: <strong>{sat_count}</strong>{sat_adj_html}</span>"
                f" &nbsp; <span class='weekend-sun'>Sun: <strong>{sun_count}</strong>{sun_adj_html}</span>"
                f"{adjustments_html}"
                f"</div>"
            )

            worker_td = """<td title='Site: {site}' class='worker {site}' 
            data-worker-id='{worker_id}'
            data-grade='{grade}'
            data-nwds='{nwds}' data-site='{site}' data-worker='{name}' 
            data-fte='{fte}' data-fte_adj='{fte_adj}' 
            data-start_date='{start_date}' 
            data-end_date='{end_date}' 
            data-oops='{oops}' 
            data-worker-targets='{targets}' 
            data-shift-counts='{worker_shift_counts}' 
            data-locum-shift-counts='{locum_shift_counts}'
            data-weekend-target='{weekend_target}' 
            data-remote-site='{remote_site}' 
            data-pair='{pair}'
            data-shift-balance-extra='{shift_balance_extra}'
            data-bank-holiday-extra='{bank_holiday_extra}'
            data-weekend-sat='{weekend_sat}'
            data-weekend-sun='{weekend_sun}'
            data-previous-shifts='{previous_shifts}'
            data-exact-shifts='{exact_shifts}'
            data-shift-diff='{shift_diff}'
            data-shift-diff-summed='{shift_diff_summed}'
            >
            <span class='name' title='{name} ({worker_id})'>{name}</span> ({grade}) [{fte}]{worker_summary_html}</td>""".format(
                site=worker.site,
                nwds=nwds,
                worker_id=worker.id,
                name=worker.name,
                fte=worker.fte,
                fte_adj=worker.get_fte(),
                start_date=worker.calculated_start_date,
                end_date=worker.calculated_end_date,
                oops=json.dumps(json.dumps(worker.oop, default=str)),
                targets=worker_targets,
                worker_shift_counts=json.dumps(shift_count_dict),
                locum_shift_counts=json.dumps(locum_shift_count_dict),
                weekend_target=worker.weekend_shift_target_number,
                remote_site=worker.remote_site,
                grade=worker.grade,
                weekend_sat=sat_count,
                weekend_sun=sun_count,
                worker_summary_html=worker_summary_html,
                previous_shifts=json.dumps(prior_map),
                exact_shifts=json.dumps(getattr(worker, "exact_shifts", {})),
                pair=worker.pair,
                shift_balance_extra=json.dumps(worker.shift_balance_extra),
                shift_diff=json.dumps(shift_diff_dict),
                shift_diff_summed=model.shift_count_diff_summed[worker.id].value,
                bank_holiday_extra=worker.bank_holiday_extra,
            )

            shift_count = (
                shift_count
                + f"#weekends_worked: {model.worker_weekend_count[worker.id].value}\\#"
            )

            if self.constraint_options_model.balance_bank_holidays:
                bank_holiday_count = model.bank_holiday_count[worker.id].value

                bank_holiday_count_w = model.bank_holiday_count_w[worker.id].value - 1
            else:
                bank_holiday_count = -1
                bank_holiday_count_w = -1

            # print(worker.name, bank_holiday_count, bank_holiday_count_w)
            timetable.append(
                f"<tr class='worker-row'>{worker_td}{''.join(shift_tds)}</tr>"
            )
            # timetable.append("<tr>{}</tr>".format("".join(w), shift_count))

            # if show_prefs:
            #     # prefs
            #     w = ["{:20}".format("Preferences")]
            #     for week, day in self.get_week_day_combinations():
            #         if model.pref_not_to_work[worker.id, week, day] > 0:
            #             w.append("Y")
            #         else:
            #             w.append("N")
            #     timetable.append("".join(w))

            # if show_unavailable:
            #     # prefs
            #     w = ["{:20}".format("Unavailable")]
            #     for week, day in self.get_week_day_combinations():
            #         if model.available[worker.id, week, day] > 0:
            #             w.append("A")
            #         else:
            #             w.append("U")
            #     timetable.append("".join(w))

        result_stream = StringIO()

        if self.results is not None:
            self.results.write(ostream=result_stream)
            self.results.write_json(ostream=result_stream)
            result_string = result_stream.getvalue()
        else:
            result_string = "Rota not run"

        # Wrap the remaining rows in a tbody and prepend the thead
        joined_timetable = header_row_html + "\n<tbody>\n" + "\n".join(timetable) + "\n</tbody>"

        hard_constrain_shift_info = []
        for worker in self.workers:
            for shift_name, limits in getattr(worker, "hard_constrain_shift_limits", {}).items():
                hard_constrain_shift_info.append({
                    "worker": worker.name,
                    "shift": shift_name,
                    **limits,
                })

        # Add warnings export
        warnings_html = "<details><summary><h2>Warnings</h2></summary><pre>\n"
        for warning_type, message in self.warnings:
            warnings_html += f"{warning_type}: {message}\n"
        warnings_html += "</pre></details>"


        # Human-readable worker details export
        worker_details_human = []
        for worker in self.get_all_workers():
            details = [
                f"Name: {worker.name}",
                f"ID: {worker.id}",
                f"Site: {worker.site}",
                f"Grade: {worker.grade}",
                f"FTE: {worker.fte}",
                f"Remote site: {getattr(worker, 'remote_site', '')}",
                f"Pair: {getattr(worker, 'pair', '')}",
                f"Start date: {getattr(worker, 'calculated_start_date', '')}",
                f"End date: {getattr(worker, 'calculated_end_date', '')}",
                f"Bank holiday extra: {getattr(worker, 'bank_holiday_extra', '')}",
                f"Shift balance extra: {getattr(worker, 'shift_balance_extra', '')}",
                f"Non-working days: {getattr(worker, 'non_working_day_list', '')}",
                f"OOP: {getattr(worker, 'oop', '')}",
                f"Unavailable: {''}",
                f"Shift targets: {json.dumps(getattr(worker, 'shift_target_number', {}))}",
            ]
            # compute Sat / Sun assigned counts for this worker
            sat_count = 0
            sun_count = 0
            for w, d in self.get_week_day_combinations():
                if d == "Sat":
                    for s in self.get_shift_names_by_week_day(w, d, return_empty_if_week_day_not_found=True):
                        if self.model.works[worker.id, w, d, s].value > 0.5:
                            sat_count += 1
                            break
                if d == "Sun":
                    for s in self.get_shift_names_by_week_day(w, d, return_empty_if_week_day_not_found=True):
                        if self.model.works[worker.id, w, d, s].value > 0.5:
                            sun_count += 1
                            break
            details.insert(
                -1,
                f"Sat assigned: {sat_count}, Sun assigned: {sun_count}",
            )
            # Build unavailability list for this worker
            unavail_list = []
            for entry in sorted(self.unavailable_to_work):
                wid, wk, d = entry
                if wid != worker.id:
                    continue
                date = self.get_date_by_week_day(wk, d)
                reason = self.unavailable_to_work_reason.get((wid, wk, d), "")
                if reason:
                    unavail_list.append(f"{date.isoformat()} ({d}) [week {wk}]: {reason}")
                else:
                    # date may be a string if fallback used
                    if isinstance(date, str):
                        unavail_list.append(f"{date} ({d}) [week {wk}]")
                    else:
                        unavail_list.append(f"{date.isoformat()} ({d}) [week {wk}]")

            # replace placeholder Unavailable entry
            details = [line for line in details if not line.startswith("Unavailable:")]
            details.insert(
                -1,  # before Shift targets
                f"Unavailable: {', '.join(unavail_list) if unavail_list else ''}",
            )
            worker_details_human.append("\n".join(details))
        worker_details_human_str = "\n\n".join(worker_details_human)


        # Collect force assigned shifts for export
        force_assigned_info = []
        for worker in self.get_all_workers():
            for assignment in getattr(worker, "forced_assignments", []):
                # Support both tuple and pydantic model
                if hasattr(assignment, "week") and hasattr(assignment, "day") and hasattr(assignment, "shift_name"):
                    week = assignment.week
                    day = assignment.day
                    shift_name = assignment.shift_name
                else:
                    week, day, shift_name = assignment
                date = self.get_date_by_week_day(week, day)
                force_assigned_info.append({
                    "worker": worker.name,
                    "week": week,
                    "day": day,
                    "date": str(date),
                    "shift": shift_name,
                })

        html = f""" 
        <body>
        <div class="table-div" id="{table_name}">
        {run_times_html}
        {heading_html}
        <table id="main-table">{joined_timetable}</table>
        </div>
        {warnings_html}
        <details>
        <summary><h2>Rota settings</h2></summary>
        <div>
        <pre>
        {json.dumps({k: _pydantic_to_dict(v) for k, v in self.constraint_options_model.model_dump().items()}, indent=4)}
        </pre>
        </details>
        </div>
        <details>
        <summary><h2>Shifts settings</h2></summary>
        <div id="shifts-container" data-shifts='{json.dumps([i.name for i in self.shifts])}' data-shifts-json='{json.dumps([_pydantic_to_dict(i) for i in self.shifts])}'>
        <div id="shifts-controls">
            <input type="text" id="shift-filter" placeholder="Filter shifts by name" style="width:100%; margin-bottom:8px;" />
        </div>
        <div id="shifts-table">
        {"<br/>".join(str(i) for i in self.shifts)}
        </div>
        </div>
        </details>
        <details>
        <summary><h2>Worker details</h2></summary>

        <div id="worker_details" data-workers='{json.dumps([json.loads(worker.model_dump_json()) for worker in self.get_all_workers()])}'>
        <span id="workers-json-target"></span>
        <pre>
        {worker_details_human_str}
        </pre>
        </div>
        </details>
        <details>
        <summary><h2>Output</h2></summary>
        <pre>
        {result_string}
        </pre>
        </details>
        <div>
        <details>
        <summary><h2>Export table</h2></summary>
        <div id="export-div">
        <div>
        </details>
        <details>
        <summary><h2>Shift timetables</h2></summary>
        <div id="shift-timetable-options">
        </div>
        <div id="shift-timetable-div">
        <div>
        </details>
        <details>
        <summary><h2>Extra</h2></summary>
        <div id="extra-options">
        </div>
        <div id="extra-div">
        <h3>Locum availability</h3>
        {"TEST"}
        {self.locum_availability_map}
        <br/>
        <details>
        <summary>Force assigned shifts</summary>
        <pre>
        {json.dumps(force_assigned_info, indent=4)}
        </pre>
        </details>
        <div>
        </details>
    <details><summary><h2>Hard Constrain Shift Min/Max</h2></summary>
    <pre>
     {json.dumps(hard_constrain_shift_info, indent=4)}
     "</pre></details>

        </div
        </body>
        """

        if include_html_tag:
            # html_static_prefix controls how CSS/JS are referenced:
            # - If html_static_prefix is None: try to use Django STATIC_URL when running inside Django, otherwise use relative links
            # - If html_static_prefix is a falsey value (e.g. empty string or False): use relative links ("timetable.css")
            # - If html_static_prefix is a string: use it as the prefix (ensuring trailing slash)
            static_prefix = None
            if html_static_prefix is not None:
                static_prefix = html_static_prefix

            if static_prefix is None:
                # Attempt to detect Django settings; if not present, fall back to relative links
                try:
                    from django.conf import settings

                    static_prefix = getattr(settings, "STATIC_URL", "")
                except Exception:
                    static_prefix = ""

            # If a prefix is provided and not empty, ensure trailing slash
            if isinstance(static_prefix, str) and static_prefix:
                if not static_prefix.endswith("/"):
                    static_prefix = static_prefix + "/"

            if static_prefix:
                css_href = f"{static_prefix}timetable.css"
                js_src = f"{static_prefix}timetable.js"
            else:
                # Relative links for standalone usage
                css_href = "timetable.css"
                js_src = "timetable.js"

            html = f"""<html>
        <head>
        <link rel="stylesheet" type="text/css" href="{css_href}">
        <link rel="stylesheet" type="text/css" href="https://code.jquery.com/ui/1.13.0/themes/base/jquery-ui.css">
        <script src="https://ajax.googleapis.com/ajax/libs/jquery/3.5.1/jquery.min.js"></script>
        <script src="https://code.jquery.com/ui/1.13.0/jquery-ui.min.js"></script>
        <script src="{js_src}" defer></script>
        </head>
            {html}</html>"""

        return html

    def get_shift_summary_dict(self):
        """"""
        timetable = {
            worker.name: {shift: "" for shift in self.get_shift_names()}
            for worker in self.workers
        }
        for worker in self.workers:
            for shift in self.get_shift_names():
                timetable[worker.name][shift] = round(
                    self.model.shift_count[worker.id, shift].value
                )

        return timetable

    def get_worker_shifts_by_date(self, worker: Worker) -> Dict[datetime.date, str]:
        shifts = {}

        n = 0
        for week, day in self.get_week_day_combinations():
            d = self.start_date + datetime.timedelta(n)
            n = n + 1

            temp = ""
            for shift in self.get_shift_names_by_week_day(week, day):
                if self.model.works[worker.id, week, day, shift].value > 0.5:
                    shifts[d] = shift

        return shifts

    def get_workers_total_shifts(self) -> Dict[str, int]:
        shifts = {}

        for worker in self.workers:
            shifts[worker.name] = len(
                [i for i in self.get_worker_shift_list(worker) if i != ""]
            )

        return shifts

    def get_shift_worker_list(self, shift_name: str) -> List[set[Worker]]:
        """
        Returns a list of workers assigned to a specific shift.
        
        Args:
            shift_name (str): The name of the shift.

        Returns:
            List[Worker]: List of workers assigned to the specified shift.
        """
        shift_list = []
        for week, day in self.get_week_day_combinations():
            workers: set = set()
            for worker in self.workers:
                if self.model.works[worker.id, week, day, shift_name].value > 0.5:
                    workers.add(worker)
            
            shift_list.append(workers)

        return shift_list


    def get_worker_shift_list(self, worker: Worker, include_locums=False, search_multiple_assignments=False) -> List:
        """
        This function returns a list of shifts assigned to a worker for each week and day.

        It was originally designed to return a list of shift names (when only a single shift per day was supported), but it can also return a set of shift names if `search_multiple_assignments` is True.
        """
        shifts = []

        for week, day in self.get_week_day_combinations():
            # d = self.start_date + datetime.timedelta(n)
            # n = n + 1

            if search_multiple_assignments:
                temp = set()
            else:
                temp = ""
            for shift in self.get_shift_names_by_week_day(week, day):
                if self.model.works[worker.id, week, day, shift].value > 0.90:
                    if search_multiple_assignments:
                        temp.add(shift)
                    else:
                        temp = shift
                elif (
                    include_locums
                    and self.model.locum_works[worker.id, week, day, shift].value > 0.90
                ):
                    if search_multiple_assignments:
                        temp.add(shift)
                    else:
                        temp = shift
            shifts.append(temp)

        return shifts

    def get_worker_shift_count(self, worker: Worker, shift_name: str) -> int:
        """Returns the number of shifts of a specific type assigned to a worker."""
        count = 0
        for week, day in self.get_week_day_combinations():
            if shift_name in self.get_shift_names_by_week_day(week, day):
                if self.model.works[worker.id, week, day, shift_name].value > 0.5:
                    count += 1
        return count

    def get_worker_days_worked(self, worker: Worker) -> int:
        """
        Returns the number of unique days on which the worker is assigned at least one shift.
        """
        days_worked = 0
        for week, day in self.get_week_day_combinations():
            for shift in self.get_shift_names_by_week_day(week, day):
                if self.model.works[worker.id, week, day, shift].value > 0.5:
                    days_worked += 1
                    break  # Only count each day once, even if multiple shifts
        return days_worked

    def get_worker_shift_list_string(self, worker: Worker) -> str:
        shifts = self.get_worker_shift_list(worker)

        # Convert shift to a string representation
        return "".join(
            [
                self.get_shift_by_name(i).get_display_char() if i != "" else "-"
                for i in shifts
            ]
        )

    def get_shift_summary(self):
        works = self.model.works
        timetable = {
            worker.get_details(): {shift: "" for shift in self.get_shift_names()}
            for worker in self.workers
        }
        # timetable = { worker.get_details() : { shift : "" for shift in ["truro_twilight"] } for worker in workers }
        t = []
        for worker in self.workers:
            l = []

            total_shifts = 0
            for shift in self.get_shift_names():
                # for shift in ["truro_twilight"]:
                c = [
                    works[worker.id, week, day, shift].value
                    for week in self.weeks
                    for day in days
                ].count(1)
                if c > 0:
                    l.append(f"{shift} ({c})")
                    total_shifts = total_shifts + c
                timetable[worker.get_details()][shift] = c
            t.append(f"{worker.get_full_details()} [{total_shifts}]: {', '.join(l)}")
        return "\n".join(t)

    def get_shift_summary_html(self):
        works = self.model.works
        timetable = {
            worker.get_details(): {
                shift.get_display_char(): "" for shift in self.get_shifts()
            }
            for worker in self.workers
        }
        # timetable = { worker.get_details() : { shift : "" for shift in ["truro_twilight"] } for worker in workers }
        t = []
        for worker in self.workers:
            l = []

            total_shifts = 0
            for shift in self.get_shift_names():
                # for shift in ["truro_twilight"]:
                c = [
                    works[worker.id, week, day, shift]
                    for week in self.weeks
                    for day in days
                ].count(1)
                if c > 0:
                    l.append(f"{shift} ({c})")
                    total_shifts = total_shifts + c
                shift_obj = self.get_shift_by_name(shift)
                timetable[worker.get_details()][shift_obj.get_display_char()] = c
            t.append(f"{worker.get_full_details()} [{total_shifts}]: {', '.join(l)}")
        return "\n".join(t)


class NoActiveSites(Exception):
    """Raised when there are no active sites"""

    pass


class NoWorkers(Exception):
    """Raised when there are no active sites"""

    pass


class InvalidShift(Exception):
    """Raised when there are no active sites"""

    pass


class WarningTermination(Exception):
    """Raised when a warning in the termination group is raised"""

    pass

class WeekDayNotFound(Exception):
    """Raised when a week/day combination is not found"""

    pass