From 07bbd9f9d3b0f52506eeb6525c5502643510eb4f Mon Sep 17 00:00:00 2001
From: Nao Pross <np@0hm.ch>
Date: Sun, 26 May 2024 16:58:49 +0200
Subject: Create pretty printing functions for some polymatrix types

---
 polymatrix/polymatrix/impl.py  | 25 +++++++++++++++++++
 polymatrix/polymatrix/index.py | 10 +++++++-
 polymatrix/polymatrix/utils.py | 55 ++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 87 insertions(+), 3 deletions(-)

diff --git a/polymatrix/polymatrix/impl.py b/polymatrix/polymatrix/impl.py
index 07b9ffa..bb42313 100644
--- a/polymatrix/polymatrix/impl.py
+++ b/polymatrix/polymatrix/impl.py
@@ -10,6 +10,31 @@ class PolyMatrixImpl(PolyMatrix):
     data: PolyMatrixDict
     shape: tuple[int, int]
 
+    def __str__(self) -> str:
+        nrows, ncols = self.shape
+        columns = []
+        for c in range(ncols):
+            rows, row_width = [], 0
+            for r in range(nrows):
+                if (r, c) in self.data:
+                    entry = str(self.data[r, c])
+                    rows.append(entry)
+
+                    if len(entry) > row_width:
+                        row_width = len(entry)
+                else:
+                    rows.append("0")
+            columns.append(tuple(row.rjust(row_width) for row in rows))
+
+        string = ""
+        for r in range(nrows):
+            column = ",  ".join(columns[c][r] for c in range(ncols))
+            string += f"[ {column} ]\n"
+
+        return string
+
+
+
 
 @dataclassabc.dataclassabc(frozen=True)
 class BroadcastPolyMatrixImpl(BroadcastPolyMatrixMixin):
diff --git a/polymatrix/polymatrix/index.py b/polymatrix/polymatrix/index.py
index 7fe3620..83eccf5 100644
--- a/polymatrix/polymatrix/index.py
+++ b/polymatrix/polymatrix/index.py
@@ -28,6 +28,11 @@ class VariableIndex(NamedTuple):
         """ Variables indices can be sorted with respect to variable index. """
         return self.index < other.index
 
+    def __str__(self) -> str:
+        if self.power == 1:
+            return f"V({self.index:d})"
+        return f"V({self.index:d})^{self.power:d}"
+
 
 class MonomialIndex(tuple[VariableIndex]):
     r"""
@@ -48,7 +53,7 @@ class MonomialIndex(tuple[VariableIndex]):
         return "MonomialIndex(" + ",".join(repr(v) for v in self) + ")"
 
     def __str__(self):
-        return "MonomialIndex(" + ",".join(f"{v.index}^{v.power}" for v in self) + ")"
+        return " * ".join(str(v) for v in sorted(self))
 
     @property
     def degree(self) -> int:
@@ -158,6 +163,9 @@ class PolyDict(UserDict[MonomialIndex, int | float]):
     def __repr__(self):
         return f"{self.__class__.__qualname__}({super().__repr__()})"
 
+    def __str__(self) -> str:
+        return " + ".join(f"{str(c)} * {str(m)}" for m, c in self.items())
+
     def __setitem__(self, key: Iterable[VariableIndex] | MonomialIndex, value: int | float):
         if not isinstance(key, MonomialIndex):
             key = MonomialIndex(key)
diff --git a/polymatrix/polymatrix/utils.py b/polymatrix/polymatrix/utils.py
index 289302a..d1f9d51 100644
--- a/polymatrix/polymatrix/utils.py
+++ b/polymatrix/polymatrix/utils.py
@@ -1,10 +1,62 @@
 import itertools
 import math
 
-from polymatrix.polymatrix.index import MonomialData, MonomialIndex, PolyDict
+from polymatrix.expressionstate import ExpressionState
+from polymatrix.polymatrix.index import MonomialData, MonomialIndex, PolyDict, PolyMatrixDict
+from polymatrix.polymatrix.mixins import PolyMatrixMixin, PolyMatrixAsDictMixin
 from polymatrix.utils.deprecation import deprecated
 
 
+def pretty_format_polydict(p: PolyDict, state: ExpressionState) -> str:
+    """ Pretty print a polynomial """
+    strings = []
+    for monomial, coeff in p.terms():
+        s = f"{coeff}"
+        for var in monomial:
+            if var.power == 1:
+                s += f" * {state.get_name(var.index)}"
+            else:
+                s += f" * {state.get_name(var.index)}^{var.power:d}"
+
+        strings.append(s)
+    return " + ".join(strings)
+
+
+def pretty_format_polymatrixdict(p: PolyMatrixDict, shape: tuple[int, int], state: ExpressionState) -> str:
+    """ Pretty print a polymatrix """
+    nrows, ncols = shape
+    columns = []
+    for c in range(ncols):
+        rows, row_width = [], 0
+        for r in range(nrows):
+            if (r, c) in p:
+                entry = pretty_format_polydict(p[r, c], state)
+                rows.append(entry)
+
+                if len(entry) > row_width:
+                    row_width = len(entry)
+            else:
+                rows.append("0")
+
+        columns.append(tuple(row.rjust(row_width) for row in rows))
+
+    string = ""
+    for r in range(nrows):
+        column = ",  ".join(columns[c][r] for c in range(ncols))
+        string += f"[ {column} ]\n"
+
+    return string
+
+
+def pretty_format_polymatrix(p: PolyMatrixMixin, state: ExpressionState) -> str:
+    """ Pretty print a polymatrix """
+    if isinstance(p, PolyMatrixAsDictMixin):
+        return pretty_format_polymatrixdict(p.data, (p.shape[0], p.shape[1]), state)
+
+    else:
+        # TODO: add pretty print for other functions
+        raise NotImplementedError
+
 
 # NP: this is a werid case of the product
 # NP: merge is not intutive IMHO, find better name
@@ -74,7 +126,6 @@ def multiply_polynomial(
 
 
 # NP: sort mononial indices with respect to variable index in state object
-
 @deprecated("With new index types you can use sorted() on MonomialIndex")
 def sort_monomial_indices(
     monomial: MonomialData,
-- 
cgit v1.2.1