from __future__ import annotations

import re

from functools import cached_property, cmp_to_key, lru_cache
from typing import (
    Any,
    cast,
    List,
    Literal,
    Optional,
    overload,
    Sequence,
    Tuple,
    TYPE_CHECKING,
    Union,
)

from sympy import Basic, Matrix, ones, Symbol, sympify

from ._backend import create_backend
from ._methods import (
    _basis_lookup,
    annotate_matrix,
    Basis,
    cartan_matrix,
    cocartan_matrix,
    quadratic_form,
    reflection_matricies,
)

BASIS = Literal["ortho", "omega", "alpha"]


class NumericSymbol(Symbol):
    """Extension of Sympy symbol that allows
    latex formatting but also tracks the underlying
    integer value. Useful for dimension representations
    of irreps"""

    _numeric_dim: int

    def __new__(cls, dim: int, fmtted_dim: str):
        obj = super().__new__(cls, fmtted_dim)
        obj._numeric_dim = int(dim)
        return obj

    def __len__(self):
        return len(str(self))

    @classmethod
    def from_symbol(cls, symbol: Symbol):
        r"""Converts from Symbol into NumericSymbol by
        regex search for digits from latex display pattern and
        returns a NumericSymbol. Will raise if no numeric is
        present in the symbol.


        Example
        =======
        >>> import liesym as ls
        >>> from sympy import Symbol
        >>> anti_fundamental = Symbol(r"\bar{4}")
        >>> assert NumericSymbol.from_symbol(anti_fundamental).numeric_dim == 4
        """
        try:
            s = symbol.__str__()
            num = re.findall(r"\d+", s)[0]
            return cls(int(num), s)
        except (IndexError, ValueError):
            raise ValueError("Could not extract numerical from Symbol")

    @property
    def numeric_dim(self) -> int:
        """Underlying integer value of this symbol"""
        return self._numeric_dim


class LieAlgebra(Basic):
    """The base class for all lie algebras. The methods and properties
    in this class are basis independent and apply in a general sense. In
    order to write down the roots as matricies and vectors, we choose a
    representation.

    """

    _series: str
    _simple_roots: list[Matrix]

    def __new__(cls, series: str, rank: int, simple_roots: list[Matrix]):
        """
        Returns a new instance of a Sympy object

        Args:
            series (str): The series type
            rank (int): The rank of the algebra
            simple_roots (list[Matrix]): Simple roots of the algebra
        """
        obj = super().__new__(cls, sympify(rank))
        obj._series = series

        assert (
            len(simple_roots) == rank
        ), "Number of simple roots do not match rank of algebra"

        obj._simple_roots = [sympify(x) for x in simple_roots]

        return obj

    @property
    def series(self) -> str:
        """Algebra series type"""
        return self._series

    @property
    def rank(self) -> int:
        """Algebra rank"""
        return self.args[0]  # type:ignore

    @property
    def dimension(self) -> int:
        """Algebra dimension

        Abstract
        """
        raise NotImplementedError()

    @property
    def n_pos_roots(self) -> int:
        """Total number of positive roots in the algebra

        Abstract
        """
        raise NotImplementedError()

    @property
    def n_roots(self) -> int:
        """Total number of roots in the algebra"""
        return 2 * self.n_pos_roots + self.rank

    @lru_cache
    def simple_roots(self, basis: BASIS = "ortho") -> list[Matrix]:
        """Returns a list of Sympy matrix (1,dimension)
        objects representing a chosen basis of the algebra.
        `simple_roots` define the entire representation of the algebra.


        Args:
            basis ("ortho" | "omega" | "alpha", optional): Basis to return simple roots in. Defaults to "ortho".
        """
        if basis == "ortho":
            return [annotate_matrix(x) for x in self._simple_roots]
        else:
            return [self._to_basis(x, basis, "ortho") for x in self.simple_roots()]

    @cached_property
    def cartan_matrix(self) -> Matrix:
        r"""For a given simple Lie algebra the elements $a_{ij}$ can be
        generated by

        .. math::
            a_{ji} = 2 \langle\alpha_i, \alpha_j\rangle / \langle\alpha_j, \alpha_j\rangle

        where $a_i$ is the i'th simple root and $\langle,\rangle$ is the scalar product.

        Sources:
            - https://en.wikipedia.org/wiki/Cartan_matrix
            - https://mathworld.wolfram.com/CartanMatrix.html

        Returns:
            Matrix: Cartan Matrix as a Sympy object
        """
        return cartan_matrix(self.simple_roots())

    @cached_property
    def cocartan_matrix(self) -> Matrix:
        """The cocartan matrix rows are generated from the coroots of
        the algebra such that multiplication by a simple root will
        generate a row of the cartan matrix.

        Returns:
            Matrix: Cocartan Matrix as a Sympy object
        """
        return cocartan_matrix(self.simple_roots())

    @cached_property
    def omega_matrix(self) -> Matrix:
        """The rows of the omega matrix are the fundamental weights
        of the algebra.

        Returns:
            Matrix: Omega Matrix as a Sympy object
        """
        return self.cocartan_matrix.pinv().T

    @cached_property
    def metric_tensor(self) -> Matrix:
        """Also known as the quadratic form, the metric tensor
        serves as the metrix for the inner product of two roots or weights
        when they are not in the orthogonal basis.

        Returns:
            Matrix: Metric Tensor as a Sympy object
        """
        return quadratic_form(self.cartan_matrix, self.simple_roots())

    @cached_property
    def reflection_matricies(self) -> List[Matrix]:
        """Returns a list of reflection matrices built from
        rotations about each simple root.

        Returns:
            list[Matrix]: list of Sympy Matrices
        """
        return reflection_matricies(self.simple_roots())

    @lru_cache
    def fundamental_weights(self, basis: BASIS = "ortho") -> List[Matrix]:
        """Returns the fundamental weights of the algebra.

        Args:
            basis ("ortho" | "omega" | "alpha", optional): Basis to return result in. Defaults to "ortho".

        Returns:
            list[Matrix]: list of Sympy Matrices
        """
        if basis == "ortho":
            return [
                annotate_matrix(self.omega_matrix.row(i))
                for i in range(self.omega_matrix.rows)
            ]
        else:
            return [
                self._to_basis(x, basis, "ortho") for x in self.fundamental_weights()
            ]

    @lru_cache
    def positive_roots(self, basis: BASIS = "ortho") -> List[Matrix]:
        """Returns the postive roots of the algebra. They are sorted
        first by their distance from the highest root and then by
        tuple ordering (convention).

        Args:
            basis ("ortho" | "omega" | "alpha", optional): Basis to return result in. Defaults to "ortho".

        Returns:
            list[Matrix]: list of Sympy Matrices
        """
        if basis == "ortho":
            return self.root_system()[: self.n_pos_roots]
        else:
            return [self._to_basis(x, basis, "ortho") for x in self.positive_roots()]

    @cached_property
    def _backend_instance(self):
        return create_backend(self)

    def orbit(
        self,
        weight: Matrix,
        stabilizers: Optional[Sequence[int]] = None,
        basis: BASIS = "ortho",
        weight_basis: BASIS = "ortho",
    ) -> List[Matrix]:
        """Returns the orbit of the weight or root by reflecting it
        a plane. A stabilizer may be passed to calculate the orbit using
        the Orbit-Stabilizer theorem.

        Basis: Ortho

        Args:
            weight (Matrix): A Matrix of shape (1, rank)
            stabilizer (Iterable of ints, optional): Per Orbit-Stabilizer
            theorem, integer iterable of simple root indexes. Defaults to None.
            basis ("ortho" | "omega" | "alpha", optional): Basis to return orbit in, Defaults to "ortho".
            weight_basis ("ortho" | "omega" | "alpha", optional): Basis of weight, this will be a fallback if weight is not annotated. Defaults to "ortho".

        Sources:
        - https://en.wikipedia.org/wiki/Coadjoint_representation#Coadjoint_orbit
        - https://en.wikipedia.org/wiki/Group_action#Orbits_and_stabilizers

        """

        return self._orbit(sympify(weight), stabilizers, basis, weight_basis)

    @lru_cache
    def _orbit(
        self,
        weight: Matrix,
        stabilizers: Optional[Sequence[int]] = None,
        basis: BASIS = "ortho",
        weight_basis: BASIS = "ortho",
    ) -> List[Matrix]:
        if basis == "ortho":
            weight = self.to_ortho(weight, weight_basis)
            return [
                self.to_ortho(x, "ortho")
                for x in self._backend_instance.orbit(weight, stabilizers)
            ]
        elif basis == "omega":
            return [
                self.to_omega(x, "ortho")
                for x in self.orbit(
                    weight, stabilizers=stabilizers, weight_basis=weight_basis
                )
            ]
        elif basis == "alpha":
            return [
                self.to_alpha(x, "ortho")
                for x in self.orbit(
                    weight, stabilizers=stabilizers, weight_basis=weight_basis
                )
            ]
        else:
            raise ValueError(f"Unsupported basis `{basis}`")

    def dim_name(self, irrep: Matrix, basis="omega") -> NumericSymbol:
        r"""Returns a sympy formatted symbol for the irrep.
        This is commonly used in physics literature. Returns
        a NumericSymbol object that is a simple extension of
        Symbol.

        Examples
        ========
        >>> from liesym import A
        >>> from sympy import Matrix
        >>> a3 = A(3)
        >>> assert str(a3.dim_name(Matrix([[1, 1, 0]]))) == '\\bar{20}'
        """
        hashable = sympify(irrep)
        annotate_matrix(hashable, basis)
        return self._dim_name(hashable, basis)

    @lru_cache
    def _dim_name(self, irrep: Matrix, basis="omega") -> NumericSymbol:
        """See docstring on LieAlgebra.dim_name for info, this is the caching
        version it calls"""
        irrep = self.to_omega(irrep, basis)

        dim = self.dim(irrep)
        max_dd = self.max_dynkin_digit(irrep)
        same_dim_irreps: list[Matrix] = self.get_irrep_by_dim(dim, max_dd=max_dd)
        num_primes = 0
        conjugate = 0
        so8label = ""

        if len(same_dim_irreps) > 1:
            # group by index
            index_pairs = {}  # type: ignore
            for i in same_dim_irreps:
                index = self._backend_instance.index_irrep(i, dim)
                index_pairs[index] = index_pairs.get(index, []) + [i]

            groups = [
                sorted(dimindex, key=cmp_to_key(self._dimindexsort))
                for dimindex in index_pairs.values()
            ]
            positions = []
            for id1, grps in enumerate(groups):
                for id2, g in enumerate(grps):
                    if g == irrep:
                        positions.append([id1, id2])
            [num_primes, conjugate] = positions[0]
            so8label = self._is_s08(irrep)
        has_conjugate = conjugate == 1 if so8label == "" else False
        return self._dim_name_fmt(dim, has_conjugate, num_primes, so8label)

    @lru_cache
    def irrep_lookup(
        self, dim: Any, max_dynkin_digit=5, basis: BASIS = "omega"
    ) -> Matrix:
        """Returns the irrep matrix for the dimension.

        Args:
            dim (Any): Can either be any sympy object.

        Raises:
            KeyError: Dim not found

        Returns:
            Matrix: Returns irrep in Omega basis


        Examples
        ========
        >>> from liesym import A
        >>> A3 = A(3)
        >>> A3.irrep_lookup(r"\\bar{4}")
        Matrix([[0, 0, 1]])
        >>> A3.irrep_lookup("4")
        Matrix([[1, 0, 0]])
        """
        if not isinstance(dim, Symbol):
            dim = Symbol(str(dim))
        if isinstance(dim, Symbol) and not isinstance(dim, NumericSymbol):
            dim = NumericSymbol.from_symbol(dim)

        n_dim = cast(NumericSymbol, dim).numeric_dim

        dd = 0
        while dd < max_dynkin_digit:
            dd += 1
            for c in self.get_irrep_by_dim(n_dim, max_dd=dd):
                if self.dim_name(c) == dim:
                    return self._to_basis(c, basis, "omega")
        raise KeyError(f"Irrep {dim} not found.")

    def conjugate(self, irrep: Matrix, basis: BASIS = "omega") -> Matrix:
        """Finds the conjugate irrep. If it is the same
        as the original irrep, you have a Real Irrep, otherwise
        it's a Complex Irrep.

        Examples
        ========
        .. code-block:: python

            from liesym import A,D
            from sympy import Matrix

            SU4 = A(3)
            irrep_20 = Matrix([[0,1,1]])
            irrep_20bar = Matrix([[1,1,0]])
            assert irrep_20 == SU4.conjugate(irrep_20bar)

            SO10 = D(5)
            irrep_10 = Matrix([[1, 0, 0, 0, 0]])
            assert irrep_10 == SO10.conjugate(irrep_10)
        """
        raw = self._backend_instance.conjugate(irrep)
        if len(raw) == 0:
            raise ValueError(f"Conjugate for {irrep} could not be found")
        return self._to_basis(raw[0], basis, "omega")

    def _is_s08(self, irrep):
        return ""

    def _dimindexsort(self, irrep1, irrep2):
        cong1 = self._congruency_class(irrep1)
        cong2 = self._congruency_class(irrep2)
        if isinstance(cong1, tuple):
            return 1 if cong1[-1] <= cong2[-1] else -1
        else:
            return -1 if cong1 <= cong2 else 1

    def _congruency_class(self, irrep):
        return 0

    def max_dynkin_digit(self, irrep: Matrix) -> int:  # type: ignore
        """Returns the max Dynkin Digit for the representations"""
        pass

    def _dim_name_fmt(self, dim: int, conj=False, primes=0, sub="") -> NumericSymbol:
        if conj:
            irrep = r"\bar{" + str(dim) + "}"
        else:
            irrep = str(dim)

        if primes > 0:
            irrep += r"^{" + " ".join([r"\prime"] * primes) + r"}"

        if sub != "":
            irrep += r"_{" + str(sub) + r"}"

        return NumericSymbol(dim, irrep)

    @overload
    def get_irrep_by_dim(
        self,
        dim: int,
        *,
        max_dd: int = 3,
        with_symbols: Literal[True],
        basis: BASIS = "omega",
    ) -> List[Tuple[Matrix, NumericSymbol]]:
        ...

    @overload
    def get_irrep_by_dim(
        self,
        dim: int,
        *,
        max_dd: int = 3,
        with_symbols: Literal[False] = False,
        basis: BASIS = "omega",
    ) -> List[Matrix]:
        ...

    @lru_cache  # type:ignore
    def get_irrep_by_dim(
        self,
        dim: int,
        max_dd: int = 3,
        with_symbols: bool = False,
        basis: BASIS = "omega",
    ) -> Union[List[Tuple[Matrix, NumericSymbol]], List[Matrix]]:
        r"""Gets all irreps by dimension and max dynkin digit. `max_dd` is . This algorithm brute forces searches by using `itertools.product`
        which can become expensive for large so searching max_dd > 3 will be
        very expensive

        Args:
            dim (int): Dimension to query
            max_dd (int, optional): The max dynkin digit to use. Defaults to 3.
            with_symbols (bool, optional): Returns list of tuples of rep and latex fmt. Defaults to False.
            basis ("ortho" | "omega" | "alpha", optional): Basis to return orbit in, Defaults to "ortho".

        Returns:
            list[Union[Matrix, Tuple[Matrix,NumericSymbol]]]: If `with_symbols=True` will return a list of tuples.


        Examples
        ========
        >>> from liesym import A
        >>> from sympy import Matrix
        >>> a3 = A(3)
        >>> expected = a3.get_irrep_by_dim(20)
        >>> result = [
        ... Matrix([[1, 1, 0]]),
        ... Matrix([[0, 1, 1]]),
        ... Matrix([[0, 2, 0]]),
        ... Matrix([[3, 0, 0]]),
        ... Matrix([[0, 0, 3]])]
        >>> assert expected == result
        >>> a3.get_irrep_by_dim(20, with_symbols=True)
        [(Matrix([[1, 1, 0]]), \bar{20}), (Matrix([[0, 1, 1]]), 20), (Matrix([[0, 2, 0]]), 20^{\prime}), (Matrix([[3, 0, 0]]), \bar{20}^{\prime \prime}), (Matrix([[0, 0, 3]]), 20^{\prime \prime})]
        """

        if basis == "omega":
            backend_results: list[Matrix] = (
                self._backend_instance.get_irrep_by_dim(dim, max_dd) or []
            )
            results = [self.to_omega(x, "omega") for x in backend_results]
            if with_symbols:
                return [(x, self.dim_name(x)) for x in results]
            return results
        else:
            omega_results = self.get_irrep_by_dim(
                dim=dim, max_dd=max_dd, with_symbols=False, basis="omega"
            )
            results = [self._to_basis(x, basis, "omega") for x in omega_results]
            if with_symbols:
                return [(x, self.dim_name(y)) for x, y in zip(results, omega_results)]
            return results

    def dim(self, irrep: Matrix, basis: BASIS = "omega") -> int:
        r"""Returns the dimension of the weight, root or irreducible representations.
        This follows Weyl's dimension formula:

        .. math::
            dim(w) = \prod_{\alpha\in\Delta^{+}} \frac{\langle \alpha, w + \rho\rangle}{\langle\alpha,\rho\rangle}

        where $\Delta^{+}$ are the positive roots and $rho$ is the sum of
        the positive roots: `[1] * rank`.

        Examples
        ========
        >>> from liesym import A
        >>> from sympy import Matrix
        >>> a2 = A(2)
        >>> assert a2.dim(Matrix([[1,0]])) == 3
        """
        irrep = self.to_omega(sympify(irrep), basis)
        return self._backend_instance.dim(irrep)

    @lru_cache
    def root_system(self, basis: BASIS = "omega") -> list[Matrix]:
        """Returns the entire rootsystem of the algebra. This
        includes the positive, negative and zeros of the algebra.

        Args:
            basis ("ortho" | "omega" | "alpha", optional): Basis of incoming weights. If not set, will implicitly set. Defaults to 'omega'.

        Returns:
            list[Matrix]: list of ordered roots.
        """
        if basis == "omega":
            return [
                self._to_basis(x, to_basis="omega", from_basis="omega")
                for x in self._backend_instance.root_system()
            ]
        else:
            return [
                self._to_basis(x, to_basis=basis, from_basis="omega")
                for x in self.root_system()
            ]

    def tensor_product_decomposition(
        self, *weights: Matrix, basis: BASIS = "omega"
    ) -> list[Matrix]:
        """Returns the tensor product between irreducible representations
        as a the tensor sum of the irreducible representations of their
        highest weights. This algorithm is based on Klimky's formula.


        Args:
            weights (Matrix: tuple[Matrix, ...]): Fundamental weights to take the tensor product between
            basis ("ortho" | "omega" | "alpha", optional): Basis of incoming weights and result. If not set, will implicitly set. Defaults to 'omega'.

        Returns:
            list[Matrix]: list of weights decomposed from the tensor product


        Examples
        ========
        >>> from liesym import A
        >>> from sympy import Matrix
        >>> a2 = A(2)
        >>> results = a2.tensor_product_decomposition([Matrix([[1,0]]), Matrix([[1,0]])])
        >>> print(results)
        [Matrix([[0, 1]]), Matrix([[2, 0]])]


        """

        if len(weights) < 2:
            raise ValueError(
                "At least two weights must be passed to take the tensor product"
            )

        w = [self.to_omega(x, basis) for x in weights]
        i = w.pop()
        j = w.pop()

        decomp = self._backend_instance.tensor_product_decomposition(i, j)

        while len(w) > 0:
            j = w.pop()
            results = []
            for i in decomp:
                # i,j reversed because pop takes from -1 index
                results += self._backend_instance.tensor_product_decomposition(j, i)
            decomp = results
        return [self._to_basis(x, basis, "omega") for x in decomp]

    def to_ortho(self, x: Matrix, basis=None) -> Matrix:
        """Rotates to orthogonal basis

        Args:
            x (Matrix): Matrix to be rotated
            basis (optional): If `basis` attribute is not set on `x` define it here. Defaults to None.

        Raises:
            ValueError: If no `x.basis` is set and None is passed to `basis` kwarg.

        Returns:
            Matrix: Matrix in orthogonal basis.
        """
        basis = _basis_lookup(basis)
        annotate_matrix(x, basis)

        if x.basis is Basis.ORTHO:  # type: ignore[attr-defined]
            r = x
        elif x.basis is Basis.OMEGA:  # type: ignore[attr-defined]
            r = x * self.omega_matrix
        elif x.basis is Basis.ALPHA:  # type: ignore[attr-defined]
            r = x * self.cartan_matrix * self.omega_matrix
        else:
            raise ValueError(
                "Basis arg cannot be None if attribute `basis` has not been set on Matrix."
            )

        r.basis = Basis.ORTHO  # type: ignore[attr-defined]
        return r

    def to_omega(self, x: Matrix, basis=None) -> Matrix:
        """Rotates to omega basis

        Args:
            x (Matrix): Matrix to be rotated
            basis (optional): If `basis` attribute is not set on `x` define it here. Defaults to None.

        Raises:
            ValueError: If no `x.basis` is set and None is passed to `basis` kwarg.

        Returns:
            Matrix: Matrix in omega basis.
        """
        annotate_matrix(x, basis)

        if x.basis is Basis.OMEGA:  # type: ignore[attr-defined]
            r = x
        elif x.basis is Basis.ORTHO:  # type: ignore[attr-defined]
            r = x * self.omega_matrix.pinv()
        elif x.basis is Basis.ALPHA:  # type: ignore[attr-defined]
            r = x * self.cartan_matrix
        else:
            raise ValueError(
                "Basis arg cannot be None if attribute `basis` has not been set on Matrix."
            )

        r.basis = Basis.OMEGA  # type: ignore[attr-defined]
        return r

    def to_alpha(self, x: Matrix, basis=None) -> Matrix:
        """Rotates to alpha basis

        Args:
            x (Matrix): Matrix to be rotated
            basis (optional): If `basis` attribute is not set on `x` define it here. Defaults to None.

        Raises:
            ValueError: If no `x.basis` is set and None is passed to `basis` kwarg.

        Returns:
            Matrix: Matrix in alpha basis.
        """
        basis = _basis_lookup(basis)
        annotate_matrix(x, basis)

        if x.basis is Basis.ALPHA:  # type: ignore[attr-defined]
            r = x
        elif x.basis is Basis.ORTHO:  # type: ignore[attr-defined]
            r = x * self.omega_matrix.pinv() * self.cartan_matrix.pinv()
        elif x.basis is Basis.OMEGA:  # type: ignore[attr-defined]
            r = x * self.cartan_matrix.pinv()
        else:
            raise ValueError(
                "Basis arg cannot be None if attribute `basis` has not been set on Matrix."
            )

        r.basis = Basis.ALPHA  # type: ignore[attr-defined]
        return r

    def _scalar_product(self, irrep1, irrep2, basis="ortho"):
        """Scalar product between two irreps."""
        irrep1 = self.to_ortho(irrep1, basis=basis)
        irrep2 = self.to_ortho(irrep2, basis=basis)

        return irrep1.dot(irrep2)

    def quadratic_casimir(
        self, irrep: Optional[Matrix] = None, basis: BASIS = "omega"
    ) -> Basic:
        r"""Returns the quadratic casimir for an arbitrary irreducible
        representation, $R$.

        .. math::

            \frac{1}{2}<R | R + \rho>

        Args:
            irrep (Matrix): Irreducible representation
            basis (str, optional): Basis of irrep. Defaults to "omega".

        Returns:
            Basic: Rational number
        """
        adj = self.adjoint_rep("omega")
        irrep = self.to_omega(irrep, basis) if irrep else adj
        rho = 2 * ones(1, self.rank)
        return self._scalar_product(irrep, irrep + rho, basis="omega") / 2

    def dynkin_index(
        self, irrep: Optional[Matrix] = None, basis: BASIS = "omega"
    ) -> Basic:
        r"""Calculates the dynkin index of an arbitrary irrep. This
        is also called in literature second-order index, $I_2(R)$. If no irrep
        is passed, basis and rep default to the algebra's adjoint rep.
        For arbitrary irrep, $R$, the dynkin index, $I_2(R)$ can be calculated
        as:

        .. math::
            I_2(R) = \frac{Dim(R)}{Dim(adj)}c(R)

        where the $dim$ is the dimension formula, $adj$ is the adjoint irrep
        and $c$ is the quadratic casimir (`Liealgebra.casimir`).
        """
        adj = self.adjoint_rep("omega")
        irrep = self.to_omega(irrep, basis) if irrep else adj
        d_r = cast(Basic, self.dim(irrep, basis="omega"))
        d_a = cast(Basic, self.dim(adj))
        c_r = self.quadratic_casimir(irrep, basis="omega")
        return d_r * c_r / d_a  # type:ignore[operator]

    def adjoint_rep(self, basis: BASIS = "ortho") -> Matrix:
        """Returns the adjoint irrep. This is the highest weight of the
        root system."""
        return self.positive_roots(basis)[0]

    def _to_basis(self, x: Matrix, to_basis: BASIS, from_basis: BASIS):
        if to_basis == "ortho":
            fn = self.to_ortho
        elif to_basis == "omega":
            fn = self.to_omega
        elif to_basis == "alpha":
            fn = self.to_alpha
        else:
            raise ValueError(f"Unsupported basis `{to_basis}`")

        return fn(x, from_basis)
