"""
模块: Operator
抽象基类，定义了抽象的量子算子，形如下式。
    $\\prod_{i=0}\\eta \\hat A_{i}^{a_i} \\hat B_{i}^{b_i}$
    $\\eta\\in\\{1, -1, i, -i\\}$
属性：
    occupy_x：哪些local site有A算子，也就是$a_i\neq 0$对应的$i$的集合
    occupy_z：哪些local site有B算子，也就是$b_i\neq 0$对应的$i$的集合
    coff：系数，只能取±1, ±i，也就是$\\eta$
方法：
    __init__：算子的抽象构造函数
    __matmul__：算子的抽象矩阵乘法
    __rmatmul__：算子的抽象右矩阵乘法
    __mul__：算子的抽象标量乘法
    __rmul__：算子的抽象右标量乘法
    __eq__：算子的抽象相等比较
    __neg__：算子的抽象取负
    __str__：算子的抽象字符串表示
    weight：算子的抽象权重
    is_hermitian：算子的抽象是否为厄米算子
    index_map：算子的抽象索引映射，将local site的索引映射到global site的索引
    get_vector：算子的抽象向量表示，将算子表示为一个向量
    copy：算子的抽象复制
    commute：算子的抽象commute关系
    get_matrix：算子的抽象矩阵表示，将算子表示为一个矩阵
    HermitianOperatorFromOccupy：从occupy_x和occupy_z构造厄米算子
    HermitianOperatorFromVector：从向量构造厄米算子
"""
import copy
from abc import abstractmethod, ABC
from typing import Iterable
import galois
import numpy as np
from numpy._typing import NDArray


class Operator(ABC):

    __slots__ = ['occupy_x', 'occupy_z', 'coff']

    ##  SECTION: ----构造函数----
    def __init__(self, occupy_x: Iterable[int], occupy_z: Iterable[int], coff: int | float | complex)-> None:
        """
        input.occupy_x：哪些local site有A算子
        input.occupy_z：哪些local site有B算子
        input.coff：系数，只能取±1, ±i
        output：None
        """

        ##  输入检查
        try:
            occupy_x = np.array([temp for temp in occupy_x], dtype=int)
            occupy_z = np.array([temp for temp in occupy_z], dtype=int)
        except Exception as e:
            raise ValueError("occupy_x and occupy_z must be a list of integers.") from e

        assert coff == -1 or coff == 1 or coff == 1j or coff == -1j, "coff must be 1, -1, 1j, or -1j."

        ##  赋值
        self.occupy_x:NDArray[np.integer] = occupy_x
        self.occupy_z:NDArray[np.integer] = occupy_z
        self.occupy_x.sort()
        self.occupy_z.sort()
        self.coff:complex = coff


    ##  SECTION: ----抽象矩阵乘法----
    @abstractmethod
    def __matmul__(self, other: 'Operator') -> 'Operator':
        """
        input.other：另一个同类算子，乘在右边
        output：算子乘积得到的算子
        """

        pass


    ##  SECTION: ----抽象右矩阵乘法----
    @abstractmethod
    def __rmatmul__(self, other: 'Operator') -> 'Operator':
        """
        input.other：另一个同类算子，乘在左边
        output：算子乘积得到的算子
        """

        pass


    ##  SECTION: ----抽象标量乘法----
    def __mul__(self, other: int | float | complex):
        """
        input.other：系数（±1, ±i）
        output：被乘了个系数的新算子
        """

        return self.__class__(self.occupy_x, self.occupy_z, self.coff * other)


    ##  SECTION: ----抽象右标量乘法----
    def __rmul__(self, other: int | float | complex):
        """
        input.other：系数（±1, ±i）
        output：被乘了个系数的新算子
        """

        return self.__class__(self.occupy_x, self.occupy_z, other * self.coff)


    ##  SECTION: ----比较两个算符是否相等----
    def __eq__(self, other:'Operator') -> bool:
        """
        input.other：另一个同类算子.
        output：相等判断结果
        """

        ##  输入检查
        assert isinstance(other, Operator), "other must be an operator."

        ##  比较
        if type(other)==type(self):
            return np.all(self.occupy_x == other.occupy_x) and np.all(self.occupy_z == other.occupy_z) and self.coff == other.coff
        else:
            return False


    ##  SECTION: ----取相反数----
    def __neg__(self):
        """
        output：新算子，系数取相反数
        """

        ##  取相反数
        return self.__class__(self.occupy_x, self.occupy_z, -self.coff)


    ##  SECTION: ----字符表示----
    @abstractmethod
    def __str__(self) -> str:
        """
        output：算子的字符表示
        """

        pass


    ##  SECTION: ----权重大小----
    @property
    def weight(self) -> int:
        """
        output: 算符中含有的local operators的数目
        """

        return len(self.occupy_x) + len(self.occupy_z)


    ##  SECTION: ----是否为厄米算子----
    @property
    @abstractmethod
    def is_hermitian(self) -> bool:
        """
        output：是否为厄米算子的判断结果
        """

        pass


    ##  SECTION: ----求对偶算符----
    @property
    def dual(self) -> 'Operator':
        """
        output：算子的对偶算符
        """

        return self.__class__(self.occupy_z.copy(), self.occupy_x.copy(), self.coff)


    ##  SECTION：----求当前算符中索引最大的值----
    @property
    def max_index(self)->int:
        """
        output：当前算符中索引最大的值
        """
        return np.hstack([self.occupy_x,self.occupy_z]).max()


    ##  SECTION: ----更改每一个索引----
    def index_map(self, indices) -> None:
        """
        input.indices：新的索引映射数组或列表，索引下标代表原索引，索引值代表新索引
        output：None
        influence：在原对象上直接修改索引到新的表示上
        示例：
            将$i\\hat X_0\\hat X_2\\hat X_3\\hat Z_3$从[0,1,2,3]映射到[3,2,1,0]：
            op.index_map([3,2,1,0])
            结果：$i\\hat X_3\\hat X_1\\hat X_0\\hat Z_0$
        """

        ##  检查输入
        try:
            indices = np.array([temp for temp in indices], dtype=int)
            assert len(indices)-1 >= self.max_index
        except Exception as e:
            raise ValueError("新索引不合法") from e

        ##  修改索引
        occupy_x_new = np.array([indices[i] for i in self.occupy_x], dtype=int)
        occupy_z_new = np.array([indices[i] for i in self.occupy_z], dtype=int)
        self.occupy_x = occupy_x_new
        self.occupy_z = occupy_z_new


    ##  SECTION: ----获取算子的向量表示----
    def get_vector(self, number: int)->galois.GF2:
        """
        input.number：固定算子所在系统的local sites数量
        output：算子的向量表示，奇数位代表$(n-1)//2$处有无A，偶数位代表$n//2$处有无B
        """

        ##  检查输入
        try:
            number = int(number)
        except Exception as e:
            raise ValueError("输入必须是整数") from e

        ##  构造向量
        vector = galois.GF2(np.zeros(number * 2, dtype=int))
        vector[self.occupy_x * 2] = 1
        vector[self.occupy_z * 2 + 1] = 1

        ##  返回结果
        return vector


    ##  SECTION: ----复制函数----
    def copy(self):
        """
        output：新的算子对象，与原算子系数、索引、向量表示均相同
        """

        return copy.deepcopy(self)


    ##  SECTION: ----判断两个算符之间是否对易----
    @staticmethod
    @abstractmethod
    def commute(A, B) -> bool:
        """
        input.A：第一个算符
        input.B：第二个算符
        output：对易性判断结果
        """

        pass


    ##  SECTION: ----获取算子矩阵表示----
    @staticmethod
    def get_matrix(ops:Iterable['Operator'], number: int):
        """
        input.ops: 算子列表，每个算子都是Operator的子类
        input.number: 系统所含有的local sites总数
        output：m*2n矩阵，m为算子列表的长度，n为系统所含有的local sites总数，奇数位代表$(n-1)//2$处有无A，偶数位代表$n//2$处有无B
        """

        ##  检查输入
        try:
            number = int(number)
        except Exception as e:
            raise ValueError("number must be an integer.") from e
        try:
            np.all([isinstance(temp, Operator) for temp in ops])
        except Exception as e:
            raise ValueError("ops必须是算符对象.") from e

        ##  竖直堆叠向量得到矩阵
        matrix = None
        for i, op in enumerate(ops):
            vector = op.get_vector(number)
            if i == 0:
                matrix = vector
            else:
                matrix = np.vstack([matrix, vector])

        if matrix is not None and matrix.ndim == 1:
            matrix = matrix.reshape(1, -1)

        ##  返回矩阵
        return matrix


    ##  SECTION: ----抽象方法：从索引创建厄米算子----
    @staticmethod
    @abstractmethod
    def HermitianOperatorFromOccupy(occupy_x, occupy_z) -> 'Operator':
        """
        input.occupy_x：哪些local site有A算子
        input.occupy_z：哪些local site有B算子
        output：新的厄米算子对象，根据weight确定系数为1或i来保证厄米性
        """

        pass


    ##  SECTION: ----抽象方法：从向量创建厄米算子----
    @staticmethod
    @abstractmethod
    def HermitianOperatorFromVector(vector) -> 'Operator':
        """
        input.vector：长度为2n的向量，奇数位代表$(n-1)//2$处有无A，偶数位代表$n//2$处有无B
        output：新的厄米算子对象，根据weight确定系数为1或i来保证厄米性
        """

        pass


class MajoranaOperator(Operator):

    ##  SECTION: ----构造函数----
    def __init__(self, occupy_x, occupy_z, coff):
        """
        input.occupy_x：哪些local site有$\\hat\\gamma$算子
        input.occupy_z：哪些local site有$\\hat\\gamma^\\prime$算子
        input.coff：系数，只能取±1, ±i
        output：None
        """

        ##  直接构造父类
        super().__init__(occupy_x, occupy_z, coff)


    ##  SECTION: ----矩阵乘积----
    def __matmul__(self, other:'MajoranaOperator')->'MajoranaOperator':
        """
        input.other：另一个同类算子，乘在右边
        output：算子乘积得到的算子
        """

        ##  检查输入
        assert isinstance(other, MajoranaOperator)

        ##  计算新的occupy
        occupy_x = np.setxor1d(self.occupy_x, other.occupy_x, assume_unique=True)
        occupy_z = np.setxor1d(self.occupy_z, other.occupy_z, assume_unique=True)
        self_occupy = np.append(self.occupy_x * 2, self.occupy_z * 2 + 1)
        other_occupy = np.append(other.occupy_x * 2, other.occupy_z * 2 + 1)
        self_occupy = np.sort(self_occupy)
        other_occupy = np.sort(other_occupy)

        ##  计算新的系数
        exchange_times = np.sum([np.count_nonzero(self_occupy > temp) for temp in other_occupy])
        if exchange_times % 2 == 1:
            factor = -1
        else:
            factor = 1

        ##  返回新的算子
        return MajoranaOperator(occupy_x, occupy_z, self.coff * other.coff * factor)


    ##  SECTION: ----右矩阵乘积----
    def __rmatmul__(self, other:'MajoranaOperator')->'MajoranaOperator':
        """
        input.other：另一个同类算子，乘在左边
        output：算子乘积得到的算子
        """

        ##  计算并返回新的算子
        return other.__matmul__(self)


    ##  SECTION: ----算符的字符串表示----
    def __str__(self)->str:
        """
        output：算符的字符串表示
        """

        ##  初始化算子字符串
        op_str=f"{self.coff}"  # 放在最前面的系数
        flag_z=0  # 用于记录当前处理到的Z算子的索引
        flag_x=0  # 用于记录当前处理到的X算子的索引

        ##  按从左到右、从X到Z的顺序填充算子字符串
        while flag_x<len(self.occupy_x) or flag_z<len(self.occupy_z):
            if len(self.occupy_z)==0 or self.occupy_x[flag_x]>=self.occupy_z[flag_z]:
                op_str+=f"\\hat\\gamma_{self.occupy_x[flag_x]}"
                flag_x+=1
            else:
                op_str+=f"\\hat\\gamma^\\prime_{self.occupy_z[flag_z]}"
                flag_z+=1

        ##  返回结果
        return op_str


    ##  SECTION: ----判断算符是否厄米----
    @property
    def is_hermitian(self)->bool:
        """
        output：判断厄米性的结果
        """

        ##  检查厄米性
        if np.mod(self.weight * (self.weight - 1) // 2,2) == 0:
            if not (self.coff == 1 or self.coff == -1):
                return False
        else:
            if not (self.coff == 1j or self.coff == -1j):
                return False

        ##  返回结果
        return True


    ##  SECTION: ----判断两个算符之间是否对易----
    @staticmethod
    def commute(A:'MajoranaOperator', B:'MajoranaOperator')->bool:
        """
        input.A：第一个算符
        input.B：第二个算符
        output：对易性判断结果
        """

        ##  检查输入
        assert isinstance(A, MajoranaOperator) and isinstance(B, MajoranaOperator)

        ##  计算重叠部分
        overlap_x = len(np.intersect1d(A.occupy_x, B.occupy_x))
        overlap_z = len(np.intersect1d(A.occupy_z, B.occupy_z))
        weight = (len(A.occupy_x) + len(A.occupy_z)) * (len(B.occupy_x) + len(B.occupy_z))
        judge = overlap_x + overlap_z + weight

        ##  检查对易性的结果
        return np.mod(judge, 2) == 0


    ##  SECTION: ----根据索引生成厄米的Majorana operator----
    @staticmethod
    def HermitianOperatorFromOccupy(occupy_x:list[int]|np.ndarray,occupy_z:list[int]|np.ndarray)->'MajoranaOperator':
        """
        input.occupy_x：哪些local site有$\\hat\\gamma$算子
        input.occupy_z：哪些local site有$\\hat\\gamma^\\prime$算子
        output：根据索引生成的厄米的Majorana operator
        """

        ##  检查输入
        try:
            occupy_x=np.array([temp for temp in occupy_x],dtype=int)
            occupy_z=np.array([temp for temp in occupy_z],dtype=int)
        except Exception as e:
            raise ValueError("occupy_x and occupy_z must be a list of integers.") from e
        assert np.all(occupy_x >= 0), "occupy_x must be non-negative integers."
        assert np.all(occupy_z >= 0), "occupy_z must be non-negative integers."

        ##  计算系数
        weight = len(occupy_x) + len(occupy_z)
        if (weight * (weight - 1) // 2) % 2 == 0:
            coff = 1
        else:
            coff = 1j

        ##  返回新的算子
        return MajoranaOperator(occupy_x, occupy_z, coff)


    ##  SECTION: ----根据向量生成厄米的Majorana operator----
    @staticmethod
    def HermitianOperatorFromVector(vector)->'MajoranaOperator':
        """
        input.vector：Majorana operator的向量表示，偶数位为$\\hat\\gamma$算子，奇数位为$\\hat\\gamma^\\prime$算子
        output：根据向量生成的厄米的Majorana operator
        """

        ##  检查输入
        try:
            GF2=galois.GF(2)
            vector=GF2(np.array([temp for temp in vector],dtype=int))
        except Exception as e:
            raise ValueError("vector must be a list of integers.") from e

        ##  计算索引
        occupy_x = np.where(vector[0::2] == 1)[0]
        occupy_z = np.where(vector[1::2] == 1)[0]

        ##  返回新的算子
        return MajoranaOperator.HermitianOperatorFromOccupy(occupy_x, occupy_z)


class PauliOperator(Operator):

    ##  SECTION: ----构造函数----
    def __init__(self, occupy_x, occupy_z, coff):
        """
        input.occupy_x：哪些local site有$\\hat X$算子
        input.occupy_z：哪些local site有$\\hat Z^\\prime$算子
        input.coff：系数，只能取±1, ±i
        output：None
        """
        super().__init__(occupy_x, occupy_z, coff)


    ##  SECTION: ----Matrix product----
    def __matmul__(self, other:'PauliOperator')->'PauliOperator':
        """
        input.other：另一个同类算子，乘在右边
        output：算子乘积得到的算子
        """

        ##  检查输入
        assert isinstance(other, PauliOperator)

        ##  计算新的occupy
        occupy_x = np.setxor1d(self.occupy_x, other.occupy_x, assume_unique=False)
        occupy_z = np.setxor1d(self.occupy_z, other.occupy_z, assume_unique=False)
        exchange_times = np.sum([np.count_nonzero(self.occupy_z == temp) for temp in other.occupy_x])

        ##  计算新的系数
        if exchange_times % 2 == 1:
            factor = -1
        else:
            factor = 1

        ##  返回新的算子
        return PauliOperator(occupy_x, occupy_z, self.coff * other.coff * factor)


    ##  SECTION: ----右矩阵乘积----
    def __rmatmul__(self, other:'PauliOperator')->'PauliOperator':
        """
        input.other：另一个同类算子，乘在左边
        output：算子乘积得到的算子
        """

        ##  计算并返回新的算子
        return other.__matmul__(self)


    ##  SECTION: ----Print----
    def __str__(self)->str:
        """
        output：算符的字符串表示
        """

        ##  初始化算子字符串
        op_str=f"{self.coff}"  # 放在最前面的系数
        flag_z=0  # 用于记录当前处理到的Z算子的索引
        flag_x=0  # 用于记录当前处理到的X算子的索引

        ##  按从左到右，从X到Z的顺序填充算子字符串
        while flag_x<len(self.occupy_x) or flag_z<len(self.occupy_z):
            if len(self.occupy_z)==0 or self.occupy_x[flag_x]>=self.occupy_z[flag_z]:
                op_str+=f"\\hat X_{self.occupy_x[flag_x]}"
                flag_x+=1
            else:
                op_str+=f"\\hat Z_{self.occupy_z[flag_z]}"
                flag_z+=1

        ##  返回结果
        return op_str


    ##  SECTION: ----判断算符是否厄米----
    @property
    def is_hermitian(self)->bool:
        """
        output：判断厄米性的结果
        """

        ##  检查厄米性
        if len(np.intersect1d(self.occupy_x, self.occupy_z, assume_unique=False)) % 2 == 0:
            if not (self.coff == 1 or self.coff == -1):
                return False
        else:
            if not (self.coff == 1j or self.coff == -1j):
                return False

        ##  返回结果
        return True


    ##  SECTION: ----判断两个算符之间是否对易----
    @staticmethod
    def commute(A:'PauliOperator',B:'PauliOperator')->bool:
        """
        input.A：第一个算符
        input.B：第二个算符
        output：对易性判断结果
        """

        ##  检查输入
        assert isinstance(A, PauliOperator) and isinstance(B, PauliOperator)

        ##  计算重叠部分
        same_time=np.sum([np.count_nonzero(A.occupy_z==temp) for temp in B.occupy_x])
        same_time+=np.sum([np.count_nonzero(A.occupy_x==temp) for temp in B.occupy_z])

        ##  返回结果
        return np.mod(same_time, 2) == 0


    ##  SECTION: ----根据索引生成厄米的Majorana operator----
    @staticmethod
    def HermitianOperatorFromOccupy(occupy_x:list[int]|np.ndarray,occupy_z:list[int]|np.ndarray)->'PauliOperator':
        """
        input.occupy_x：哪些local site有$\\hat\\gamma$算子
        input.occupy_z：哪些local site有$\\hat\\gamma^\\prime$算子
        output：根据索引生成的厄米的Majorana operator
        """

        ##  检查输入
        try:
            occupy_x=np.array([temp for temp in occupy_x],dtype=int)
            occupy_z=np.array([temp for temp in occupy_z],dtype=int)
        except Exception as e:
            raise ValueError("occupy_x and occupy_z must be a list of integers.") from e

        ##  计算系数
        weight=len(np.intersect1d(occupy_x,occupy_z, assume_unique=False))
        if weight % 2 == 0:
            coff=1
        else:
            coff=1j

        ##  返回结果
        return PauliOperator(occupy_x,occupy_z,coff)


    ##  SECTION: ----根据向量生成厄米的Majorana operator----
    @staticmethod
    def HermitianOperatorFromVector(vector)->'PauliOperator':
        """
        input.vector：Majorana operator的向量表示，偶数位为$\\hat\\gamma$算子，奇数位为$\\hat\\gamma^\\prime$算子
        output：根据向量生成的厄米的Majorana operator
        """

        ##  检查输入
        try:
            GF2=galois.GF(2)
            vector=GF2(np.array([temp for temp in vector],dtype=int))
        except Exception as e:
            raise ValueError("vector must be a list of integers.") from e

        ##  计算索引
        pauli_x = np.where(vector[0::2] == 1)[0]
        pauli_z = np.where(vector[1::2] == 1)[0]

        ##  返回新的算子
        return PauliOperator.HermitianOperatorFromOccupy(pauli_x, pauli_z)
