From 0514f80a77c839025cdf068db4f06f961799601e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Th=C3=A9ophile=20Bastian?= <contact@tobast.fr>
Date: Sat, 4 Mar 2023 11:45:48 +0100
Subject: [PATCH] Partition tasks using toughness score

---
 repartir_taches/entrypoint.py | 103 ++++++++++++++++++++++++++++------
 1 file changed, 87 insertions(+), 16 deletions(-)

diff --git a/repartir_taches/entrypoint.py b/repartir_taches/entrypoint.py
index b8bde7a..0b11802 100644
--- a/repartir_taches/entrypoint.py
+++ b/repartir_taches/entrypoint.py
@@ -2,13 +2,19 @@ import argparse
 import typing as t
 import random
 from pathlib import Path
+import logging
 import jinja2 as j2
+import prtpy
 
 from .config import Task, Category, Config
 from . import util
 
+logger = logging.getLogger(__name__)
 
-def constituer_groupes(choristes: list[str]) -> list[list[str]]:
+Group: t.TypeAlias = list[str]
+
+
+def constituer_groupes(choristes: list[str]) -> list[Group]:
     """Répartir aléatoirement les gens en groupes"""
     TAILLE_GROUPE: int = 4
     nb_choristes = len(choristes)
@@ -36,19 +42,67 @@ def constituer_groupes(choristes: list[str]) -> list[list[str]]:
     return groupes
 
 
-def assigner_taches(task: Category | Task, group_count: int, cur_group: int = 0) -> int:
+class AssignError(Exception):
+    """Levé lorsque les tâches ont été mal partitionnées"""
+
+
+def assigner_taches(root_task: Category | Task, group_count: int):
     """Assigne les tâches aux groupes (round-robin)"""
-    if isinstance(task, Task):
-        task.assigned = list(
-            map(
-                lambda x: x % group_count,
-                range(cur_group, cur_group + task.nb_groups),
-            )
-        )
-        return (cur_group + task.nb_groups) % group_count
-    for subtask in task.tasks:
-        cur_group = assigner_taches(subtask, group_count, cur_group)
-    return cur_group
+
+    TaskId = t.NewType("TaskId", int)
+    UniqueTask: t.TypeAlias = tuple[TaskId, int]
+
+    def flatten(task: Category | Task) -> list[Task]:
+        if isinstance(task, Task):
+            return [task]
+        out = []
+        for subtask in task.tasks:
+            out += flatten(subtask)
+        return out
+
+    all_tasks = flatten(root_task)
+
+    def pp_assigned_toughness(repart: list[list[UniqueTask]]) -> str:
+        """Pretty-print the assigned toughness for each group"""
+        out = []
+        for grp_id, grp in enumerate(repart):
+            toughness: int = sum(map(lambda x: all_tasks[x[0]].tough, grp))
+            out.append(f"{grp_id:2d}: {toughness:>3d}")
+        return "\n".join(out)
+
+    opt_input: dict[UniqueTask, int] = {}
+    for task_id, task in enumerate(all_tasks):
+        for rep in range(task.nb_groups):
+            opt_input[(TaskId(task_id), rep)] = task.tough
+    repart: list[list[UniqueTask]] = prtpy.partition(
+        algorithm=prtpy.partitioning.greedy,
+        numbins=group_count,
+        items=opt_input,
+    )
+
+    # Sanity-check
+    for g_id, grp in enumerate(repart):
+        taskset: set[TaskId] = set()
+        for (task_id, _) in grp:
+            if task_id in taskset:
+                raise AssignError(
+                    f"Le groupe {g_id + 1} a deux fois la tâche {task.qualified_name}"
+                )
+            taskset.add(task_id)
+
+    # Actually assign
+    for g_id, grp in enumerate(repart):
+        for (task_id, _) in grp:
+            task = all_tasks[task_id]
+            if task.assigned is None:
+                task.assigned = [g_id]
+            else:
+                task.assigned.append(g_id)
+
+    logger.debug("Assigned toughness mapping:\n%s", pp_assigned_toughness(repart))
+
+    # Check that all tasks are assigned
+    assert all(map(lambda x: x.assigned is not None and x.assigned, all_tasks))
 
 
 def export_short_md(config: Config, groupes: list[list[str]]) -> str:
@@ -115,7 +169,15 @@ def export_latex(config: Config, groupes: list[list[str]]) -> str:
     return template.render(**env)
 
 
-def main():
+def repartition(config: Config) -> list[Group]:
+    """Crée des groupes et assigne des tâches"""
+    groupes: list[Group] = constituer_groupes(config.choristes)
+    assigner_taches(config.taches, len(groupes))
+
+    return groupes
+
+
+def main() -> None:
     parser = argparse.ArgumentParser("Répartition des tâches")
     parser.add_argument("taches", help="Fichier yaml contenant les tâches")
     parser.add_argument("choristes", help="Fichier CSV contenant les choristes")
@@ -130,15 +192,24 @@ def main():
         "--to-short-md",
         help="Exporter vers un fichier Markdown (pour vérification uniquement)",
     )
+    parser.add_argument(
+        "-g",
+        "--debug",
+        dest="loglevel",
+        action="store_const",
+        const=logging.DEBUG,
+        default=logging.INFO,
+    )
     args = parser.parse_args()
 
+    logging.basicConfig(level=args.loglevel)
+
     config = Config(args.taches, args.choristes)
 
     if args.bare_tasks:
         util.write_to_file(args.bare_tasks, export_bare_tasks_md(config))
 
-    groupes = constituer_groupes(config.choristes)
-    assigner_taches(config.taches, len(groupes))
+    groupes = repartition(config)
 
     if args.to_tex:
         util.write_to_file(args.to_tex, export_latex(config, groupes))