Center and Border pencil marks, plus added data types

pydoku.py

@@ -1,6 +1,9 @@
 import pygame as pg
 from pygame import Rect, Vector2 as v2
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Set
+
 import os
 from websudoku import *
 
@@ -9,6 +12,24 @@ class Cursor:
     row: int = 5
     col: int = 5
 
+@dataclass
+class Cell:
+    value: int = 0
+    given: bool = False
+
+@dataclass
+class PencilMark:
+    center: Set[int] = field(default_factory=set)
+    border: Set[int] = field(default_factory=set)
+
+class InputMethod(Enum):
+    """
+    This enum helps representing Snyder notation in the cells
+    """
+    FILL = 1,
+    CENTER = 2,
+    BORDER = 3
+
 
 ##################
 # PYGAME INIT
@@ -44,16 +65,26 @@ GRID_Y = screen.get_height() / 2 - 9 * CELL_SIZE // 2
 cursorSurface = pg.Surface((CURSOR_SIZE, CURSOR_SIZE), pg.SRCALPHA)
 
 cursor = Cursor()
+input_method = InputMethod.FILL
 
-sudoku = get_sudoku_puzzle(Difficulty.EVIL)
+# TODO: We have to make this async
+# sudoku = get_sudoku_puzzle(Difficulty.EVIL)
+sudoku = get_sudoku_puzzle_test()
 
 board = []
 
 for i in range(81):
     if sudoku.editmask[i] == 0:
-        board.append((sudoku.cheat[i], True))
+        board.append(Cell(sudoku.cheat[i], True))
     else:
-        board.append((0, False))
+        board.append(Cell())
+
+# pencil_marks = [PencilMark() for _ in range(81)]
+pencil_marks = []
+for i in range(81):
+    pm = PencilMark() if not board[i].given else None
+    pencil_marks.append(pm)
+
 
 def draw_grid():
     for row in range(9):
@@ -85,22 +116,45 @@ def draw_cursor():
                 CURSOR_SIZE, CURSOR_SIZE)
     screen.blit(cursorSurface, rect)
 
-filled_num_font = pg.font.Font(None, 60)
+def draw_pm_border(row, col, idx):
+    third = CELL_SIZE // 3
+    for n in pencil_marks[idx].border:
+        digit = pm_border_font.render(str(n), True, "red")
+        padding = 6
+        basepos = (GRID_X + CELL_SIZE * col + padding, GRID_Y + CELL_SIZE * row + padding)
+        pos = basepos[0] + ((n - 1) % 3) * third , basepos[1] + ((n - 1) // 3) * third
+        screen.blit(digit, pos)
+
+def draw_pm_center(row, col, idx):
+    third = CELL_SIZE // 3
+    nums = "".join(map(str, sorted(pencil_marks[idx].center)))
+    digits = pm_center_font.render(nums, True, "green")
+    pos = (GRID_X + CELL_SIZE * col + CELL_SIZE // 2 - digits.get_width() // 2,
+           GRID_Y + CELL_SIZE * row + CELL_SIZE // 2 - digits.get_height() // 2.3)
+    screen.blit(digits, pos)
+
 def draw_numbers():
     for row in range(9):
         for col in range(9):
-            num = board[row * 9 + col][0]
-            if num == 0:
-                continue
-            digit = filled_num_font.render(str(num), True, "black")
+            idx = row * 9 + col
+            cell = board[idx]
+            if cell.value > 0:
+                color = "black" if cell.given else "dark blue"
+                digit = filled_num_font.render(str(cell.value), True, color)
+
                 pos = (GRID_X + CELL_SIZE * col + CELL_SIZE // 2 - digit.get_width() // 2,
                     # Here we divide the height by 2.3 because there seems to be some extra
                     # padding at the bottom of the surface
                     GRID_Y + CELL_SIZE * row + CELL_SIZE // 2 - digit.get_height() // 2.3)
 
                 screen.blit(digit, pos)
+                continue
+            if pencil_marks[idx].border:
+                draw_pm_border(row, col, idx)
+            if pencil_marks[idx].center:
+                draw_pm_center(row, col, idx)
+                # continue
 
-title = title_font.render("Pydoku", True, "black")
 def draw_hud():
     screen.blit(title, (screen.get_width() // 2 - title.get_width() // 2, 70))
 
@@ -136,16 +190,42 @@ while running:
             if num.isdigit() and num != '0':
                 num = int(num)
                 idx = (cursor.row - 1) * 9 + cursor.col - 1
-                if not board[idx][1]:
-                    board[idx] = (num, False)
+                if not board[idx].given:
+                    match input_method:
+                        case InputMethod.FILL:
+                            board[idx].value = num
+                        case InputMethod.CENTER:
+                            if num in pencil_marks[idx].center:
+                                pencil_marks[idx].center.remove(num)
+                            # Max 3 center pencil marks
+                            elif len(pencil_marks[idx].center) < 3:
+                                pencil_marks[idx].center.add(num)
+                        case InputMethod.BORDER:
+                            if num in pencil_marks[idx].border:
+                                pencil_marks[idx].border.remove(num)
+                            else:
+                                pencil_marks[idx].border.add(num)
 
             if event.key == pg.K_BACKSPACE:
                 idx = (cursor.row - 1) * 9 + cursor.col - 1
-                if not board[idx][1]:
-                    board[idx] = (0, False)
+                if not board[idx].given:
+                    board[idx].value = 0
+
+            match event.key:
+                case pg.K_f:
+                    input_method = InputMethod.FILL
+                case pg.K_c:
+                    input_method = InputMethod.CENTER
+                case pg.K_b:
+                    input_method = InputMethod.BORDER
 
             if event.key == pg.K_F1:
                 print(sudoku)
+            if event.key == pg.K_F2:
+                idx = (cursor.row - 1) * 9 + cursor.col - 1
+                if not board[idx].given:
+                    pencil_marks[idx].border = {i for i in range(1, 10)}
+
 
         
     screen.fill('cornflower blue')