path: root/mdpoly/abc.py

""" Abstract Base Classes of MDPoly """
from __future__ import annotations
from typing import TYPE_CHECKING

from .index import Number, Shape, MatrixIndex, PolyIndex
from .constants import NUMERICS_EPS
from .errors import AlgebraicError
from .util import iszero

if TYPE_CHECKING:
    from .state import State

from typing import Self, TypeVar, Generic, Any, Iterable, Sequence
from enum import Enum, auto
from copy import copy
from abc import ABC, abstractmethod
from dataclassabc import dataclassabc
from hashlib import sha256


# ┏━╸╻ ╻┏━┓┏━┓┏━╸┏━┓┏━┓╻┏━┓┏┓╻┏━┓
# ┣╸ ┏╋┛┣━┛┣┳┛┣╸ ┗━┓┗━┓┃┃ ┃┃┗┫┗━┓
# ┗━╸╹ ╹╹  ╹┗╸┗━╸┗━┛┗━┛╹┗━┛╹ ╹┗━┛


class Algebra(Enum):
    """ Types of algebras. """
    none = auto()
    poly_ring = auto()
    matrix_ring = auto()


class Expr(ABC):
    """ Merkle binary tree to represent a mathematical expression. """

    # --- Properties ---

    @property
    def is_leaf(self) -> bool:
        """ Whether an expression is a Leaf of the tree. """
        return False

    @property
    @abstractmethod
    def left(self) -> Self:
        """ Left child of node binary tree. """

    @left.setter
    @abstractmethod
    def left(self, left: Self) -> None:
        """ Setter for left child of binary tree. """
        # TODO: If shape is replaced with functools.cached_property
        # we will need to update shape here.

    @property
    @abstractmethod
    def right(self) -> Self:
        """ Right child of node in binary tree. """

    @right.setter
    @abstractmethod
    def right(self, right: Self) -> None:
        """ Setter for right child of binary tree. """

    @property
    @abstractmethod
    def algebra(self) -> Algebra:
        """ Specifies to which algebra belongs the expression. 
        
        This is used to provide nicer error messages and to avoid accidental
        mistakes (like adding a scalar and a matrix) which are hard to debug
        as expressions are lazily evalated (i.e. without this the exception for
        adding a scalar to a matrix does not occurr at the line where the two
        are added but rather where :py:meth:`mdpoly.abc.Expr.to_repr` is
        called).
        """

    @property
    @abstractmethod
    def shape(self) -> Shape:
        # TODO: Test on very large expressions. If there is a performance hit
        # change to functools.cached_property.
        """ Computes the shape of the expression. This method is also used to
        check if the operation described by *Expr* can be perfomed with *left*
        and *right*. If the shape cannot be computed because of an algebraic
        problem raises :py:class:`mdpoly.errors.AlgebraicError`. """

    # --- Magic Methods ---

    def __repr__(self):
        name = self.__class__.__qualname__
        return f"{name}(left={self.left}, right={self.right})"

    def __iter__(self) -> Iterable[Expr]:
        yield from self.children()

    def __eq__(self, other: object) -> bool:
        if isinstance(other, Expr):
            return self.subtree_hash() == other.subtree_hash()
        return False

    # --- Methods for polynomial expression ---

    @abstractmethod
    def to_repr(self, repr_type: type, state: State) -> tuple[Repr, State]:
        """ Convert to a concrete representation 

        To convert an abstract object from the expression tree, we ... TODO: finish.

        See also :py:class:`mdpoly.state.State`.
        """

    # --- Methods for tree data structure ---

    def subtree_hash(self) -> bytes:
        h = sha256()
        h.update(self.left.subtree_hash())
        h.update(self.right.subtree_hash())
        return h.digest()

    def children(self) -> Sequence[Expr]:
        """ Iterate over the two nodes """
        if self.is_leaf:
            return tuple()
        return self.left, self.right

    def leaves(self) -> Iterable[Expr]:
        """ Returns the leaves of the tree. This is done recursively and in
        :math:`\mathcal{O}(n)`."""
        if self.left.is_leaf:
            yield self.left
        else:
            yield from self.left.leaves()

        if self.right.is_leaf:
            yield self.right
        else:
            yield from self.right.leaves()

    def replace(self, old: Self, new: Self) -> Self:
        """ Create a new expression wherein all expression equal to ``old`` are
        replaced with ``new``.

        This can be used to convert variables into parameters or constants, and
        vice versa. For example:

        .. code:: py

            x, y = Variable.from_names("x, y")
            poly_2d = x ** 2 + 2 * x * y + y ** 2

            # Suppose we want to consider only x as variable to reduce poly_2d
            # to a one-dimensional polynomial
            y_p = Parameter("y_p")
            poly_1d = poly_2d.replace(y, y_p)

        It can also be used for more advanced replacements since ``old`` and
        ``new`` are free to be any expression.
        """
        def replace_all(node):
            if node == old:
                return new 

            if isinstance(node, Leaf):
                return node

            node_copy = copy(node)
            node_copy.left = replace_all(node.left)
            node_copy.right = replace_all(node.right)
            return node_copy

        return replace_all(self)

    def rotate_left(self) -> Expr:
        """ Perform a left tree rotation. """
        root = copy(self) # current root
        pivot = copy(self.right) # new parent

        root.right = pivot.left
        pivot.left = root
        
        return pivot

    def rotate_right(self) -> Expr:
        """ Perform a right tree rotation. """
        root = copy(self) # current root
        pivot = copy(root.left) # new parent

        root.left = pivot.right
        pivot.right = root
        
        return pivot


    # --- Private methods ---

    @staticmethod
    def _assert_same_algebra(left: Expr, right: Expr) -> None:
        if not isinstance(left, Expr) or not isinstance(right, Expr):
            return

        if left.algebra != right.algebra:
            raise AlgebraicError("Cannot perform algebraic operation between "
                                 f"{left} and {right} because they have different "
                                 f"algebras {left.algebra} and {right.algebra}.")

    @staticmethod
    def _wrap(if_type: type, wrapper_type: type, obj: Any, *args, **kwargs) -> Expr:
        """ Wrap non-expr objects. 

        Suppose ``x`` is of type *Expr*, then we would like to be able to do
        things like ``x + 1``, so this function can be called in operator
        overloadings to wrap the 1 into a :py:class:`mdpoly.leaves.Const`. If
        ``obj`` is already of type *Expr*, this function does nothing. The
        arguments *args* and  *kwargs* are forwarded to the constructor of
        *wrapper_type*.
        """
        # Do not wrap if is alreay an expression
        if isinstance(obj, Expr):
            return obj

        if not isinstance(obj, if_type):
            raise TypeError(f"Cannot wrap {obj} with type {wrapper_type} because "
                            f"it is not of type {if_type}.")

        return wrapper_type(obj, *args, **kwargs)

    # --- Operator overloading ---

    def __add__(self, other: Any) -> Self:
        raise NotImplementedError

    def __radd__(self, other: Any) -> Self:
        raise NotImplementedError

    def __sub__(self, other: Any) -> Self:
        raise NotImplementedError

    def __rsub__(self, other: Any) -> Self:
        raise NotImplementedError

    def __neg__(self) -> Self:
        raise NotImplementedError

    def __mul__(self, other: Any) -> Self:
        raise NotImplementedError

    def __rmul__(self, other: Any) -> Self:
        raise NotImplementedError

    def __pow__(self, other: Any) -> Self:
        raise NotImplementedError

    def __matmul__(self, other: Any) -> Self:
        raise NotImplementedError

    def __rmatmul__(self, other: Any) -> Self:
        raise NotImplementedError

    def __truediv__(self, other: Any) -> Self:
        raise NotImplementedError

    def __rtruediv__(self, other: Any) -> Self:
        raise NotImplementedError


class Leaf(Expr):
    """ Leaf of the binary tree. """

    # --- Properties --- 

    @property
    @abstractmethod
    def name(self) -> str:
        """ Name of the leaf. """

    # --- Magic methods ---

    def __repr__(self) -> str:
        return self.name

    # -- Overloads --- 

    @property
    def is_leaf(self):
        """ Overloads :py:meth:`mdpoly.abc.Expr.is_leaf`. """
        return True

    @property
    def left(self):
        """ Overloads left child with function that does nothing. (Leaves do
        not have children). """

    @left.setter
    def left(self, left: Self) -> None:
        """ Overloads left child setter with function that does nothing.
        (Leaves do not have children). """

    @property
    def right(self):
        """ Overloads right child with function that does nothing. (Leaves do
        not have children). """

    @right.setter
    def right(self, right: Self) -> None:
        """ Overloads right child setter with function that does nothing.
        (Leaves do not have children). """

    def subtree_hash(self) -> bytes:
        h = sha256()
        h.update(self.__class__.__qualname__.encode("utf-8"))
        h.update(bytes(self.algebra.value))
        h.update(self.name.encode("utf-8")) # type: ignore
        h.update(bytes(self.shape))
        return h.digest()


T = TypeVar("T")


class Const(Leaf, Generic[T]):
    """ Leaf to represent a constant. TODO: desc. """

    @property
    @abstractmethod
    def value(self) -> T:
        """ Value of the constant. """

    def __repr__(self) -> str:
        if not self.name:
            return repr(self.value)

        return self.name


class Var(Leaf):
    """ Leaf to reprsent a Variable. TODO: desc """


class Param(Leaf):
    """ Parameter. TODO: desc """


@dataclassabc(frozen=True)
class Nothing(Leaf):
    """
    A leaf that is nothing. This is a placeholder (eg. for unary operators).
    """
    name: str = "."
    shape: Shape = Shape(0, 0)
    algebra: Algebra = Algebra.none

    def to_repr(self, repr_type: type, state: State) -> tuple[Repr, State]:
        raise ValueError("Nothing cannot be represented.")

    def __repr__(self) -> str:
        return "Nothing"


# ┏━┓┏━╸╻  ┏━┓╺┳╸╻┏━┓┏┓╻┏━┓
# ┣┳┛┣╸ ┃  ┣━┫ ┃ ┃┃ ┃┃┗┫┗━┓
# ╹┗╸┗━╸┗━╸╹ ╹ ╹ ╹┗━┛╹ ╹┗━┛


class Rel(ABC):
    """ Relation between two expressions. """
    lhs: Expr
    rhs: Expr


# ┏━┓┏━╸┏━┓┏━┓┏━╸┏━┓┏━╸┏┓╻╺┳╸┏━┓╺┳╸╻┏━┓┏┓╻┏━┓
# ┣┳┛┣╸ ┣━┛┣┳┛┣╸ ┗━┓┣╸ ┃┗┫ ┃ ┣━┫ ┃ ┃┃ ┃┃┗┫┗━┓
# ╹┗╸┗━╸╹  ╹┗╸┗━╸┗━┛┗━╸╹ ╹ ╹ ╹ ╹ ╹ ╹┗━┛╹ ╹┗━┛


class Repr(ABC):
    r""" Representation of a multivariate matrix polynomial expression.

    Concretely, a representation must be able to store a mathematical object
    like :math:`Q \in \mathbb{R}^2[x, y]`, for example

    .. math::
        Q = \begin{bmatrix}
            x^2 + 1 &  5y \\
                  0 &   y^2
        \end{bmatrix}

    See also :py:mod:`mdpoly.representations`.

    """
    @abstractmethod
    def at(self, entry: MatrixIndex, term: PolyIndex) -> Number:
        """ Access polynomial coefficient. """

    @abstractmethod
    def set(self, entry: MatrixIndex, term: PolyIndex, value: Number) -> None:
        """ Set value of polynomial coefficient. """

    def is_zero(self, entry: MatrixIndex, term: PolyIndex, eps: Number = NUMERICS_EPS) -> bool:
        """ Check if a polynomial coefficient is zero. """
        return iszero(self.at(entry, term), eps=eps)

    @abstractmethod
    def set_zero(self, entry: MatrixIndex, term: PolyIndex) -> None:
        """ Set a coefficient to zero (delete it). """
        # Default implementation for refence, you should do something better!
        self.set(entry, term, 0.)

    @abstractmethod
    def entries(self) -> Iterable[MatrixIndex]:
        """ Return indices to non-zero entries of the matrix. """

    @abstractmethod
    def terms(self, entry: MatrixIndex) -> Iterable[PolyIndex]:
        """ Return indices to non-zero terms in the polynomial at the given
        matrix entry. """

    def zero_out_small(self, eps: Number = NUMERICS_EPS) -> None:
        """ Zero out (set to zero) values that are smaller than ``eps``
        See also :py:data:`mdpoly.constants.NUMERICS_EPS`. """
        for entry in self.entries():
            for term in self.terms(entry):
                if self.is_zero(entry, term, eps=eps):
                    self.set_zero(entry, term)

    def __iter__(self) -> Iterable[tuple[MatrixIndex, PolyIndex, Number]]:
        """ Iterate over non-zero entries of the representations. """
        for entry in self.entries():
            for term in self.terms(entry):
                yield entry, term, self.at(entry, term)