1 files changed, 82 insertions, 6 deletions

diff --git a/cameltris.py b/cameltris.py
index cac81f2..d2ed6f8 100755
--- a/cameltris.py
+++ b/cameltris.py
@@ -6,8 +6,62 @@ import time
 import pygame
 
 
-def zpiece():
-    return ((blue_square, blue_square, None), (None, blue_square, blue_square))
+class Piece():
+    def rotate(self):
+        pass
+
+    def rotate_clockwise(self):
+        self.rotate()
+
+    def rotate_counter_clockwise(self):
+        self.rotate()
+
+
+class ZPiece(Piece):
+    def __init__(self):
+        self.elements = (blue_square, blue_square, None), (None, blue_square, blue_square), (None, None, None)
+        self.vertical = False
+
+    def rotate(self):
+        if self.vertical:
+            self.elements = (blue_square, blue_square, None), (None, blue_square, blue_square), (None, None, None)
+        else:
+            self.elements = (None, None, blue_square), (None, blue_square, blue_square), (None, blue_square, None)
+
+        self.vertical = not self.vertical
+
+
+class SPiece(Piece):
+    def __init__(self):
+        self.elements = (None, None, None), (None, green_square, green_square), (green_square, green_square, None)
+        self.vertical = False
+
+    def rotate(self):
+        if self.vertical:
+            self.elements = (None, green_square, green_square), (green_square, green_square, None), (None, None, None)
+        else:
+            self.elements = (None, green_square, None), (None, green_square, green_square), (None, None, green_square)
+
+        self.vertical = not self.vertical
+
+
+class SquarePiece(Piece):
+    def __init__(self):
+        self.elements = ((white_square, white_square), (white_square, white_square))
+
+
+class IPiece(Piece):
+    def __init__(self):
+        self.elements = (None, None, None, None), (None, None, None, None), (red_square, red_square, red_square, red_square), (None, None, None, None)
+        self.vertical = False
+
+    def rotate(self):
+        if self.vertical:
+            self.elements = (None, None, None, None), (None, None, None, None), (red_square, red_square, red_square, red_square), (None, None, None, None)
+        else:
+            self.elements = (None, None, red_square, None), (None, None, red_square, None), (None, None, red_square, None), (None, None, red_square, None)
+
+        self.vertical = not self.vertical
 
 
 def refresh_screen():
@@ -19,7 +73,7 @@ def refresh_screen():
                 screen.blit(element, (col_idx * 50 + 1, row_idx * 50 + 1))
 
     # Display the current piece
-    for row_idx, row in enumerate(current_piece):
+    for row_idx, row in enumerate(current_piece.elements):
         for col_idx, element in enumerate(row):
             if element is not None:
                 screen.blit(element, ((col_idx + current_piece_position[1]) * 50 + 1, (row_idx + current_piece_position[0]) * 50 + 1))
@@ -29,6 +83,8 @@ pygame.init()
 black = (0, 0, 0)
 white = (0xff, 0xff, 0xff)
 blue = (0x10, 0x20, 0xbb)
+green = (0x20, 0xbb, 0x10)
+red = (0xbb, 0x10, 0x20)
 
 white_square = pygame.Surface((48, 48))
 white_square.fill(white)
@@ -36,19 +92,35 @@ white_square.fill(white)
 blue_square = pygame.Surface((48, 48))
 blue_square.fill(blue)
 
+green_square = pygame.Surface((48, 48))
+green_square.fill(green)
+
+red_square = pygame.Surface((48, 48))
+red_square.fill(red)
+
 screen = pygame.display.set_mode((500, 1000))
 
 grid = [[None for _ in range(10)] for _ in range(20)]
 
-# Let's start with a Z piece
-current_piece = zpiece()
-current_piece_position = [0, 4]
+# Let's start with a I piece
+current_piece = ZPiece()
+
+for row_id, row in enumerate(current_piece.elements):
+    if list(filter(lambda x: x is not None, row)):
+        break
+
+initial_y_position = -row_id
+initial_x_position = (len(grid[0]) // 2) - (len(current_piece.elements[0]) / 2)
+
+current_piece_position = [initial_y_position, initial_x_position]
 
 while True:
     for event in pygame.event.get():
         if event.type == pygame.QUIT:
             sys.exit()
         elif event.type == pygame.KEYDOWN:
+            if event.key in (pygame.K_ESCAPE, pygame.K_q):
+                sys.exit()
             if event.key in (pygame.K_DOWN, pygame.K_j):
                 current_piece_position[0] += 1
             if event.key in (pygame.K_UP, pygame.K_k):
@@ -57,6 +129,10 @@ while True:
                 current_piece_position[1] += 1
             if event.key in (pygame.K_LEFT, pygame.K_h):
                 current_piece_position[1] -= 1
+            if event.key == pygame.K_s:
+                current_piece.rotate_clockwise()
+            if event.key == pygame.K_d:
+                current_piece.rotate_counter_clockwise()
 
     refresh_screen()
     pygame.display.flip()