app/engine/solver.py

"""
Constraint-Satisfaction Solver for Smart Scheduling.

Algorithm
---------
1.  Load all existing Confirmed + Locked sessions from DB → build OccupationMap.
2.  Walk forward day-by-day from start_date, visiting only days in `days` list.
3.  On each date, try slots in heuristic order (load-balance + clustering).
4.  For each candidate (date, slot):
      Hard constraints:
        a. Mentor is free on that (date, slot)
        b. All enrolled students are free on that (date, slot)
        c. At least one room is available on that (date, slot)
5.  First passing candidate → add to proposals & update occupation map.
6.  Stop when `to_place` proposals are collected or horizon is exhausted.
7.  Return placed proposals + unplaced count + reasons.
"""

from datetime import date, timedelta

from app.constants import (
    ABBR_TO_WEEKDAY,
    SLOT_DEFINITIONS,
    SESSION_STATUS_CONFIRMED,
    SESSION_STATUS_LOCKED,
)
from app.engine.heuristics import order_slots_for_mentor, prefer_adjacent_for_students
from app.logging_config import get_logger
from app.schemas.ingest import AutoGenerateResponse, ScheduleProposal

logger = get_logger(__name__)

# Weekday number → abbreviation (for readable log messages)
_WEEKDAY_ABBR = {v: k for k, v in ABBR_TO_WEEKDAY.items()}

# Type alias for OccupationMap
OccupationMap = dict[tuple[date, int], dict]   # (date, slot_id) → {"mentor_ids", "room_ids", "student_ids"}


def _build_occupation_map(existing_sessions: list) -> OccupationMap:
    occ: OccupationMap = {}
    for s in existing_sessions:
        if s.status not in (SESSION_STATUS_CONFIRMED, SESSION_STATUS_LOCKED):
            continue
        key = (s.date, s.slot_id)
        if key not in occ:
            occ[key] = {"mentor_ids": set(), "room_ids": set(), "student_ids": set()}
        cell = occ[key]
        # mentor is on the batch
        if s.batch and s.batch.mentor_id:
            cell["mentor_ids"].add(s.batch.mentor_id)
        if s.room_id:
            cell["room_ids"].add(s.room_id)
        # enrolled students
        if s.batch:
            for student in s.batch.students:
                cell["student_ids"].add(student.id)
    return occ


def find_optimal_slots(
    *,
    batch_id: str,
    batch_name: str,
    mentor_id: str,
    mentor_name: str,
    mentor_off_days: list[str],
    enrolled_student_ids: set[str],
    start_date: date,
    days: list[str],                   # e.g. ["Sat", "Sun"]
    num_sessions: int,                 # how many to place
    rooms: list,                       # list of Room ORM objects
    existing_sessions: list,           # all Session ORM objects (with batch + students loaded)
    preferred_slot_id: int | None = None,  # try this slot first each day
) -> AutoGenerateResponse:

    to_place = num_sessions
    placed: list[ScheduleProposal] = []
    reasons: list[str] = []

    target_weekdays = {ABBR_TO_WEEKDAY[d] for d in days if d in ABBR_TO_WEEKDAY}
    mentor_off_weekdays = {ABBR_TO_WEEKDAY[d] for d in mentor_off_days if d in ABBR_TO_WEEKDAY}

    occ = _build_occupation_map(existing_sessions)

    logger.info(
        "solver.start batch=%s mentor=%s students=%d need=%d start=%s days=%s preferred_slot=%s",
        batch_name, mentor_name, len(enrolled_student_ids),
        to_place, start_date, days, preferred_slot_id,
    )
    logger.debug(
        "solver.occupation_map size=%d (existing confirmed+locked sessions loaded)",
        len(occ),
    )

    # Walk forward up to 2 years (safety cap — solver will stop at to_place)
    MAX_DAYS_AHEAD = 730
    current_date = start_date

    while to_place > 0 and (current_date - start_date).days <= MAX_DAYS_AHEAD:
        wd = current_date.weekday()

        # Skip days not in requested set
        if wd not in target_weekdays:
            current_date += timedelta(days=1)
            continue

        # Skip mentor's off days
        if wd in mentor_off_weekdays:
            logger.debug("solver.skip date=%s reason=mentor_off_day weekday=%s",
                         current_date, _WEEKDAY_ABBR.get(wd, wd))
            current_date += timedelta(days=1)
            continue

        # Determine slot order using heuristics
        slot_order = order_slots_for_mentor(mentor_id, current_date, occ)
        slot_order = prefer_adjacent_for_students(
            enrolled_student_ids, current_date, occ, slot_order
        )
        # If a preferred slot was requested, promote it to the front
        if preferred_slot_id and preferred_slot_id in slot_order:
            slot_order = [preferred_slot_id] + [s for s in slot_order if s != preferred_slot_id]

        logger.debug("solver.try_date date=%s slot_order=%s", current_date, slot_order)

        placed_today = False
        slot_block_reasons: list[str] = []

        for slot_id in slot_order:
            cell = occ.get((current_date, slot_id), {})

            # Hard constraint a: mentor free
            if mentor_id in cell.get("mentor_ids", set()):
                msg = f"slot {slot_id}: mentor {mentor_name} already has a session"
                slot_block_reasons.append(msg)
                logger.debug("solver.slot_blocked date=%s %s", current_date, msg)
                continue

            # Hard constraint b: all enrolled students free
            busy_students = enrolled_student_ids & cell.get("student_ids", set())
            if busy_students:
                msg = f"slot {slot_id}: {len(busy_students)} enrolled student(s) have a conflict"
                slot_block_reasons.append(msg)
                logger.debug("solver.slot_blocked date=%s %s student_ids=%s",
                             current_date, msg, busy_students)
                continue

            # Hard constraint c: at least one room available
            occupied_rooms = cell.get("room_ids", set())
            free_room = next((r for r in rooms if r.id not in occupied_rooms), None)
            if not free_room:
                msg = f"slot {slot_id}: all {len(rooms)} room(s) occupied"
                slot_block_reasons.append(msg)
                logger.debug("solver.slot_blocked date=%s %s", current_date, msg)
                reasons.append(f"{current_date} slot {slot_id}: no rooms available")
                continue

            # ── All constraints satisfied — record proposal ──
            slot_start, slot_end = SLOT_DEFINITIONS[slot_id]
            placed.append(ScheduleProposal(
                batch_id=batch_id,
                batch_name=batch_name,
                mentor_id=mentor_id,
                mentor_name=mentor_name,
                room_id=free_room.id,
                room_name=free_room.name,
                date=current_date,
                slot_id=slot_id,
                slot_start=slot_start,
                slot_end=slot_end,
            ))

            # Update occupation map so subsequent proposals on same date are aware
            key = (current_date, slot_id)
            if key not in occ:
                occ[key] = {"mentor_ids": set(), "room_ids": set(), "student_ids": set()}
            occ[key]["mentor_ids"].add(mentor_id)
            occ[key]["room_ids"].add(free_room.id)
            occ[key]["student_ids"].update(enrolled_student_ids)

            to_place -= 1
            placed_today = True
            logger.info(
                "solver.placed batch=%s date=%s slot=%d room=%s remaining=%d",
                batch_name, current_date, slot_id, free_room.name, to_place,
            )
            break   # one session per day per batch

        if not placed_today:
            reason_summary = "; ".join(slot_block_reasons) if slot_block_reasons else "no slots tried"
            reasons.append(f"{current_date}: no valid slot found — {reason_summary}")
            logger.warning(
                "solver.date_blocked batch=%s date=%s reasons=[%s]",
                batch_name, current_date, reason_summary,
            )

        current_date += timedelta(days=1)

    if to_place > 0:
        logger.warning(
            "solver.incomplete batch=%s placed=%d could_not_place=%d "
            "— check solver.date_blocked lines above for per-day reasons",
            batch_name, num_sessions - to_place, to_place,
        )
    else:
        logger.info(
            "solver.complete batch=%s all %d sessions placed",
            batch_name, num_sessions,
        )

    return AutoGenerateResponse(
        placed=placed,
        could_not_place=to_place,
        reasons=list(dict.fromkeys(reasons)),   # deduplicate, preserve order
    )