Simplify tehai_mjai() and add test cases

python/mjai/bot/base.py

@@ -206,30 +206,21 @@ def tehai_mjai(self) -> list[str]:
             ["1m", "2m", "6m", "9m", "1p", "3p", "4p", "3s", "4s", "5s",
              "7s", "9s", "5z", "6z"]
         """
-        zi_map = ["E", "S", "W", "N", "P", "F", "C"]
-        ms, ps, ss, zis, akas = [], [], [], [], []
         tiles = []
         for tile_idx, tile_count in enumerate(self.player_state.tehai):
-            if tile_count and tile_idx == 4 and self.akas_in_hand[0]:
-                akas.append("5mr")
-                ms += [f"{tile_idx + 1}m"] * tile_count
-            elif tile_count and tile_idx == 4 + 9 and self.akas_in_hand[1]:
-                akas.append("5pr")
-                ps += [f"{tile_idx + 1}m"] * tile_count
-            elif tile_count and tile_idx == 4 + 18 and self.akas_in_hand[2]:
-                akas.append("5sr")
-                ss += [f"{tile_idx + 1}m"] * tile_count
-            elif tile_count and tile_idx < 9:
-                ms += [f"{tile_idx + 1}m"] * tile_count
-            elif tile_count and tile_idx < 18:
-                ps += [f"{tile_idx - 9 + 1}p"] * tile_count
-            elif tile_count and tile_idx < 27:
-                ss += [f"{tile_idx - 18 + 1}s"] * tile_count
-            else:
-                for _ in range(tile_count):
-                    zis.append(zi_map[tile_idx - 27])
-
-        tiles = ms + ps + ss + zis + akas
+            if tile_idx == 4 and self.akas_in_hand[0]:
+                tile_count -= 1
+                tiles.append("5mr")
+            elif tile_idx == 4 + 9 and self.akas_in_hand[1]:
+                tile_count -= 1
+                tiles.append("5pr")
+            elif tile_idx == 4 + 18 and self.akas_in_hand[2]:
+                tile_count -= 1
+                tiles.append("5sr")
+
+            for _ in range(tile_count):
+                tiles.append(MJAI_VEC34_TILES[tile_idx])
+
         return tiles
 
     @property

tests/mjai/bot/test_base.py

@@ -9,45 +9,33 @@
 def test_tehai_mjai():
     bot = Bot(player_id=0)
     bot.player_state = MagicMock()
-    bot.player_state.tehai = [
-        1,
-        1,
-        1,
-        1,
-        2,
-        1,
-        1,
-        1,
-        1,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        2,
-        2,
-        0,
-        0,
-        0,
-        0,
-        0,
-    ]
+
+    # Case1
+    bot.player_state.tehai = list(
+        map(int, list("1111111110000000000000000000000000"))
+    )
+    bot.player_state.akas_in_hand = [True, False, False]
+    assert bot.tehai == "123406789m"
+    assert bot.tehai_mjai.count("5m") == 0
+    assert bot.tehai_mjai.count("5mr") == 1
+
+    # Case2
+    bot.player_state.tehai = list(
+        map(int, list("1111211110000000000000000000000000"))
+    )
     bot.player_state.akas_in_hand = [True, False, False]
-    assert "5m" in bot.tehai_mjai
-    assert "5mr" in bot.tehai_mjai
+    assert bot.tehai == "1234056789m"
+    assert bot.tehai_mjai.count("5m") == 1
+    assert bot.tehai_mjai.count("5mr") == 1
+
+    # Case3
+    bot.player_state.tehai = list(
+        map(int, list("1111311110000000000000000000000000"))
+    )
+    bot.player_state.akas_in_hand = [False, False, False]
+    assert bot.tehai == "12345556789m"
+    assert bot.tehai_mjai.count("5m") == 3
+    assert bot.tehai_mjai.count("5mr") == 0
 
 
 def test_find_improving_tiles():