# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/07_shape.ipynb.

# %% auto 0
__all__ = ['ShapeMeta', 'Shape']

# %% ../nbs/07_shape.ipynb 4
from typing import Tuple, Any

# %% ../nbs/07_shape.ipynb 5
from .pttype import PTMeta, PTType

# %% ../nbs/07_shape.ipynb 7
class ShapeMeta(PTMeta):
    """
    Metaclass for Shape. It contains logic for instance checking.
    """

    __args__: Tuple[Any]
    _parameterized: bool

    def __getitem__(cls, item):
        if not isinstance(item, tuple):
            raise TypeError("Shape expects a tuple argument.")
        return super().__getitem__(item)

    def __str__(cls):
        shape_args = cls.__args__[0]
        return f"Shape[{shape_args}]"

# %% ../nbs/07_shape.ipynb 8
class Shape(PTType, metaclass=ShapeMeta):
    """
    Represents a shape expression, which can include integers, slices, ellipses, and asterisks.
    """

    __args__ = (Any,)

    def __init__(self, *args):
        self.args = args

    def __eq__(self, other):
        if isinstance(other, tuple):
            return self.match(other)
        elif isinstance(other, Shape):
            return self.args == other.args
        else:
            return False

    def match(self, shape):
        """
        Check if a shape matches this shape expression.
        """
        args = list(self.args)
        shape = list(shape)

        while args and shape:
            arg = args.pop(0)

            if arg == ...:
                if args:
                    while len(shape) >= len(args) and shape[-1] == args[-1]:
                        shape.pop()
                        args.pop()
                else:
                    shape = []
            elif arg == '*':
                shape.pop(0)
            elif isinstance(arg, int):
                if shape.pop(0) != arg:
                    return False
            else:
                return False

        return not args and not shape

    def broadcast(self, other):
        """
        Broadcast this shape expression with another shape or shape expression.
        """
        if isinstance(other, Shape):
            other = other.args
        elif not isinstance(other, tuple):
            raise TypeError("Can only broadcast with another shape or shape expression.")

        args = list(self.args)
        other = list(other)

        result = []
        while args or other:
            if args and other:
                arg = args.pop(0)
                other_arg = other.pop(0)

                if arg == other_arg or arg == '*' or other_arg == '*':
                    result.append(arg)
                elif arg == ...:
                    while len(other) >= len(args) and other[-1] == args[-1]:
                        result.append(other.pop())
                        args.pop()
                    result.append(arg)
                elif other_arg == ...:
                    while len(args) >= len(other) and args[-1] == other[-1]:
                        result.append(args.pop())
                        other.pop()
                    result.append(other_arg)
                else:
                    raise ValueError("Incompatible shapes for broadcasting.")
            elif args:
                result.append(args.pop(0))
            else:
                result.append(other.pop(0))

        return Shape(*result)

    def align(self, other):
        """
        Align this shape expression with another shape or shape expression.
        """
        if isinstance(other, Shape):
            other = other.args
        elif not isinstance(other, tuple):
            raise TypeError("Can only align with another shape or shape expression.")

        args = list(self.args)
        other = list(other)

        result = []
        while args or other:
            if args and other:
                arg = args.pop(0)
                other_arg = other.pop(0)
