From 0e3011c406a61471dfa1435352d9faf071b033e5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zo=C3=A9=20Cassiop=C3=A9e=20Gauthier?= <hello@blorp.dev>
Date: Sat, 30 Mar 2024 23:17:34 -0400
Subject: [PATCH] Allow attackers of any types, but with a charged move of the
 specified type

---
 src/pogo_scaled_estimators/calculator.py | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/pogo_scaled_estimators/calculator.py b/src/pogo_scaled_estimators/calculator.py
index 223af8b..14d6293 100644
--- a/src/pogo_scaled_estimators/calculator.py
+++ b/src/pogo_scaled_estimators/calculator.py
@@ -55,11 +55,12 @@ class Calculator:
             f"""SELECT DISTINCT(attacker) FROM estimators WHERE charged_move IN ("{'","'.join(charged_moves)}")"""
         )
         attackers = {row[0]: {"RAID_LEVEL_3": [], "RAID_LEVEL_5": [], "RAID_LEVEL_MEGA": []} for row in res.fetchall()}
+        movesets = {}
 
         defenders = functools.reduce(operator.iconcat, raid_bosses.values(), [])
         res = self.db.execute(
             f"""
-            SELECT e.raid_tier, e.defender, e.attacker, MIN(e.estimator) / m.min_estimator
+            SELECT e.raid_tier, e.defender, e.attacker, e.estimator / m.min_estimator, e.quick_move, e.charged_move
             FROM estimators e
             INNER JOIN (
                 SELECT defender, MIN(estimator) AS min_estimator
@@ -69,12 +70,17 @@ class Calculator:
                     AND defender IN ("{'","'.join(defenders)}")
                 GROUP BY defender
             ) AS m ON e.defender = m.defender
-            WHERE party = ? AND level = ? AND attacker IN ("{'","'.join(attackers.keys())}")
-            GROUP BY e.defender, e.attacker
+            INNER JOIN (
+                SELECT defender, attacker, MIN(estimator) as min_estimator
+                FROM estimators
+                WHERE party = ? AND level = ? AND charged_move IN ("{'","'.join(charged_moves)}")
+                GROUP BY defender, attacker
+            ) AS ms ON e.defender = ms.defender AND e.attacker = ms.attacker AND e.estimator = ms.min_estimator
+            WHERE e.attacker IN ("{'","'.join(attackers.keys())}")
         """,
             (party, party, level),
         )
-        for raid_tier, _defender, attacker, estimator in res.fetchall():
+        for raid_tier, _defender, attacker, estimator, fast_move, charged_move in res.fetchall():
             if raid_tier == "RAID_LEVEL_MEGA_5":
                 simplified_raid_tier = "RAID_LEVEL_MEGA"
             elif raid_tier == "RAID_LEVEL_ULTRA_BEAST":
@@ -82,6 +88,7 @@ class Calculator:
             else:
                 simplified_raid_tier = raid_tier
             attackers[attacker][simplified_raid_tier].append(estimator)
+            movesets[attacker] = (fast_move, charged_move)
 
         ase = {}
         for attacker, estimators in attackers.items():
@@ -92,7 +99,8 @@ class Calculator:
                 + 0.50 * sum(estimators["RAID_LEVEL_5"]) / len(estimators["RAID_LEVEL_5"])
                 + 0.35 * sum(estimators["RAID_LEVEL_MEGA"]) / len(estimators["RAID_LEVEL_MEGA"])
             )
-            print(f"{attacker},{ase}")
+            fast_move, charged_move = movesets[attacker]
+            print(f"{attacker},{ase},{fast_move},{charged_move}")
 
     def _raid_bosses(self):
         raid_tiers = []
@@ -166,7 +174,6 @@ class Calculator:
             "includeLegendary": "true",
             "includeShadow": "true",
             "includeMegas": "true",
-            "attackerTypes": self.attacker_types,
             "primalAssistants": "",
             "numParty": str(party),
         }