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


class BicycleCode:
    def __init__(self, majorana_number:int, weight:int, logical_number:int, seed:int)->None:
        ##  SECTION：----数据预处理-----
        majorana_number=int(majorana_number)
        weight=int(weight)
        logical_number=int(logical_number)
        seed=int(seed)

        ##  SECTION：----生成校验矩阵-----
        np.random.seed(seed)
        H, diff_set = generate_matrix(majorana_number, weight, majorana_number-logical_number)
        self.seed=seed
        self.weight=weight
        self.H=H


#%%  CHAPTER：====构造Majorana bicycle code====
class MajoranaBicycleCode(MajoranaCSSCode,BicycleCode):
    """""
    self.majorana_number：int对象，fermionic sites的数目（偶数）
    self.weight：int对象，稳定子算符权重 (偶数)
    self.logical_number：int对象，逻辑位数目 (偶数)
    self.seed：int对象，随机种子
    """""

    __slots__ = ["majorana_number", "weight", "logical_number", "seed"]

    ##  SECTION：----构造方法-----
    def __init__(self, majorana_number:int, weight:int, logical_number:int, seed:int)->None:
        """""
        input.majorana_number：fermionic sites的数目（偶数）
        input.weight：稳定子算符权重 (偶数)
        input.logical_number：逻辑位数目 (偶数)
        input.seed：随机种子
        """""
        BicycleCode.__init__(self,majorana_number, weight, logical_number, seed)
        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,majorana_number)

    def __str__(self):
        return f"MajoranaBicycleCode({self.majorana_number},{self.weight},{self.majorana_number-self.logical_number},{self.seed})"

class PauliBicycleCode(PauliCSSCode,BicycleCode):
    """""
    self.pauli_number：int对象，qubits的数目（偶数）
    self.weight：int对象，稳定子算符权重 (偶数)
    self.logical_number：int对象，逻辑位数目 (偶数)
    self.seed：int对象，随机种子
    """""

    __slots__ = ["pauli_number", "weight", "logical_number", "seed"]

    #%%  CHAPTER：====构造方法====
    def __init__(self, pauli_number:int, weight:int, logical_number:int, seed:int)->None:
        """""
        input.pauli_number：qubits的数目（偶数）
        input.weight：稳定子算符权重 (偶数)
        input.logical_number：逻辑位数目 (偶数)
        input.seed：随机种子
        """""

        BicycleCode.__init__(self,pauli_number, weight, logical_number, seed)
        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,pauli_number)


#%%   CHAPTER：====创建01循环矩阵=====
def create_circulant_matrix(difference_set, size):
    """""
    input.difference_set：差集（独特差值）
    input.size：矩阵大小 n
    output：numpy.ndarray (n,n) 01循环矩阵
    """""
    matrix = np.zeros((size, size), dtype=int)
    for i in range(size):
        for d in difference_set:
            j = (i + d) % size
            matrix[i, j] = 1
    return matrix


#%%   CHAPTER：====生成满足唯一差值特性的01循环矩阵====
def generate_matrix(physical_number, row_weight, row_number):
    """""
    input.N：总参数（偶数）
    input.k：行权目标（偶数）
    input.M：保留行数 (< N/2)
    output：(GF(2)矩阵, 差集)
    """""

    # 验证参数
    if physical_number % 2 != 0:
        raise ValueError("physical_number必须是偶数")
    if row_weight % 2 != 0:
        raise ValueError("row_weight必须是偶数")
    if row_number >= physical_number / 2:
        raise ValueError("row_number必须小于physical_number/2")

    n = physical_number // 2  # 循环矩阵大小
    k_half = row_weight // 2  # C的行权重

    # 步骤1: 生成满足唯一差值特性的差集
    difference_set = generate_difference_set(n, k_half)

    # 步骤2: 创建循环矩阵C和其转置
    C = create_circulant_matrix(difference_set, n)
    C_T = C.T

    # 步骤3: 构造H0 = [C, C_T]
    H0 = np.hstack((C, C_T))

    # 步骤4: 删除行以实现均匀列权重
    rows_to_keep = list(range(H0.shape[0]))

    # 计算初始列权重
    col_weights = np.sum(H0, axis=0)

    # 计算需要删除的行数
    rows_to_delete = H0.shape[0] - row_number

    # 贪心算法删除行，使列权重均匀
    for _ in range(rows_to_delete):
        best_row = -1
        best_variance = float('inf')

        # 尝试删除每一行，找到使列权重方差最小的行
        for row in rows_to_keep:
            # 临时删除该行
            temp_weights = col_weights - H0[row, :]
            variance = np.var(temp_weights)

            if variance < best_variance:
                best_variance = variance
                best_row = row

        # 删除最佳行
        rows_to_keep.remove(best_row)
        col_weights -= H0[best_row, :]

    # 创建最终矩阵
    H = H0[rows_to_keep, :]

    # 转换为GF(2)矩阵
    GF2 = galois.GF2
    H_gf2 = GF2(H)

    return H_gf2, difference_set


#%% CHAPTER：====生成满足唯一差值特性的差集====
def generate_difference_set(n, k):
    """""生成满足唯一差值特性的差集
    input.n：环大小
    input.k：差集目标大小
    output：排序后的差集列表
    """""

    # 初始化差集
    diff_set = [0]

    # 记录已使用的差值
    used_differences = set()

    while len(diff_set) < k:
        candidate = np.random.randint(1, n - 1)
        valid = True

        # 检查与现有元素的所有差值是否唯一
        for d in diff_set:
            diff1 = (candidate - d) % n
            diff2 = (d - candidate) % n

            if diff1 in used_differences or diff2 in used_differences:
                valid = False
                break

        if valid:
            # 添加新元素并记录差值
            diff_set.append(candidate)
            for d in diff_set[:-1]:
                diff1 = (candidate - d) % n
                diff2 = (d - candidate) % n
                used_differences.add(diff1)
                used_differences.add(diff2)

    return sorted(diff_set)