import galois
import numpy as np
from mip import Model, BINARY, minimize, xsum


##  CHAPTER: ====对逻辑算符基矢正交归一化====
def orthogonalize(matrix):
    """""
    input.matrix：可迭代的GF(2)行向量列表/矩阵
    output：列表[GF(2)行向量]，两两正交的基
    """""
    GF = galois.GF2

    B_i = [GF(v.copy()) for v in matrix]
    ortho_basis = []  # 存储正交基

    while True:

        ##  在B_i中查找odd-weight向量
        length = len(B_i)
        flag = None
        cases = True
        for i in range(length):
            if np.mod(np.count_nonzero(B_i[i]), 2) == 1:
                flag = i
                ortho_basis.append(B_i[i])
                ##  将其余向量对该向量做正交化处理
                for j in range(length):
                    if j != i and np.dot(B_i[j], B_i[i]) == 1:
                        B_i[j] = B_i[j] + B_i[i]
                break
            if i == length - 1:
                cases = False
        if not cases:
            break
        else:
            next_B = B_i[:flag] + B_i[flag + 1:]
            B_i = next_B
            if len(B_i) == 0:
                break

    while True:
        length = len(B_i)
        for i in range(length):
            if np.dot(B_i[i], B_i[0]) == 1:
                temp = ortho_basis[0].copy()
                ortho_basis[0] = ortho_basis[0] + B_i[i] + B_i[0]
                B_i[0] = temp + B_i[0]
                B_i[i] = temp + B_i[i]
                ortho_basis.append(B_i[0])
                ortho_basis.append(B_i[i])
                B_i = B_i[1:i] + B_i[i + 1:]
                break
        if len(B_i) == 0:
            break

    return ortho_basis


##  CHAPTER: ====计算logical operators的最小Hamming weight====
def mip_distance_caculator(H, logicOp) -> int:
    H=np.array(H,dtype=int)
    logicOp=np.array(logicOp,dtype=int)
    n = H.shape[1]  # number of qubits
    m = H.shape[0]  # number of stabilizers
    wstab = np.max([np.sum(H[i, :]) for i in range(m)])  # maximum stabilizer weight
    wlog = np.count_nonzero(logicOp)  # weight of the logical operator
    num_anc_stab = int(np.ceil(np.log2(wstab)))  # how many slack variables are needed to express orthogonality constraints modulo two
    num_anc_logical = int(np.ceil(np.log2(wlog)))  # how many slack variables are needed to express orthogonality constraints modulo two
    num_var = n + m * num_anc_stab + num_anc_logical  # total number of variables

    ##  创建整数规划问题
    model = Model()
    model.verbose = 0
    x = [model.add_var(var_type=BINARY) for _ in range(num_var)]
    model.objective = minimize(xsum(x[i] for i in range(n)))

    ##  与H的正交性约束
    for row in range(m):
        weight = [0] * num_var
        supp = np.nonzero(H[row, :])[0]
        for q in supp:
            weight[q] = 1
        cnt = 1
        for q in range(num_anc_stab):
            weight[n + row * num_anc_stab + q] = -(1 << cnt)
            cnt += 1
        model += xsum(weight[i] * x[i] for i in range(num_var)) == 0

    ##  与逻辑算符的奇数重叠约束
    d_min=H.shape[1]
    for i in range(len(logicOp)):
        supp = np.nonzero(logicOp[i])[0]
        weight = [0] * num_var
        for q in supp:
            weight[q] = 1
        cnt = 1
        for q in range(num_anc_logical):
            weight[n + m * num_anc_stab + q] = -(1 << cnt)
            cnt += 1
        model += xsum(weight[i] * x[i] for i in range(num_var)) == 1
        model.optimize()
        opt_val = int(np.sum([x[i].x for i in range(n)]))
        if opt_val<d_min:
            d_min=opt_val
    return d_min


##  CHAPTER: ====估算logical operators的最小Hamming weight====
def random_distance_caculator_css(H:galois.GF2,logicOp:galois.GF2,number:int):

    ##  输入检查
    try:
        H=galois.GF2(H)
        logicOp=galois.GF2(logicOp)
    except Exception as e:
        raise ValueError("输入的矩阵不合法") from e

    ##  计算零空间
    null_space=H.null_space()
    rows, cols = null_space.shape  # 获取零空间矩阵的维度信息
    dist_bound = cols + 1  # 初始化距离上界为最大可能值（列数+1）
    vec_count = 0  # 计数器：记录找到当前最小权重的向量数量

    ##  随机搜索
    for i in range(number):
        per = np.random.permutation(cols)  # 生成随机排列，用于随机化搜索顺序
        null_space1 = null_space[:, per]  # 对零空间矩阵的列进行随机排列
        null_space2 = null_space1.row_reduce()  # 对排列后的矩阵进行行约简（高斯消元）
        null_space2 = null_space2[:, np.argsort(per)]  # 将列顺序恢复为原始顺序

        for j in range(rows):
            temp_vec = null_space2[j, :]  # 获取当前行向量
            temp_weight = np.count_nonzero(temp_vec)  # 计算向量的Hamming权重（非零元素个数）

            ##  确保结果处于[0,dist_bound]范围内
            if 0 < temp_weight <= dist_bound:

                ##  判断对象是否是逻辑算符
                if np.any(logicOp @ temp_vec):
                    if temp_weight < dist_bound:
                        dist_bound = temp_weight
                        vec_count = 1
                    elif temp_weight == dist_bound:
                        vec_count += 1

            ##  若最小权重为2，则直接返回2
            if dist_bound <= 2:
                return 2

    ##  返回当前上界
    return dist_bound


##  CHAPTER: ====Solve Linear Equations====
def solve(matrix, vector):
    """
    input.matrix：galois.GF(2) 或 01矩阵，形状 (m, n)
    input.vector：galois.GF(2) 或 01向量，形状 (n,)
    output：galois.GF(2) 向量解；若无解返回 None
    """

    # PART: Details
    assert matrix.shape[1] == vector.shape[0]
    A = matrix.T
    b = vector.reshape(-1, 1)

    # PART: Details
    aug = np.concatenate((A, b), axis=1)
    n, m_plus_1 = aug.shape
    m = m_plus_1 - 1
    rank = 0

    for col in range(m):

        pivot_row = None
        for i in range(rank, n):
            if aug[i, col] == 1:
                pivot_row = i
                break
        if pivot_row is None:
            continue  # 该列没有主元，跳过

        if pivot_row != rank:
            aug[[rank, pivot_row], :] = aug[[pivot_row, rank], :]

        for i in range(n):
            if i != rank and aug[i, col] == 1:
                aug[i, :] ^= aug[rank, :]
        rank += 1

    for i in range(rank, n):
        if aug[i, -1] == 1:
            return None  # 无解的情况

    GF=galois.GF(2)
    solution = GF.Zeros(m)
    leading_cols = []

    for i in range(rank):
        for j in range(m):
            if aug[i, j] == 1:
                leading_cols.append(j)
                break

    for i in range(rank):
        col = leading_cols[i]
        solution[col] = aug[i, -1]

        for j in range(col + 1, m):
            if aug[i, j] == 1:
                solution[col] ^= solution[j]

    # PART: Details
    return solution


##  CHAPTER: ====Subspace Intersection====
def cap(matrix1, matrix2):
    """""
    input.matrix1：GF(2)矩阵，行张成子空间1
    input.matrix2：GF(2)矩阵，行张成子空间2
    output：GF(2)矩阵，表示交集的行空间基
    """""

    ##  PART: Details
    m = matrix1.shape[0]

    ## basis1^T * x + basis2^T * y = 0 (), [basis1^T | basis2^T]
    aug_matrix = galois.GF(2)(np.concatenate((matrix1.T, matrix2.T), axis=1))

    nullspace = aug_matrix.null_space()

    ab_space = nullspace[:, :m]

    ## basis1
    if len(ab_space) == 0:
        return galois.GF(2).Zeros((0, matrix1.shape[1]))

    intersection_vectors = ab_space @ matrix1

    rref_intersection = intersection_vectors.row_reduce()

    ##  0
    nz_mask = np.any(rref_intersection != 0, axis=1)
    rref_basis = rref_intersection[nz_mask]

    # PART: Details
    return rref_basis


##  CHAPTER: ====Subspace Difference====
def minus(matrix1, matrix2):
    """""
    input.matrix1：GF(2)矩阵，行张成子空间1
    input.matrix2：GF(2)矩阵，行张成子空间2
    output：列表[ndarray]，为子空间差的基向量（01行向量）
    """""

    # PART: Details
    if matrix1 is None:
        return None
    if matrix2 is None or len(matrix2) == 0:
        return galois.GF(2)(matrix1)
    assert matrix1.shape[1] == matrix2.shape[1]

    intersect = cap(matrix1, matrix2)
    if len(intersect) == 0:
        return galois.GF(2)(matrix1)

    result = []
    for i in range(len(matrix1)):
        rank = np.linalg.matrix_rank(intersect)
        intersect = np.vstack((intersect, matrix1[i]))
        if np.linalg.matrix_rank(intersect) > rank:
            result.append(matrix1[i])

    # PART: Details
    return galois.GF(2)(np.array(result,dtype=int))


##  CHAPTER: ====Subspace Direct Sum====
def direct_sum(matrix1, matrix2):
    """
    input.matrix1：GF(2)矩阵
    input.matrix2：GF(2)矩阵
    output：GF(2)矩阵，直和后的独立基
    """

    # PART: Details
    assert isinstance(matrix1, galois.GF(2))
    assert isinstance(matrix2, galois.GF(2))

    # PART: Details
    result = matrix1[0].copy()

    for i in range(1, len(matrix1)):
        if len(result.shape)>1:
            rank = np.linalg.matrix_rank(result)
        else:
            rank=1
        temp = np.vstack((result, matrix1[i]))
        if np.linalg.matrix_rank(temp) > rank:
            result = temp
    for i in range(len(matrix2)):
        if len(result.shape)>1:
            rank=np.linalg.matrix_rank(result)
        else:
            rank=1
        temp = np.vstack((result, matrix2[i]))
        if np.linalg.matrix_rank(temp) > rank:
            result = temp

    # PART: Details
    return galois.GF(2)(result)