"""
模块: QuantumCode
定义了量子码，形如下式
$$
\\begin{cases}
\\mathcal{S}_X=\\text{span}(\\{\\hat S_i\\})\\\\
\\mathcal{S}_Z=\\text{span}(\\{\\hat S_i\\})\\\\
\\end{cases}
$$
这里$\\hat S_i$是stabilizer generators，也就是某一类Operator
------------------------------------------------------------------------------------------------
属性：
    QuantumCode的属性：
        self.generators：Stabilizer generators
        self.physical_number：物理模式数
        self.checker_number：Stabilizer generators数目
        self._logical_operators：逻辑算子，计算预存
        self._distance：码距，计算预存

        PauliCode的属性：
            None

        MajoranaCode的属性：
            None

        QuantumCSSCode的属性：
            self.generators_x：X-type stabilizer generators
            self.generators_z：Z-type stabilizer generators
            self.checker_number_x：X-type stabilizer generators的数量
            self.checker_number_z：Z-type stabilizer generators的数量
            self._logical_operators_x：X-type logical operators
            self._logical_operators_z：Z-type logical operators
            self._distance_x：X-type stabilizer generators的码距，计算预存
            self._distance_z：Z-type stabilizer generators的码距，计算预存

            MajoranaCSSCode的属性：
                None

            PauliCSSCode的属性：
                None
-------------------------------------------------------------------------
方法：
    QuantumCode的方法：
        __init__（魔法方法）：量子码的构造函数
        check_matrix（属性方法）：量子CSS码的矩阵表示，将量子CSS码表示为一个矩阵
        rank（属性方法）：量子CSS码的rank，也就是Stabilizer generators的线性无关组的维数
        logical_number（属性方法）：逻辑模式数
        even_or_odd（属性方法）：逻辑算子的奇偶性，也就是有无odd-weight logical operator
        index_map：量子码的索引映射，将local site的索引映射到新的编号的site的索引
        copy：量子码的复制
        FromCheckMatrix（类方法）：从校验矩阵构造量子码

        PauliCode的方法：
            __init__（魔法方法，重构自QuantumCode）：Pauli stabilizer code的构造函数

        MajoranaCode的属性：
            __init__（魔法方法，重构自QuantumCode）：Majorana stabilizer code的构造函数

        QuantumCSSCode的方法：
            __init__（魔法方法，重构自QuantumCode）：CSS code的构造函数
            distance（重构自QuantumCode）：码距，也就是最小的非零子空间维数
            logical_operators（重构自QuantumCode）：逻辑算子，处于logical subspace的一组线性无关算子
            check_matrix_x（属性方法）：X-type stabilizer generators的矩阵表示
            check_matrix_z（属性方法）：Z-type stabilizer generators的矩阵表示
            rank_x（属性方法）：X-type stabilizer generators的rank
            rank_z（属性方法）：Z-type stabilizer generators的rank
            logical_operators_x（属性方法）：X-type logical operators
            logical_operators_z（属性方法）：Z-type logical operators
            distance_x（属性方法）：X-type stabilizer generators的码距，也就是最小的非零子空间维数
            distance_z（属性方法）：Z-type stabilizer generators的码距，也就是最小的非零子空间维数
            FromCheckMatrix（类方法，重构自QuantumCode）：从校验矩阵构造量子码

            MajoranaCSSCode的方法：
                __init__（魔法方法，重构自QuantumCSSCode）：Majorana CSS code的构造函数
                logical_operators_z（属性方法，重构自QuantumCSSCode）：Z-type logical operators
                distance_z（属性方法，重构自QuantumCSSCode）：Z-type stabilizer generators的码距，也就是最小的非零子空间维数

            PauliCSSCode的属性：
                __init__（魔法方法，重构自QuantumCSSCode）：Pauli CSS code的构造函数
-------------------------------------------------------------------------
例子：
    $S_0=i\\hat X_0\\hat X_1\\hat X_2\\hat X_3$
    $S_1=i\\hat X_1\\hat X_2\\hat X_5\\hat X_6$
    $S_2=i\\hat X_0\\hat X_2\\hat X_4\\hat X_6$
    $S_3=i\\hat Z_0\\hat Z_1\\hat Z_2\\hat Z_3$
    $S_4=i\\hat Z_1\\hat Z_2\\hat Z_5\\hat Z_6$
    $S_5=i\\hat Z_0\\hat Z_2\\hat Z_4\\hat Z_6$
    n=7
"""
import copy
from typing import Iterable
import galois
import numpy as np
from clifforder.Code.tools import mip_distance_caculator, orthogonalize, solve
from clifforder.Code.Operator import Operator, MajoranaOperator, PauliOperator


class QuantumCode:

    _OpType=None

    ##  SECTION: ----构造函数----
    def __init__(self,generators:Iterable[Operator],physical_number:int):
        """
        input.generators：Stabilizer generators
        input.physical_number：物理模式数目
        output：None
        """

        ##  检查输入
        try:
            generators = [temp for temp in generators]
            physical_number = int(physical_number)
            assert np.all([isinstance(temp, Operator) for temp in generators]), "generators必须都是算符"
            assert np.all([temp.max_index<physical_number for temp in generators]), "generators的最大索引必须小于physical_number"
        except Exception as e:
            raise ValueError("输入参数有误") from e

        ##  赋值
        self.generators=generators
        self.physical_number=physical_number
        self.checker_number=len(self.generators)
        self._logical_operators=None
        self._distance=None

        ##  检查对易性
        assert np.all(self.check_matrix@self.check_matrix.T==0), "generators必须是对易的"

    ##  SECTION: ----Get check matrix----
    @property
    def check_matrix(self)->galois.GF2:
        """
        output：校验矩阵
        """

        ##  校验矩阵
        return Operator.get_matrix(self.generators, self.physical_number)


    ##  SECTION: ----稳定子秩----
    @property
    def rank(self)->int:
        """
        output：稳定子秩，也就是约束个数
        """

        ##  稳定子秩
        return np.linalg.matrix_rank(self.check_matrix)


    ##  SECTION: ----逻辑模式数----
    @property
    def logical_number(self)->int:
        """
        output：逻辑模式数，也就是物理模式数减去稳定子秩
        """

        ##  逻辑模式数
        return self.physical_number-self.rank


    ##  SECTION: ----码距----
    @property
    def distance(self)->int:
        """
        output：Stabilizer generators的码距
        """

        ##  如果没有计算需要先计算
        if self._distance is not None:
            self._distance = mip_distance_caculator(self.check_matrix, self.logical_operators)

        ##  返回结果
        return self._distance


    ##  SECTION: ----逻辑算子----
    @property
    def logical_operators(self)->list[PauliOperator|MajoranaOperator]:
        ##  计算逻辑算子
        if self._logical_operators is None:

            ##  从零空间筛选与行空间线性无关的基矢
            codewords = self.check_matrix.null_space()  # 零空间
            independent_null_basis_list = []  # 初始化线性无关基矢组
            matrix = self.check_matrix  # 行空间和已有基矢合成的基矢组
            for vec in codewords:
                rank_before = np.linalg.matrix_rank(matrix)
                matrix = np.vstack([matrix, vec])
                if np.linalg.matrix_rank(matrix) == rank_before + 1:
                    independent_null_basis_list.append(vec)

            ##  Odd codes需要正交归一化
            if self.even_or_odd == 'even':
                basis_list = independent_null_basis_list
            else:
                basis_list = orthogonalize(independent_null_basis_list)
            self._logical_operators = [self._OpType.HermitianOperatorFromOccupy([],np.where(basis_list[i] == 1)[0]) for i in range(len(basis_list))]

        ##  返回逻辑算子
        return self._logical_operators


    ##  SECTION: ----更新索引方式----
    def index_map(self, index_map,physical_number):
        """
        input.indices：新的索引映射数组或列表，索引下标代表原索引，索引值代表新索引
        input.physical_number：新的物理模式数
        output：None
        influence：在原对象上直接修改索引到新的表示上
        """

        ##  检查输入
        try:
            indices = np.array([temp for temp in index_map], dtype=int)
            assert len(indices)==self.physical_number, "新索引映射的长度必须等于原来的物理模式数"
            assert np.array(indices,dtype=int).max()<=physical_number-1, "新索引映射的最大值必须小于新的物理模式数"
        except Exception as e:
            raise ValueError("新索引不合法") from e

        ##  更新索引
        for generator in self.generators:
            generator.index_map(index_map)
        self.physical_number=physical_number


    ##  SECTION: ----复制函数----
    def copy(self):
        """
        output：新的量子码对象，与原量子码完全相同
        """

        ##  复制
        return copy.deepcopy(self)


    ##  SECTION: ----从校验矩阵构造量子码----
    @classmethod
    def FromCheckMatrix(cls,check_matrix,check_matrix_z=None):
        """
        input.check_matrix：整个code的校验矩阵
        output：Code的实例
        """

        ##  检查输入
        try:
            check_matrix = galois.GF2(np.array(check_matrix, dtype=int))
            assert np.all(check_matrix @ check_matrix.T == 0), "check matrix必须是self-dual matrix"
        except Exception as e:
            raise ValueError("check matrix不合法") from e

        ##  从check matrix中提取stabilizer generators
        generators = []
        for i in range(len(check_matrix)):
            generators.append(cls._OpType.HermitianOperatorFromVector(check_matrix[i]))

        ##  构造code
        return cls(generators, check_matrix.shape[1])


    ##  SECTION: ----判断奇偶性----
    @property
    def even_or_odd(self):
        """
        output：根据逻辑算子的奇偶性判断结果，"even" or "odd"
        """

        ##  观察ones是否在check matrix的行空间中
        H=self.check_matrix
        ones=galois.GF2(np.ones(H.shape[1],dtype=int))
        if solve(H,ones) is None:
            return "odd"
        else:
            return "even"


class QuantumCSSCode(QuantumCode):
    """
    An example:
    $S_0=i\\hat X_0\\hat X_1\\hat X_2\\hat X_3$
    $S_1=i\\hat X_1\\hat X_2\\hat X_5\\hat X_6$
    $S_2=i\\hat X_0\\hat X_2\\hat X_4\\hat X_6$
    $S_3=i\\hat Z_0\\hat Z_1\\hat Z_2\\hat Z_3$
    $S_4=i\\hat Z_1\\hat Z_2\\hat Z_5\\hat Z_6$
    $S_5=i\\hat Z_0\\hat Z_2\\hat Z_4\\hat Z_6$
    n=7
    """

    ##  SECTION: ----构造函数----
    def __init__(self, generators_x: Iterable[Operator], generators_z: Iterable[Operator], physical_number: int):
        """
        input.generators_x：list of Operator, X stabilizers
        input.generators_z：list of Operator, Z stabilizers
        input.physical_number：number of physical modes
        """

        ##  检查输入
        try:
            generators_x = [temp.copy() for temp in generators_x]
            generators_z = [temp.copy() for temp in generators_z]
            physical_number = int(physical_number)
            assert np.all([isinstance(temp, Operator) for temp in generators_x]), "generators_x必须都是算符"
            assert np.all([temp.max_index < physical_number for temp in generators_x]), "generators_x的最大索引必须小于physical_number"
            assert np.all([isinstance(temp, Operator) for temp in generators_z]), "generators_z必须都是算符"
            assert np.all([temp.max_index < physical_number for temp in generators_z]), "generators_z的最大索引必须小于physical_number"
        except Exception as e:
            raise ValueError("输入参数有误") from e

        ##  赋值
        self.generators_x: list[MajoranaOperator | PauliOperator] = generators_x
        self.generators_z: list[MajoranaOperator | PauliOperator] = generators_z
        self.checker_number_x = len(generators_x)
        self.checker_number_z = len(generators_z)
        self._logical_operators_x = None
        self._logical_operators_z = None
        self._distance_x = None
        self._distance_z = None
        super().__init__(generators_x + generators_z, physical_number)


    ##  SECTION: ----获取逻辑算子----
    @property
    def logical_operators(self) -> list[PauliOperator|MajoranaOperator]:
        """
        output：逻辑算符基组
        """
        return self._logical_operators_x + self._logical_operators_z


    ##  SECTION: ----获取X-type logical operators----
    @property
    def logical_operators_x(self) -> list[PauliOperator|MajoranaOperator]:
        """
        output：np.array[MajoranaOperator]，independent \\hat gamma logical operators
        """

        ##  计算逻辑算子
        if self._logical_operators_x is None:

            ##  从零空间筛选与行空间线性无关的基矢
            codewords = self.check_matrix_x.null_space()  # 零空间
            independent_null_basis_list = []  # 初始化线性无关基矢组
            matrix = self.check_matrix_x  # 行空间和已有基矢合成的基矢组
            for vec in codewords:
                rank_before = np.linalg.matrix_rank(matrix)
                matrix = np.vstack([matrix, vec])
                if np.linalg.matrix_rank(matrix) == rank_before + 1:
                    independent_null_basis_list.append(vec)

            ##  Odd codes需要正交归一化
            if self.even_or_odd == 'even':
                basis_list = independent_null_basis_list
            else:
                basis_list = orthogonalize(independent_null_basis_list)
            self._logical_operators_x = [self._OpType.HermitianOperatorFromOccupy(np.where(basis_list[i] == 1)[0], []) for i in
                                         range(len(basis_list))]

        ##  返回逻辑算子
        return self._logical_operators_x


    ##  SECTION: ----获取Z-type logical operators----
    @property
    def logical_operators_z(self) -> list[PauliOperator|MajoranaOperator]:
        """
        output：np.array[MajoranaOperator]，independent \\hat gamma logical operators
        """

        ##  计算逻辑算子
        if self._logical_operators_x is None:

            ##  从零空间筛选与行空间线性无关的基矢
            codewords = self.check_matrix_z.null_space()  # 零空间
            independent_null_basis_list = []  # 初始化线性无关基矢组
            matrix = self.check_matrix_z  # 行空间和已有基矢合成的基矢组
            for vec in codewords:
                rank_before = np.linalg.matrix_rank(matrix)
                matrix = np.vstack([matrix, vec])
                if np.linalg.matrix_rank(matrix) == rank_before + 1:
                    independent_null_basis_list.append(vec)

            ##  Odd codes需要正交归一化
            if self.even_or_odd == 'even':
                basis_list = independent_null_basis_list
            else:
                basis_list = orthogonalize(independent_null_basis_list)
            self._logical_operators_z = [self._OpType.HermitianOperatorFromOccupy([], np.where(basis_list[i] == 1)[0]) for i in
                                         range(len(basis_list))]

        ##  返回逻辑算子
        return self._logical_operators_x


    ##  SECTION: ----获取X-type stabilizer generators的rank----
    @property
    def rank_x(self) -> int:
        """
        output：H_X的秩
        """

        ##  返回rank_x
        return np.linalg.matrix_rank(self.check_matrix_x)


    ##  SECTION: ----Get rank of Z generators----
    @property
    def rank_z(self) -> int:
        """
        output：H_Z的秩
        """

        ##  返回rank_z
        return np.linalg.matrix_rank(self.check_matrix_z)


    ##  SECTION: ----获取X-type stabilizer generators的矩阵表示----
    @property
    def check_matrix_x(self) -> galois.GF2:
        """
        output：H_X的矩阵表示
        """

        ##  从stabilizer generators计算矩阵
        matrix = Operator.get_matrix(self.generators_x, self.physical_number)
        matrix = matrix[:, 0::2]

        ##  返回H_X的矩阵表示
        return matrix


    ##  SECTION: ----获取Z-type stabilizer generators的矩阵表示----
    @property
    def check_matrix_z(self) -> galois.GF2:
        """
        output：H_Z的矩阵表示
        """

        ##  从stabilizer generators计算矩阵
        matrix = Operator.get_matrix(self.generators_z, self.physical_number)
        matrix = matrix[:, 1::2]

        ##  返回H_Z的矩阵表示
        return matrix


    ##  SECTION: ----计算码距----
    @property
    def distance(self):
        """
        output：int, get code distance by mip method
        """

        ##  返回码距较小值
        return min(self.distance_x, self.distance_z)


    ##  SECTION: ----获取X-type stabilizer generators的码距----
    @property
    def distance_x(self) -> int:
        """
        output：Z-type stabilizer generators的码距
        """

        ##  如果没有计算需要先计算
        if self._distance_x is None:
            self._distance_x = mip_distance_caculator(self.check_matrix_x, Operator.get_matrix(self.logical_operators_x,self.physical_number))

        ##  返回结果
        return self._distance_x


    ##  SECTION: ----获取Z-type stabilizer generators的码距----
    @property
    def distance_z(self) -> int:
        """
        output：Z-type stabilizer generators的码距
        """

        ##  如果没有计算需要先计算
        if self._distance_z is None:
            self._distance_z = mip_distance_caculator(self.check_matrix_z, Operator.get_matrix(self.logical_operators_z,self.physical_number))

        ##  返回结果
        return self._distance_z


    ##  SECTION: ----从校验矩阵构造量子码----
    @classmethod
    def FromCheckMatrix(cls, check_matrix_x, check_matrix_z=None):
        """
        input.check_matrix_x：$m$为Majorana stabilizer generators的数量，这里设定的是一个self-dual的Majorana stabilizer code
        input.check_matrix_z：$m$为Majorana stabilizer generators的数量，这里设定的是一个self-dual的Majorana stabilizer code
        output：MajoranaCSSCode实例
        """

        ##  检查输入
        try:
            if check_matrix_z is None:
                check_matrix_x = galois.GF2(np.array(check_matrix_x, dtype=int))
                check_matrix_z = galois.GF2(np.array(check_matrix_x, dtype=int))
                assert np.all(check_matrix_x @ check_matrix_x.T == 0), "check matrix必须是self-dual matrix"
            else:
                check_matrix_x = galois.GF2(np.array(check_matrix_x, dtype=int))
                check_matrix_z = galois.GF2(np.array(check_matrix_z, dtype=int))
                assert np.all(check_matrix_x @ check_matrix_z.T == 0), "check matrix必须对易"
        except Exception as e:
            raise ValueError("check matrix不合法") from e

        ##  从check matrix中提取stabilizer generators
        generators_x = []
        generators_z = []
        for i in range(len(check_matrix_x)):
            generators_x.append(cls._OpType.HermitianOperatorFromOccupy(np.where(check_matrix_x[i] == 1)[0], []))
        for i in range(len(check_matrix_z)):
            generators_z.append(cls._OpType.HermitianOperatorFromOccupy([], np.where(check_matrix_z[i] == 1)[0]))

        ##  构造code
        return cls(generators_x, generators_z, check_matrix_x.shape[1])


class MajoranaCode(QuantumCode):

    _OpType=MajoranaOperator

    ##  SECTION: ----构造函数----
    def __init__(self,generators:Iterable[MajoranaOperator],physical_number:int):
        """
        input.generators：Stabilizer generators
        input.physical_number：Fermionic sites数目
        output：None
        """

        ##  检查输入
        try:
            generators = [temp.copy() for temp in generators]
            physical_number = int(physical_number)
            assert np.all([isinstance(temp, MajoranaOperator) for temp in generators]), "generators必须都是Majorana算符"
            assert np.all([temp.max_index<physical_number for temp in generators]), "generators的最大索引必须小于physical_number"
        except Exception as e:
            raise ValueError("输入参数有误") from e

        ##  直接初始化父类
        super().__init__(generators, physical_number)


class MajoranaCSSCode(MajoranaCode, QuantumCSSCode):
    _OpType = MajoranaOperator

    ##  SECTION: ----Constructor----
    def __init__(self, generators_x, generators_z, physical_number):
        """
        input.generators_x：list of MajoranaOperator, X Majorana stabilizer
        input.generators_z：list of MajoranaOperator, Z Majorana stabilizer
        input.physical_number：fermionic site number
        """

        ##  检查输入
        try:
            np.all([isinstance(generator, MajoranaOperator) for generator in generators_x])
            np.all([isinstance(generator, MajoranaOperator) for generator in generators_z])
            generators_x: list[MajoranaOperator] = [generator for generator in generators_x]
            generators_z: list[MajoranaOperator] = [generator for generator in generators_z]
        except ValueError as e:
            raise ValueError("generators不合法") from e

        try:
            physical_number = int(physical_number)
            assert physical_number > 0, "physical_number must be a positive integer."
        except Exception as e:
            raise ValueError("physical_number must be a positive integer.") from e

        ##  Initialize code
        QuantumCSSCode.__init__(self, generators_x, generators_z, physical_number)


    ##  SECTION: ----计算\\hat gamma^\\prime code distance----
    @property
    def distance_z(self)->int:
        """
        output：Z-type stabilizer generators的码距
        """

        self._distance_z = self.distance_x
        return self._distance_z


    ##  SECTION: ----计算\\hat gamma^\\prime logical operators----
    @property
    def logical_operators_z(self) -> list[PauliOperator|MajoranaOperator]:
        """
        output：$\\hat\\gamma^\\prime$ type的独立逻辑算子
        """

        ##  直接从$\\hat\\gamma$-type logical operators计算$\\hat\\gamma^\\prime$-type logical operators
        self._logical_operators_z = [temp.dual for temp in self._logical_operators_x]
        return self._logical_operators_z


class PauliCode(QuantumCode):

    _OpType=PauliOperator

    ##  SECTION: ----构造函数----
    def __init__(self, generators: Iterable[PauliOperator], physical_number: int):
        """
        input.generators：Stabilizer generators
        input.physical_number：Qubits数目
        output：None
        """

        ##  检查输入
        try:
            generators = [temp.copy() for temp in generators]
            physical_number = int(physical_number)
            assert np.all([isinstance(temp, PauliOperator) for temp in generators]), "generators必须都是Pauli算符"
            assert np.all([temp.max_index < physical_number for temp in generators]), "generators的最大索引必须小于physical_number"
        except Exception as e:
            raise ValueError("输入参数有误") from e

        ##  直接初始化父类
        super().__init__(generators, physical_number)


##  CHAPTER: ====Pauli CSS codes====
class PauliCSSCode(PauliCode, QuantumCSSCode):

    _OpType = PauliOperator

    ##  SECTION: ----构造函数----
    def __init__(self, generators_x: Iterable[PauliOperator], generators_z: Iterable[PauliOperator], physical_number: int) -> None:
        """
        input.generators_x：list of PauliOperator, X stabilizers
        input.generators_z：list of PauliOperator, Z stabilizers
        input.physical_number：number of physical qubits
        """

        ##  检查输入
        try:
            np.all([isinstance(generator, PauliOperator) for generator in generators_x]), '稳定子都要是Pauli算子'
            np.all([isinstance(generator, PauliOperator) for generator in generators_z]), '稳定子都要是Pauli算子'
            generators_x: list[PauliOperator] = [generator for generator in generators_x]
            generators_z: list[PauliOperator] = [generator for generator in generators_z]
        except ValueError as e:
            raise ValueError("generators不合法") from e

        try:
            physical_number = int(physical_number)
            assert physical_number > 0, "Physical number必须是正整数."
        except Exception as e:
            raise ValueError("Physical number不合法") from e

        ##  直接初始化父类
        QuantumCSSCode.__init__(self, generators_x, generators_z, physical_number)
