From bd334c9abea732dc911032c4d2e39db3e62e2c17 Mon Sep 17 00:00:00 2001
From: Thomas Lindner <tom@dl6tom.de>
Date: Fri, 30 Dec 2022 17:32:55 +0100
Subject: [PATCH] implement castling

---
 src/blunderboard/movegenerator.py | 100 +++++++++++++++++++++++++++---
 1 file changed, 93 insertions(+), 7 deletions(-)

diff --git a/src/blunderboard/movegenerator.py b/src/blunderboard/movegenerator.py
index a37d6b4..925473a 100644
--- a/src/blunderboard/movegenerator.py
+++ b/src/blunderboard/movegenerator.py
@@ -72,18 +72,104 @@ class TakeState(State):
 
 class TakeTakeState(State):
     def __init__(
-        self, blunder_evaluator: BlunderEvaluator, from_field: str, to_field: str
+        self, blunder_evaluator: BlunderEvaluator, from_field: str, from2_field: str
     ):
         super().__init__(blunder_evaluator)
         self.from_field = from_field
-        self.to_field = to_field
+        self.from2_field = from2_field
 
     def put(self, row: int, column: int) -> State:
-        field = coords_to_field(row, column)
-        if self.to_field == field:
-            move = self.from_field + field
-        elif self.from_field == field:
-            move = self.to_field + field
+        to_field = coords_to_field(row, column)
+        if self.from2_field == to_field:
+            move = self.from_field + to_field
+        elif self.from_field == to_field:
+            move = self.from2_field + to_field
+        elif (
+            self.from_field[1] == self.from2_field[1]
+            and self.from_field[1] == to_field[1]
+        ):
+            # king-side castling
+            if (
+                self.from_field in ["e1", "e8"]
+                and self.from2_field in ["h1", "h8"]
+                and to_field in ["f1", "f8", "g1", "g8"]
+            ):
+                return KingCastleState(
+                    self.blunder_evaluator, self.from_field, self.from2_field, to_field
+                )
+            # queen-side castling
+            if (
+                self.from_field in ["e1", "e8"]
+                and self.from2_field in ["a1", "a8"]
+                and to_field in ["c1", "c8", "d1", "d8"]
+            ):
+                return QueenCastleState(
+                    self.blunder_evaluator, self.from_field, self.from2_field, to_field
+                )
+            print("ignored invalid put")
+            return self
+        else:
+            print("ignored invalid put")
+            return self
+        print("move %s" % move)
+        self.blunder_evaluator.move(move)
+        return InitState(self.blunder_evaluator)
+
+
+class KingCastleState(State):
+    def __init__(
+        self,
+        blunder_evaluator: BlunderEvaluator,
+        from_field: str,
+        from2_field: str,
+        to_field: str,
+    ):
+        super().__init__(blunder_evaluator)
+        self.from_field = from_field
+        self.from2_field = from2_field
+        self.to_field = to_field
+
+    def put(self, row: int, column: int) -> State:
+        to2_field = coords_to_field(row, column)
+        if self.to_field[1] == to2_field[1]:
+            if to2_field in ["f1", "g1"]:
+                move = "e1g1"
+            elif to2_field in ["f8", "g8"]:
+                move = "e8g8"
+            else:
+                print("ignored invalid put")
+                return self
+        else:
+            print("ignored invalid put")
+            return self
+        print("move %s" % move)
+        self.blunder_evaluator.move(move)
+        return InitState(self.blunder_evaluator)
+
+
+class QueenCastleState(State):
+    def __init__(
+        self,
+        blunder_evaluator: BlunderEvaluator,
+        from_field: str,
+        from2_field: str,
+        to_field: str,
+    ):
+        super().__init__(blunder_evaluator)
+        self.from_field = from_field
+        self.from2_field = from2_field
+        self.to_field = to_field
+
+    def put(self, row: int, column: int) -> State:
+        to2_field = coords_to_field(row, column)
+        if self.to_field[1] == to2_field[1]:
+            if to2_field in ["c1", "d1"]:
+                move = "e1c1"
+            elif to2_field in ["c8", "d8"]:
+                move = "e8c8"
+            else:
+                print("ignored invalid put")
+                return self
         else:
             print("ignored invalid put")
             return self