Add arrows kana type and improved grid repr

1 files changed, 71 insertions, 41 deletions

diff --git a/solver.py b/solver.py
index ff53960..b21ea88 100755
--- a/solver.py
+++ b/solver.py
@@ -18,18 +18,22 @@ import zlib
 import docopt
 import yaml
 
-
 class Kana:
 
+    types_without_kana = ('void', 'empt')
+    types_valid_for_chain = ('norm', 'froz', 'rock', 'ar_u', 'ar_l', 'ar_d', 'ar_r')
+    types_with_kana = types_valid_for_chain + ('myst', )
+    types_all = types_with_kana + types_without_kana
+
     def __init__(self, type_name, kana=None):
         self.type_name = type_name
         self.kana = kana
 
         #print(type_name)
         #print(kana)
-        assert type_name in ('void', 'norm', 'empt', 'froz', 'rock', 'myst')
-        if type_name in ('norm', 'rock', 'froz', 'myst'):
-            assert kana[0] in ('k', 's', 'n')
+        assert type_name in self.types_all
+        if type_name in self.types_with_kana:
+            assert kana[0] in ('k', 's', 'n', 'h', 't', 'm')
             assert kana[1] in ('a', 'i', 'u', 'e', 'o')
 
     def __repr__(self):
@@ -38,27 +42,11 @@ class Kana:
     def __eq__(self, other):
         return self.type_name == other.type_name and self.kana == other.kana
 
-kana_void = Kana("void")
-
-
-def is_swappable(kana1, kana2):
-    table_ok = {
-            "norm": ("norm", "empt", "froz"),
-            "froz": ("norm", "empt"),
-    }
-    if kana1.type_name in table_ok:
-        if kana2.type_name in table_ok[kana1.type_name]:
-            return True
-    elif kana2.type_name in table_ok:
-        if kana1.type_name in table_ok[kana2.type_name]:
-            return True
-    return False
+kana_void = Kana('void')
 
 
 class KanaGrid:
 
-    valid_types_for_chain = ("froz", "norm", "rock")
-
     def __init__(self, size, grid, action_count=0, score=0, parent=None):
 
         self.width = size[0]
@@ -77,13 +65,43 @@ class KanaGrid:
                 parent=self.parent,
                 )
 
+    def is_swappable(self, pos1, pos2):
+        kana1 = self.get_kana(pos1)
+        kana2 = self.get_kana(pos2)
+        table_ok = {
+            'norm': ('norm', 'empt', 'froz', 'ar_u', 'ar_r', 'ar_d', 'ar_l'),
+            'empt': ('norm',         'froz', 'ar_u', 'ar_r', 'ar_d', 'ar_l'),
+            'froz': ('norm', 'empt',         'ar_u', 'ar_r', 'ar_d', 'ar_l'),
+            'ar_u': ('norm', 'empt', 'froz',                 'ar_d'        ),
+            'ar_r': ('norm', 'empt', 'froz',                         'ar_l'),
+            'ar_d': ('norm', 'empt', 'froz', 'ar_u'                        ),
+            'ar_l': ('norm', 'empt', 'froz',         'ar_r'                ),
+        }
+        if kana1.type_name in table_ok:
+            if kana2.type_name in table_ok[kana1.type_name]:
+                ar_vect_ok = {
+                        'ar_u': ( 0, -1),
+                        'ar_l': (-1,  0),
+                        'ar_d': ( 0,  1),
+                        'ar_r': ( 1,  0),
+                }
+                vect1_target = ar_vect_ok.get(kana1.type_name, None)
+                vect2_target = ar_vect_ok.get(kana2.type_name, None)
+                if vect1_target or vect2_target:
+                    vect1 = (pos2[0] - pos1[0], pos2[1] - pos1[1])
+                    vect2 = (pos1[0] - pos2[0], pos1[1] - pos2[1])
+                    if vect1 != vect1_target and vect2 != vect2_target:
+                        return False
+                return True
+        return False
+
     def action(self, pos, action_type):
-        kana = self.get_kana(pos)
         if action_type == "reveal":
-            if kana.type_name == "myst":
+            kana = self.get_kana(pos)
+            if kana.type_name == 'myst':
                 new_grid = self.copy()
                 new_grid.action_count += 1
-                new_grid.set_kana(pos, Kana("norm", kana.kana))
+                new_grid.set_kana(pos, Kana('norm', kana.kana))
                 return new_grid
         elif action_type in ("up", "right", "down", "left"):
             if action_type == "up":
@@ -94,8 +112,7 @@ class KanaGrid:
                 pos_dest = (pos[0], pos[1]+1)
             elif action_type == "left":
                 pos_dest = (pos[0]-1, pos[1])
-            kana_dest = self.get_kana(pos_dest)
-            if is_swappable(kana, kana_dest):
+            if self.is_swappable(pos, pos_dest):
                 new_grid = self.copy()
                 new_grid.action_count += 1
                 new_grid.swap_kana(pos, pos_dest)
@@ -105,7 +122,7 @@ class KanaGrid:
         for y in range(self.height):
             for x in range(self.width):
                 kana = self.get_kana((x, y))
-                if kana.type_name in self.valid_types_for_chain:
+                if kana.type_name in Kana.types_valid_for_chain:
                     yield (x, y)
 
     def populate_chain(self, pos1, chain_positions):
@@ -123,7 +140,7 @@ class KanaGrid:
             if pos2 in chain_positions:
                 continue
             kana2 = self.get_kana(pos2)
-            if kana2.type_name in self.valid_types_for_chain:
+            if kana2.type_name in Kana.types_valid_for_chain:
                 if is_kana_compatible(kana1, kana2):
                     self.populate_chain(pos2, chain_positions)
 
@@ -172,25 +189,24 @@ class KanaGrid:
     def swap_kana(self, pos1, pos2):
         kana_dst = self.get_kana(pos2)
 
-        if kana_dst.type_name == "froz":
+        if kana_dst.type_name == 'froz':
             pos_tmp = pos1
             pos1 = pos2
             pos2 = pos_tmp
             kana_dst = self.get_kana(pos2)
 
-        # important
         kana_src = self.get_kana(pos1)
         pos_src = pos1
         pos_dst = pos2
         vect = (pos2[0] - pos1[0], pos2[1] - pos1[1])
 
-        while is_swappable(kana_src, kana_dst):
+        while self.is_swappable(pos_src, pos_dst):
             #print("swap between src %s (%s) dst %s (%s)"
             #      % (kana_src, pos_src, kana_dst, pos_dst))
             self.set_kana(pos_src, kana_dst)
             self.set_kana(pos_dst, kana_src)
 
-            if kana_src.type_name != "froz":
+            if kana_src.type_name != 'froz':
                 break
 
             pos_src = pos_dst
@@ -226,20 +242,34 @@ class KanaGrid:
 
 
 def repr_grid(grid, grid_size):
+    indicator_map = {
+            'ar_u': ('\x1b[30m∧\x1b[0m', '\x1b[30m∧\x1b[0m'),
+            'ar_r': ('\x1b[30m>\x1b[0m', '\x1b[30m>\x1b[0m'),
+            'ar_d': ('\x1b[30m∨\x1b[0m', '\x1b[30m∨\x1b[0m'),
+            'ar_l': ('\x1b[30m<\x1b[0m', '\x1b[30m<\x1b[0m'),
+            'froz': (' \x1b[36m', '\x1b[0m '),
+            'rock': (' \x1b[1;40m', '\x1b[0m '),
+            'myst': ('\x1b[33m?', '?\x1b[0m'),
+    }
     lines = []
     kana_iter = iter(grid)
     for y in range(grid_size[1]):
-        line = ""
+        line = ''
         for x in range(grid_size[0]):
             kana = next(kana_iter)
-            if kana.type_name == "void":
-                line += "      "
-            elif kana.type_name == "empt":
-                line += "|    |"
-            elif kana.type_name == "myst":
-                line += "| ?? |"
-            elif kana.type_name in ("froz", "norm", "rock"):
-                line += "| %s |" % kana.kana
+            tname = kana.type_name
+            if tname == 'void':
+                line += '      '
+            elif tname == 'empt':
+                line += '|    |'
+            elif tname in indicator_map:
+                line += '|%s%s%s|' % (
+                        indicator_map[tname][0],
+                        kana.kana,
+                        indicator_map[tname][1],
+                )
+            elif kana.type_name in Kana.types_with_kana:
+                line += '| %s |' % kana.kana
         lines.append(line)
     return '\n'.join(lines)