ENCOA-312

ielts_be/helpers/difficulty.py

@@ -0,0 +1,40 @@
+import math
+import random
+from typing import Optional, List, Iterator
+
+from ielts_be.configs.constants import EducationalContent
+
+
+class DifficultyHelper:
+
+    def __init__(self, difficulties: Optional[List[str]]):
+        self.difficulties = difficulties
+        self.distributed: Optional[Iterator[str]] = None
+
+    def distribute_for_count(self, count: int) -> None:
+        if not self.difficulties or count == 0:
+            return
+
+        result = []
+        remaining = count
+        difficulties_count = len(self.difficulties)
+
+        for i, diff in enumerate(self.difficulties):
+            if i == difficulties_count - 1:
+                slots = remaining
+            else:
+                slots = math.ceil(remaining / (difficulties_count - i))
+
+            result.extend([diff] * slots)
+            remaining -= slots
+
+        self.distributed = iter(result)
+
+    def pick_difficulty(self, difficulty: Optional[str]) -> str:
+        if difficulty:
+            return difficulty if difficulty != "Random" else random.choice(EducationalContent.DIFFICULTIES)
+
+        if self.distributed:
+            return next(self.distributed)
+
+        return random.choice(EducationalContent.DIFFICULTIES)