From 66a63d3f6b150ef5d59b9bfbfc1b5c4154bec497 Mon Sep 17 00:00:00 2001
From: LU Jialin <luxxxlucy@gmail.com>
Date: Thu, 7 Nov 2024 22:20:11 -0800
Subject: [PATCH] print the grid in terminal, modified based on PR #2

---
 arckit/vis.py | 44 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 43 insertions(+), 1 deletion(-)

diff --git a/arckit/vis.py b/arckit/vis.py
index 1c3d622..798be69 100644
--- a/arckit/vis.py
+++ b/arckit/vis.py
@@ -1,6 +1,7 @@
 import drawsvg
 import numpy as np
 import io
+import rich
 
 cmap = [
     '#252525', # black
@@ -262,4 +263,45 @@ def output_drawing(d: drawsvg.Drawing, filename: str, context=None):
         import cairosvg
         cairosvg.svg2pdf(bytestring=buffer.getvalue(), write_to=filename)
     else:
-        raise ValueError(f'Unknown file extension for {filename}')
\ No newline at end of file
+        raise ValueError(f'Unknown file extension for {filename}')
+
+def print_grid(grid: np.ndarray):
+    """
+    Print a grid to the terminal using rich library
+
+    Parameters
+    ----------
+    grid : np.ndarray
+        the standard grid format for Task 
+    """
+    CELL_WIDTH = 2
+
+    def get_color(color_str):
+        color_str = color_str.strip('#')
+        return rich.color.Color.from_rgb(*bytes.fromhex(color_str))
+
+    # Translate 'cmap' to rich style with respective color as background
+    rich_scmap = {
+        i: rich.style.Style(bgcolor=get_color(color))
+        for i, color in enumerate(cmap)
+    }
+
+    # Create and populate rich table
+    table = rich.table.Table.grid(expand=False)
+    
+    height, width = grid.shape
+    
+    # Add columns
+    for x in range(width):
+        table.add_column()
+        table.columns[x]._cells = [''] * height
+    
+    table.rows = [rich.table.Row()] * height
+
+    # Populate cells
+    for y in range(height):
+        for x in range(width):
+            color_idx = grid[y, x]
+            table.columns[x]._cells[y] = rich.text.Text(' ' * CELL_WIDTH, style=rich_scmap[color_idx])
+
+    rich.print(table)
\ No newline at end of file