{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 手术调度问题\n",
    "\n",
    "**类别：** 排程\n",
    "\n",
    "使用 OptAgent 的 Python 接口描述变量、约束与目标。\n",
    "\n",
    "问题与原始示例来源：[Hexaly Code Templates](https://www.hexaly.com/templates/surgery-scheduling-problem)。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 问题描述\n",
    "\n",
    "**在手术调度问题中**，我们考虑一家医院，其拥有固定数量的可用手术室以及一组待安排的手术。每台手术具有给定的处理时间和所需的护士数量。一台手术只有在手术室可用且有足够的护士时才可开始。此外，对护士的班次还有最早开始时间、最晚结束时间以及最长持续时间的约束。手术调度问题的目标是寻找一种手术顺序，使 makespan（即所有手术处理完成的时间）最小化。\n",
    "\n",
    "### 学习要点\n",
    "\n",
    "- 使用 OptAgent 的 `interval` 决策变量表示手术时间\n",
    "- 使用 `list` 决策变量表示手术室和护士的手术顺序\n",
    "- 使用集合 lambda 将顺序决策与 interval 不重叠约束关联起来\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据\n",
    "\n",
    "数据文件的格式如下：\n",
    "\n",
    "- 第一行：手术室数量、护士数量、手术数量\n",
    "- 第二行：每台手术的最早开始时间（单位：小时）\n",
    "- 第三行：每台手术的最晚结束时间（单位：小时）\n",
    "- 第四行：每台手术的持续时间（单位：分钟）\n",
    "- 第五行：每台手术所需的护士数量\n",
    "- 第六行：每位护士班次的最早开始时间（单位：小时）\n",
    "- 第七行：每位护士班次的最晚结束时间（单位：小时）\n",
    "- 第八行：班次的最大持续时长（单位：小时）\n",
    "- 接下来的每一行给出一台手术与各手术室的不兼容标记，1 表示不兼容，0 表示兼容\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 建模思路\n",
    "\n",
    "每台手术使用一个具有开始、结束和固定持续时间的 interval；每个手术室和每位护士分别使用一个 list 表示其手术顺序。手术室 lists 通过 `partition` 约束使每台手术恰好分配到一个兼容手术室，相邻 interval 的顺序约束保证手术室内不重叠。\n",
    "\n",
    "护士 lists 可以共同包含同一手术。模型通过 `contains` 汇总参与每台手术的护士数量，并对每位护士的首台和末台手术约束班次起止时间及最大持续时长；护士序列中的相邻手术同样不得重叠。目标是最小化所有手术结束时间的最大值。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Python 实现\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from optagent import OptModel, solve\n",
    "\n",
    "\n",
    "def read_instance(filename):\n",
    "    lines = Path(filename).read_text(encoding=\"utf-8\").splitlines()\n",
    "    num_rooms, num_nurses, num_surgeries = map(int, lines[0].split())\n",
    "\n",
    "    min_start = [int(value) * 60 for value in lines[1].split()]\n",
    "    max_end = [int(value) * 60 for value in lines[2].split()]\n",
    "    duration = [int(value) for value in lines[3].split()]\n",
    "    needed_nurses = [int(value) for value in lines[4].split()]\n",
    "    shift_earliest_start = [int(value) * 60 for value in lines[5].split()]\n",
    "    shift_latest_end = [int(value) * 60 for value in lines[6].split()]\n",
    "    max_shift_duration = int(lines[7].split()[0]) * 60\n",
    "    incompatible_rooms = [\n",
    "        [int(value) for value in lines[8 + surgery].split()]\n",
    "        for surgery in range(num_surgeries)\n",
    "    ]\n",
    "\n",
    "    return (\n",
    "        num_rooms,\n",
    "        num_nurses,\n",
    "        num_surgeries,\n",
    "        min_start,\n",
    "        max_end,\n",
    "        needed_nurses,\n",
    "        shift_earliest_start,\n",
    "        shift_latest_end,\n",
    "        max_shift_duration,\n",
    "        incompatible_rooms,\n",
    "        duration,\n",
    "    )\n",
    "\n",
    "\n",
    "def main(input_file, output_file=None, time_limit=20):\n",
    "    (\n",
    "        num_rooms,\n",
    "        num_nurses,\n",
    "        num_surgeries,\n",
    "        min_start,\n",
    "        max_end,\n",
    "        needed_nurses,\n",
    "        shift_earliest_start,\n",
    "        shift_latest_end,\n",
    "        max_shift_duration,\n",
    "        incompatible_rooms,\n",
    "        duration,\n",
    "    ) = read_instance(input_file)\n",
    "\n",
    "    model = OptModel()\n",
    "\n",
    "    # Each surgery is assigned to exactly one compatible room.\n",
    "    surgery_order = [\n",
    "        model.list(num_surgeries)\n",
    "        for room in range(num_rooms)\n",
    "    ]\n",
    "    rooms = model.array(surgery_order)\n",
    "    model.constraint(model.partition(rooms))\n",
    "    for surgery in range(num_surgeries):\n",
    "        for room in range(num_rooms):\n",
    "            if incompatible_rooms[surgery][room]:\n",
    "                model.constraint(\n",
    "                    model.contains(surgery_order[room], surgery) == 0,\n",
    "                )\n",
    "\n",
    "    selected_room = [model.find(rooms, surgery) for surgery in range(num_surgeries)]\n",
    "\n",
    "    surgeries = [\n",
    "        model.interval(\n",
    "            min_start[surgery],\n",
    "            max_end[surgery],\n",
    "        )\n",
    "        for surgery in range(num_surgeries)\n",
    "    ]\n",
    "    for surgery in range(num_surgeries):\n",
    "        model.constraint(\n",
    "            surgeries[surgery].length() == duration[surgery],\n",
    "        )\n",
    "    surgery_array = model.array(surgeries)\n",
    "\n",
    "    def make_no_overlap_lambda(sequence):\n",
    "        return model.lambda_function(\n",
    "            lambda position: surgery_array[sequence[position]]\n",
    "            < surgery_array[sequence[position + 1]]\n",
    "        )\n",
    "\n",
    "    # A room can process only one surgery at a time.\n",
    "    for room, sequence in enumerate(surgery_order):\n",
    "        adjacent_positions = model.range(0, sequence.count() - 1)\n",
    "        model.constraint(\n",
    "            model.and_(\n",
    "                adjacent_positions, make_no_overlap_lambda(sequence)\n",
    "            ),\n",
    "        )\n",
    "\n",
    "    nurse_order = [\n",
    "        model.list(num_surgeries)\n",
    "        for nurse in range(num_nurses)\n",
    "    ]\n",
    "    for nurse, sequence in enumerate(nurse_order):\n",
    "        count = sequence.count()\n",
    "        first_surgery_start = model.iif(\n",
    "            count > 0,\n",
    "            surgery_array[sequence[0]].start,\n",
    "            shift_earliest_start[nurse],\n",
    "        )\n",
    "        last_surgery_end = model.iif(\n",
    "            count > 0,\n",
    "            surgery_array[sequence[count - 1]].end,\n",
    "            shift_earliest_start[nurse],\n",
    "        )\n",
    "        model.constraint(\n",
    "            first_surgery_start >= shift_earliest_start[nurse],\n",
    "        )\n",
    "        model.constraint(\n",
    "            last_surgery_end <= shift_latest_end[nurse],\n",
    "        )\n",
    "        model.constraint(\n",
    "            last_surgery_end - first_surgery_start <= max_shift_duration,\n",
    "        )\n",
    "\n",
    "        adjacent_positions = model.range(0, count - 1)\n",
    "        model.constraint(\n",
    "            model.and_(\n",
    "                adjacent_positions, make_no_overlap_lambda(sequence)\n",
    "            ),\n",
    "        )\n",
    "\n",
    "    nurse_order_array = model.array(nurse_order)\n",
    "\n",
    "    def make_nurse_assignment_lambda(surgery):\n",
    "        return model.lambda_function(\n",
    "            lambda nurse: model.contains(nurse_order_array[nurse], surgery)\n",
    "        )\n",
    "\n",
    "    for surgery in range(num_surgeries):\n",
    "        assigned_nurses = model.sum(\n",
    "            model.range(0, num_nurses),\n",
    "            make_nurse_assignment_lambda(surgery),\n",
    "        )\n",
    "        model.constraint(\n",
    "            assigned_nurses >= needed_nurses[surgery],\n",
    "        )\n",
    "\n",
    "    makespan = model.max([surgery.end for surgery in surgeries])\n",
    "    model.minimize(makespan)\n",
    "\n",
    "    solution = solve(model, time_limit_s=float(time_limit))\n",
    "    if not solution.feasible:\n",
    "        print(f\"No feasible surgery schedule found; Status = {solution.status}\")\n",
    "        return solution\n",
    "\n",
    "    nurses_by_surgery = [[] for _ in range(num_surgeries)]\n",
    "    for nurse, sequence in enumerate(nurse_order):\n",
    "        for surgery in sequence.value:\n",
    "            nurses_by_surgery[int(surgery)].append(nurse)\n",
    "\n",
    "    output_lines = []\n",
    "    for surgery in range(num_surgeries):\n",
    "        interval = surgeries[surgery].value\n",
    "        output_lines.append(\n",
    "            f\"{surgery}\\t\\t{int(selected_room[surgery].value)}\"\n",
    "            f\"\\t{interval.start()}\\t{interval.end()}\"\n",
    "            f\"\\t{nurses_by_surgery[surgery]}\"\n",
    "        )\n",
    "\n",
    "    print(\n",
    "        f\"Makespan = {makespan.value}; Status = {solution.status}\\n\"\n",
    "        + \"\\n\".join(output_lines)\n",
    "    )\n",
    "    if output_file is not None:\n",
    "        Path(output_file).write_text(\n",
    "            \"\\n\".join(output_lines) + \"\\n\", encoding=\"utf-8\"\n",
    "        )\n",
    "    return solution\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INSTANCE_DIR = Path.cwd() / \"instances\"\n",
    "print(\"Instances:\", INSTANCE_DIR)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solution = main(\n",
    "    INSTANCE_DIR / \"instancesurgery.txt\",\n",
    "    time_limit=1,\n",
    ")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
