import galois
import numpy as np


##  CHAPTER: ====求解线性方程====
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: ----检查输入----
    assert matrix.shape[1] == vector.shape[0]
    A = matrix.T
    b = vector.reshape(-1, 1)

    ##  PART: ----求解线性方程----
    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: ----返回解----
    return solution