import numpy as np
from PIL import Image
import random

def get_G():
    img = Image.open("output2g.png").convert("L")
    P = (np.array(img) < 128).astype(np.uint8)
    return P

def test_error_trapping(P, num_data_err, num_parity_err):
    N_data = 16384
    N_parity = 2048
    
    data_errs = random.sample(range(N_data), num_data_err)
    parity_errs = random.sample(range(N_parity), num_parity_err)
    
    synd = np.zeros(N_parity, dtype=np.uint8)
    for i in data_errs:
        synd ^= P[i]
    for i in parity_errs:
        synd[i] ^= 1
        
    print(f"Generated {num_data_err} data errs and {num_parity_err} parity errs.")
    print(f"Initial Syndrome Weight: {np.sum(synd)}")
    
    # Try to find the data errors by checking weight(S ^ P_i)
    # This only works efficiently if num_data_err == 1
    if num_data_err == 1:
        for i in range(N_data):
            test_w = np.sum(synd ^ P[i])
            if test_w <= 20: 
                print(f"  Found data error {i} with residual weight {test_w} (matches parity errs)")
    elif num_data_err == 2:
        print("  Skipping search for t=2 data errors (O(N^2))")

if __name__ == "__main__":
    P = get_G()
    test_error_trapping(P, 1, 4)
