"""
模块: SurfaceCode
给定边长d，生成对应的surface code
------------------------------------------------------------------------------------------------
SurfaceCode属性：
    self.d：code distance，也就是正方晶格边长
    self.H_X：X stabilizer check matrix
    self.H_Z：Z stabilizer check matrix

    PauliSurfaceCode(PauliCSSCode,SurfaceCode)：
        None
------------------------------------------------------------------------------------------------
SurfaceCode方法：
    __init__（魔法方法）：构造physical sites和check matrices

    PauliSurfaceCode方法：
        __init__（魔法方法，重构自SurfaceCode和PauliCSSCode）：构造physical sites和check matrices
"""
import numpy as np
from galois import GF2
from clifforder.Code.QuantumCode import PauliCSSCode
from clifforder.Code.Operator import PauliOperator


##  CHAPTER: ====构造surface code的check matrices====
class SurfaceCode:
    __slots__ = ['d', 'H_X', 'H_Z']

    ##  SECTION: ----构造函数----
    def __init__(self, d: int) -> None:
        """
        input.d：一个奇数，代表了surface code的边长
        output：None
        """

        ##  输入检查
        try:
            self.d = int(d)
        except Exception as e:
            raise ValueError('输入的边长不合法') from e

        ##  初始参数
        num_qubits = d * d
        stabilizers_x = []
        stabilizers_z = []

        ##  生成Check Matrices
        for i in range(-1, d):
            for j in range(-1, d):
                #  Neighbors
                neighbors = []
                #  Check 4 corners around (i+0.5, j+0.5)
                candidates = [(i, j), (i + 1, j), (i, j + 1), (i + 1, j + 1)]

                for cx, cy in candidates:
                    if 0 <= cx < d and 0 <= cy < d:
                        neighbors.append(cx * d + cy)

                weight = len(neighbors)
                if weight == 0:
                    continue

                ##  Parity check
                if (i + j) % 2 == 0:
                    if weight == 4:
                        stabilizers_z.append(neighbors)
                    elif weight == 2:
                        if j == -1 or j == d - 1:
                            stabilizers_z.append(neighbors)
                else:
                    if weight == 4:
                        stabilizers_x.append(neighbors)
                    elif weight == 2:
                        if i == -1 or i == d - 1:
                            stabilizers_x.append(neighbors)

        ##  生成check matrix
        self.H_X = GF2(np.zeros((len(stabilizers_x), num_qubits), dtype=int))
        for k, stab in enumerate(stabilizers_x):
            self.H_X[k][stab] = 1

        self.H_Z = GF2(np.zeros((len(stabilizers_z), num_qubits), dtype=int))
        for k, stab in enumerate(stabilizers_z):
            self.H_Z[k][stab] = 1

    ##  TODO：----绘制surface code的图形----
    def plot(self):
        pass


##  CHAPTER: ====Pauli surface codes====
class PauliSurfaceCode(PauliCSSCode, SurfaceCode):

    ##  SECTION: ----构造函数----
    def __init__(self, d: int):
        
        ##  构造Geometry和Check Matrices
        SurfaceCode.__init__(self, d)

        ##  根据校验矩阵生成稳定子生成元
        generators_x = []
        generators_z = []

        for i in range(len(self.H_X)):
            generators_x.append(PauliOperator.HermitianOperatorFromOccupy(np.where(self.H_X[i] != 0)[0], []))

        for i in range(len(self.H_Z)):
            generators_z.append(PauliOperator.HermitianOperatorFromOccupy([], np.where(self.H_Z[i] != 0)[0]))

        #  Note: Bypass strict Euclidean orthogonality check in QuantumCode
        self.generators_x = generators_x
        self.generators_z = 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

        self.generators = generators_x + generators_z
        self.physical_number = self.H_X.shape[1]
        self.checker_number = len(self.generators)
        self._logical_operators = None
        self._distance = None

        if not np.all(self.H_X @ self.H_Z.T == 0):
             raise ValueError("generators必须是对易的")

        ##  先验给出逻辑算符
        self._logical_operators_x = [PauliOperator.HermitianOperatorFromOccupy([x * d for x in range(d)], [])]
        self._logical_operators_z = [PauliOperator.HermitianOperatorFromOccupy([], [y for y in range(d)])]

    ##  SECTION: ----获取X-type logical operators----
    @property
    def logical_operators_x(self):
        return self._logical_operators_x

    ##  SECTION: ----获取Z-type logical operators----
    @property
    def logical_operators_z(self):
        return self._logical_operators_z
