"""""
模块作用：实现构造
"""""
import galois
import numpy as np
from extendedstim.Code.PrimitiveCode.MajoranaCSSCode import MajoranaCSSCode
from extendedstim.Code.PrimitiveCode.PauliCSSCode import PauliCSSCode
from extendedstim.Physics.MajoranaOperator import MajoranaOperator
from extendedstim.Physics.PauliOperator import PauliOperator


# %%  CHAPTER：====构造校验矩阵的类====
class DoubleBicycleCode:

    ##  SECTION：----构造方法----
    def __init__(self,n:int,a_occupy:list,b_occupy:list)->None:
        """""
        input.n：循环矩阵的大小
        input.a_occupy：循环矩阵A中1的位置
        input.b_occupy：循环矩阵B中1的位置
        output：无（构造对象）
        """""

        ##  PART：----数据预处理----
        self.n=n
        self.a_occupy=a_occupy.copy()
        self.b_occupy=b_occupy.copy()
        self.a_occupy.sort()
        self.b_occupy.sort()

        ##  PART：----构造校验矩阵----
        a_list = [1 if i in self.a_occupy else 0 for i in range(n)]
        b_list = [1 if i in self.b_occupy else 0 for i in range(n)]
        H_up = np.hstack([cyclic_matrix(n, a_list), cyclic_matrix(n, b_list), cyclic_matrix(n, a_list).T, cyclic_matrix(n, b_list).T])
        H_down = np.hstack([cyclic_matrix(n, b_list).T, cyclic_matrix(n, a_list).T, cyclic_matrix(n, b_list), cyclic_matrix(n, a_list)])
        H = np.vstack([H_up, H_down])
        self.H=H

    ##  SECTION：----打印代码----
    def __str__(self):
        n = self.n

        S_str_a = 'A='
        for i in self.a_occupy:
            if i == 0:
                S_str_a += f"I_{{{n}}}"
            elif i == 1:
                S_str_a += f"S_{{{n}}}"
            else:
                S_str_a += f"S^{{{i}}}_{{{n}}}"

        S_str_b = 'B='
        for i in self.b_occupy:
            if i == 0:
                S_str_b += f"I_{{{n}}}"
            elif i == 1:
                S_str_b += f"S_{{{n}}}"
            else:
                S_str_b += f"S^{{{i}}}_{{{n}}}"

        return '$' + S_str_a + '$' + '\n' + '$' + S_str_b + '$'


# %%  CHAPTER：====Majorana double bicycle code====
class MajoranaDoubleBicycleCode(MajoranaCSSCode,DoubleBicycleCode):

    ##  SECTION：----构造方法----
    def __init__(self, n: int, a_occupy: list, b_occupy: list) -> None:
        DoubleBicycleCode.__init__(self,n,a_occupy,b_occupy)
        generators_x=[]
        generators_z=[]
        for i in range(len(self.H)):
            generators_x.append(MajoranaOperator.HermitianOperatorFromOccupy(np.where(self.H[i]!=0)[0],[]))
            generators_z.append(MajoranaOperator.HermitianOperatorFromOccupy([],np.where(self.H[i] != 0)[0]))
        MajoranaCSSCode.__init__(self,generators_x,generators_z,self.n*4)


# %%  CHAPTER：====Pauli double bicycle code====
class PauliDoubleBicycleCode(PauliCSSCode,DoubleBicycleCode):

    ##  SECTION：----构造方法----
    def __init__(self, n: int, a_occupy: list, b_occupy: list) -> None:
        DoubleBicycleCode.__init__(self,n,a_occupy,b_occupy)
        generators_x=[]
        generators_z=[]
        for i in range(len(self.H)):
            generators_x.append(PauliOperator.HermitianOperatorFromOccupy(np.where(self.H[i]!=0)[0],[]))
            generators_z.append(PauliOperator.HermitianOperatorFromOccupy([],np.where(self.H[i] != 0)[0]))
        PauliCSSCode.__init__(self,generators_x,generators_z,self.n*4)


# %%  CHAPTER：====用到的静态函数====
##  SECTION：----生成循环移位矩阵----
def shift(number:int, shift:int):
    """""
    input.number：矩阵维度
    input.shift：循环移位步长
    output：GF(2)矩阵，置换矩阵
    """""

    S = np.zeros((number, number), dtype=int)
    for i in range(number):
        S[i, (i + shift) % number] = 1
    return galois.GF(2)(S)


##  SECTION：----创建GF(2)循环矩阵----
def cyclic_matrix(n, a_list):
    """""
    input.n：尺寸
    input.a_list：长度为n的01序列，指示多项式系数
    output：GF(2)循环矩阵 C(n,a)
    """""
    assert len(a_list) == n
    return galois.GF2(np.sum([a*(np.linalg.matrix_power(shift(n,1),i)) for i,a in enumerate(a_list)],axis=0))