import numpy as np
from PIL import Image

def get_P_block(bx, by):
    img = Image.open("output2g.png").convert("L")
    P = (np.array(img) < 128).astype(np.uint8)
    block = P[by*256 : (by+1)*256, bx*256 : (bx+1)*256]
    return block[0, :] # First row (generator polynomial)

def poly_mul(a, b):
    # Multiplication in GF(2)[x]/(x^256 + 1)
    res = np.zeros(256, dtype=np.uint8)
    for i in range(256):
        if a[i]:
            for j in range(256):
                if b[j]:
                    res[(i+j)%256] ^= 1
    return res

def find_sparse_h2(p_poly):
    # Search for a sparse h2(x) such that h2(x) * p(x) = h1(x) is sparse.
    # We try all weight-2 polynomials (h2 = 1 + x^k)
    best_w = 256
    best_h2 = None
    best_h1 = None
    
    for k in range(1, 256):
        h2 = np.zeros(256, dtype=np.uint8)
        h2[0] = 1
        h2[k] = 1
        h1 = poly_mul(h2, p_poly)
        w = np.sum(h1)
        if w < best_w:
            best_w = w
            best_h2 = h2
            best_h1 = h1
            if w < 20: # Arbitrary threshold for "sparse enough"
                print(f"  Success! h2 = 1 + x^{k} gives h1 weight {w}")
                return h2, h1
                
    # Try weight-3?
    # ...
    return None, None

def main():
    print("Trying to find a sparse H derivation (h2(x) * P(x) = h1(x))...")
    # Check a few blocks
    for bx in range(8):
        p = get_P_block(bx, 0)
        print(f"Analyzing Block (0, {bx}) polynomial (weight {np.sum(p)})...")
        h2, h1 = find_sparse_h2(p)
        if h2 is not None:
            print(f"FOUND SPARSE H for block-column {bx}!")
            break
    else:
        print("No sparse H2 factors of weight 2 found.")

if __name__ == "__main__":
    main()
