From 11807628fab4cacd849ae1fc572ad2dc95dfd8d9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20Ke=C3=9Fler?= <git@maximilian-kessler.de>
Date: Fri, 11 Oct 2024 21:38:29 +0200
Subject: [PATCH] expand cli command solve

---
 src/hanabi/cli.py                  |  6 ++++--
 src/hanabi/live/instance_finder.py | 20 ++++++++++++--------
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/hanabi/cli.py b/src/hanabi/cli.py
index e4f5805..6b36698 100755
--- a/src/hanabi/cli.py
+++ b/src/hanabi/cli.py
@@ -98,8 +98,8 @@ def subcommand_download(
         logger.info("Successfully exported games for all variants")
 
 
-def subcommand_solve(var_id, timeout: int):
-    instance_finder.solve_unknown_seeds(var_id, timeout)
+def subcommand_solve(var_id: int, seed_class: int, num_players: Optional[int], timeout: int, ):
+    instance_finder.solve_unknown_seeds(var_id, seed_class, num_players, timeout)
 
 
 def subcommand_gen_config():
@@ -165,6 +165,8 @@ def add_solve_subparser(subparsers):
     parser = subparsers.add_parser('solve', help='Seed solving')
     parser.add_argument('var_id', type=int, help='Variant id to solve instances from.', default=0)
     parser.add_argument('--timeout', '-t', type=int, help='Timeout [s] for individual seeds.', default=150)
+    parser.add_argument('--class', '-c', type=int, dest='seed_class', help='Class of seed to analyze. 0 stands for hanab.live seeds', default=0)
+    parser.add_argument('--num_players', '-p', type=int, help='Restrict to number of players. If not specified, all player counts are analyzed.', default = None)
 
 def add_decompress_subparser(subparsers):
     parser = subparsers.add_parser('decompress', help='Decompress a hanab.live JSON-encoded replay link')
diff --git a/src/hanabi/live/instance_finder.py b/src/hanabi/live/instance_finder.py
index f42d862..a64ecc2 100644
--- a/src/hanabi/live/instance_finder.py
+++ b/src/hanabi/live/instance_finder.py
@@ -188,15 +188,19 @@ def solve_seed(seed, num_players, deck, var_name: str, timeout: Optional[int] =
         traceback.print_exc()
 
 
-def solve_unknown_seeds(variant_id, timeout: Optional[int] = 150):
+def solve_unknown_seeds(variant_id, seed_class: int = 0, num_players: Optional[int] = None, timeout: Optional[int] = 150):
     variant_name = variants.variant_name(variant_id)
-    database.cur.execute(
-        "SELECT seeds.seed, num_players, array_agg(suit_index order by deck_index asc), array_agg(rank order by deck_index asc) "
-        "FROM seeds "
-        "INNER JOIN decks ON seeds.seed = decks.seed "
-        "WHERE variant_id = (%s) AND num_players = 2 AND class = 1 AND feasible is null "
-        "GROUP BY seeds.seed order by num",
-        (variant_id,)
+    query = "SELECT seeds.seed, num_players, array_agg(suit_index order by deck_index asc), array_agg(rank order by deck_index asc) "\
+            "FROM seeds "\
+            "INNER JOIN decks ON seeds.seed = decks.seed "\
+            "WHERE variant_id = (%s) "\
+            "AND class = (%s) "\
+            "AND feasible IS NULL "
+    if num_players is not None:
+        query += "AND num_players = {} ".format(num_players)
+    query += "GROUP BY seeds.seed ORDER BY num"
+    database.cur.execute(query,
+        (variant_id, seed_class)
     )
     res = database.cur.fetchall()
     data = []