add /best, /show, /endplan

dudlebot.py

@@ -4,6 +4,7 @@ import re
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from enum import Enum
+from operator import itemgetter
 from typing import List, Dict
 
 import telegram
@@ -22,11 +23,18 @@ class DudleBot:
         dispatcher.add_handler(CommandHandler('start', self.tg_start))
         dispatcher.add_handler(CommandHandler('newplan', self.tg_newplan))
         dispatcher.add_handler(CommandHandler('plan', self.tg_plan))
+        dispatcher.add_handler(CommandHandler('show', self.tg_show))
+        dispatcher.add_handler(CommandHandler('best', self.tg_best))
+        dispatcher.add_handler(CommandHandler('endplan', self.tg_end))
 
     def tg_start(self, update, _):
         update.message.reply_text("Hello! May luck be with you and your plans!")
 
     def tg_newplan(self, update, context):
+        if context.chat_data.get("plan", None):
+            update.message.reply_text("There's a plan unresolved still! (Speak /endplan to fix it.)")
+            return
+
         periods = {
             'thisweek': datetime.today() - timedelta(days=datetime.today().weekday() % 6),
             'nextweek': datetime.today() + timedelta(days=7) - timedelta(days=datetime.today().weekday() % 6)
@@ -78,6 +86,33 @@ class DudleBot:
 
         self._reply_with_plan(update, context)
 
+    def tg_show(self, update, context):
+        if not context.chat_data.get("plan", None):
+            update.message.reply_text("No plan created yet! (Speak /newplan to do so.)")
+            return
+
+        self._reply_with_plan(update, context)
+
+    def tg_best(self, update, context):
+        if not context.chat_data.get("plan", None):
+            update.message.reply_text("No plan created yet! (Speak /newplan to do so.)")
+            return
+
+        update.message.reply_text(self._best_of_plan(context.chat_data["plan"]))
+
+    def tg_end(self, update, context):
+        if not context.chat_data.get("plan", None):
+            update.message.reply_text("No plan created yet! (Speak /newplan to do so.)")
+            return
+
+        plan = context.chat_data.pop("plan")
+        response = "Poll ended!"
+
+        if len(plan.entries) > 0:
+            response += "\n\n" + self._best_of_plan(plan)
+
+        update.message.reply_text(response)
+
     @staticmethod
     def _reply_with_plan(update, context):
         plan: Plan = context.chat_data['plan']
@@ -110,6 +145,37 @@ class DudleBot:
 
         update.message.reply_text(formatted_plan, parse_mode=telegram.ParseMode.MARKDOWN)
 
+    @staticmethod
+    def _best_of_plan(plan):
+        days = [plan.start + timedelta(days=i) for i in range(plan.duration.days)]
+        results = [0] * plan.duration.days
+        for entry in plan.entries.values():
+            for idx, response in enumerate(entry.responses):
+                if response == PlanResponse.YES:
+                    results[idx] += 1
+                # elif response == PlanResponse.UNKNOWN:
+                #     results[idx] += 0.01
+
+        sorted_days = sorted(zip(days, results), key=itemgetter(1), reverse=True)
+        best_days = []
+
+        last_result, different_results = None, 0
+        for day, result in sorted_days:
+            if result == 0 or (len(best_days) >= 3 and different_results > 1):
+                break
+
+            best_days.append((day, result))
+
+            if result != last_result:
+                different_results += 1
+                last_result = result
+
+        response = f"Best {len(best_days)} result{'s' if len(best_days) > 1 else ''}:\n\n"
+        for idx, dayresult in enumerate(best_days):
+            day, result = dayresult
+            response += f"{idx + 1}. {day.strftime('%A')} ({result})\n"
+        return response
+
     def start(self):
         self.logger.info("Starting DudleBot...")
         self.updater.start_polling()