From 1ef25924a0e0ca1624ef0ef48d0e6c59e5b94065 Mon Sep 17 00:00:00 2001
From: Yohann D'ANELLO <ynerant@crans.org>
Date: Thu, 16 Sep 2021 20:46:34 +0200
Subject: [PATCH] [WEI] Display status bar with tqdm

Signed-off-by: Yohann D'ANELLO <ynerant@crans.org>
---
 apps/wei/forms/surveys/wei2021.py             | 30 +++++++++++++++----
 apps/wei/management/commands/wei_algorithm.py | 16 ++++++++--
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/apps/wei/forms/surveys/wei2021.py b/apps/wei/forms/surveys/wei2021.py
index 8c1e2945..a046b7ea 100644
--- a/apps/wei/forms/surveys/wei2021.py
+++ b/apps/wei/forms/surveys/wei2021.py
@@ -171,7 +171,7 @@ class WEISurveyAlgorithm2021(WEISurveyAlgorithm):
     def get_bus_information_class(cls):
         return WEIBusInformation2021
 
-    def run_algorithm(self):
+    def run_algorithm(self, display_tqdm=False):
         """
         Gale-Shapley algorithm implementation.
         We modify it to allow buses to have multiple "weddings".
@@ -196,11 +196,26 @@ class WEISurveyAlgorithm2021(WEISurveyAlgorithm):
             free_seats = bus.size - WEIMembership.objects.filter(bus=bus, registration__first_year=False).count()
             quotas[bus] = 4 + int(non_men_total / registrations.count() * free_seats)
 
-        # Repartition for non men people first
-        self.make_repartition(non_men, quotas)
-        self.make_repartition(men)
+        tqdm_obj = None
+        if display_tqdm:
+            from tqdm import tqdm
+            tqdm_obj = tqdm(total=len(non_men), desc="Non-hommes")
 
-    def make_repartition(self, surveys, quotas=None):
+        # Repartition for non men people first
+        self.make_repartition(non_men, quotas, tqdm_obj=tqdm_obj)
+
+        if display_tqdm:
+            tqdm_obj.close()
+
+            from tqdm import tqdm
+            tqdm_obj = tqdm(total=len(men), desc="Hommes")
+
+        self.make_repartition(men, tqdm_obj=tqdm_obj)
+
+        if display_tqdm:
+            tqdm_obj.close()
+
+    def make_repartition(self, surveys, quotas=None, tqdm_obj=None):
         free_surveys = surveys.copy()  # Remaining surveys
         while free_surveys:  # Some students are not affected
             survey = free_surveys[0]
@@ -235,6 +250,11 @@ class WEISurveyAlgorithm2021(WEISurveyAlgorithm):
                         free_surveys.append(least_preferred_survey)
                         survey.select_bus(bus)
                         survey.save()
+                        free_surveys.remove(survey)
                         break
             else:
                 raise ValueError(f"User {survey.registration.user} has no free seat")
+
+            if tqdm_obj is not None:
+                tqdm_obj.n = len(surveys) - len(free_surveys)
+                tqdm_obj.refresh()
diff --git a/apps/wei/management/commands/wei_algorithm.py b/apps/wei/management/commands/wei_algorithm.py
index 238bf13c..93f919f0 100644
--- a/apps/wei/management/commands/wei_algorithm.py
+++ b/apps/wei/management/commands/wei_algorithm.py
@@ -24,7 +24,14 @@ class Command(BaseCommand):
         sid = transaction.savepoint()
 
         algorithm = CurrentSurvey.get_algorithm_class()()
-        algorithm.run_algorithm()
+
+        try:
+            from tqdm import tqdm
+            display_tqdm = True
+        except ImportError:
+            display_tqdm = False
+
+        algorithm.run_algorithm(display_tqdm=display_tqdm)
 
         output = options['output']
         registrations = algorithm.get_registrations()
@@ -34,8 +41,13 @@ class Command(BaseCommand):
         for bus, members in per_bus.items():
             output.write(bus.name + "\n")
             output.write("=" * len(bus.name) + "\n")
+            order = -1
             for r in members:
-                output.write(r.user.username + "\n")
+                survey = CurrentSurvey(r)
+                for order, (b, _score) in enumerate(survey.ordered_buses()):
+                    if b == bus:
+                        break
+                output.write(f"{r.user.username} ({order + 1})\n")
             output.write("\n")
 
         if not options['doit']: