import numpy as np
from typing import Union, List
import galois

__all__ = ["PauliOp"]

GF = galois.GF(2)


class PauliOp:
    """Class for representing Pauli operators as matrices and strings.

    Args
    ----
    op: Union[np.ndarray, str, List[str]]
        The Pauli operator to be represented. Can be a matrix, a string, or a list of strings.

    Examples
    --------
    Create a Pauli operator from a matrix

    .. ipython:: python

        op = mp.PauliOp(np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]))
        print(op)

    Create a Pauli operator from a string

    .. ipython:: python

        op = mp.PauliOp('XIZ;ZII;IIZ;IZI')
        print(op)

    Create a Pauli operator from a list of strings

    .. ipython:: python

        op = mp.PauliOp(['XIZ', 'ZII', 'IIZ', 'IZI'])
        print(op)

    Group
    -----
    operators
    """

    def __init__(self, op: Union[np.ndarray, str, List[str]]):
        """Initialize a PauliOp object."""

        if isinstance(op, np.ndarray):
            if op.shape[1] % 2 != 0:
                raise ValueError(
                    "Tableau representation must have an even number of columns"
                )
            self.mat = GF(op)
            self.txt = self._mat_to_txt(op)

        elif isinstance(op, str):
            op = op.replace(" ", "")
            if op[-1] == ";":
                op = op[:-1]

            op_list = op.split(";")

            if not all([len(op) == len(op_list[0]) for op in op_list]):
                raise ValueError("All Pauli operators must be the same length")

            self.txt = op.replace(";", "\n")
            self.mat = self._txt_to_mat(op_list)

        elif isinstance(op, list):
            if not all([len(pauliop) == len(op[0]) for pauliop in op]):
                raise ValueError("All Pauli operators must be the same length")

            self.txt = "\n".join(op)
            self.mat = self._txt_to_mat(op)

        else:
            raise ValueError(
                "PauliOp must be initialized with a string or a numpy array"
            )

        if self.mat.shape[0] == 1:
            self._array = [self]
        else:
            self._array = [PauliOp(op) for op in self.txt.split("\n")]

    def __repr__(self):
        return self.txt

    def __getitem__(self, key):
        paulis = self._array[key]
        if isinstance(key, slice):
            return PauliOp(";".join([pauli.txt for pauli in paulis]))
        return paulis

    def _mat_to_txt(self, op):
        n_qubits = op.shape[1] // 2
        txt = ""
        for i in range(op.shape[0]):
            for j in range(op.shape[1] // 2):
                if op[i, j] == 1 and op[i, j + n_qubits] == 0:
                    txt += "X"
                elif op[i, j] == 0 and op[i, j + n_qubits] == 1:
                    txt += "Z"
                elif op[i, j] == 1 and op[i, j + n_qubits] == 1:
                    txt += "Y"
                else:
                    txt += "I"
                if j == n_qubits - 1 and i != op.shape[0] - 1:
                    txt += "\n"
        return txt

    def _txt_to_mat(self, op_list):
        n_qubits = len(op_list[0])
        n_ops = len(op_list)

        mat = np.zeros((n_ops, 2 * n_qubits), dtype=int)

        for i, op in enumerate(op_list):
            for j, char in enumerate(op):
                if char == "X":
                    mat[i, j] = 1
                elif char == "Z":
                    mat[i, j + n_qubits] = 1
                elif char == "Y":
                    mat[i, j] = 1
                    mat[i, j + n_qubits] = 1
        return GF(mat)

    def commutator(self, other):
        """Returns the commutator of two Pauli operators."""
        return self.mat + other.mat

    def symplectic_prod(self, other):
        """Returns the symplectic product of two Pauli operators."""
        x1 = self.mat[:, : self.mat.shape[1] // 2]
        x2 = other.mat[:, : other.mat.shape[1] // 2]
        z1 = self.mat[:, self.mat.shape[1] // 2 :]
        z2 = other.mat[:, other.mat.shape[1] // 2 :]
        return x1 @ z2.T + z1 @ x2.T

    def append(self, other):
        """Appends a Pauli operator to the end of another Pauli operator."""
        self.mat = np.vstack((self.mat, other.mat))
        self.txt = self._mat_to_txt(self.mat)
