import numpy as np
from PIL import Image

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

def test_mld(P, num_errors, threshold):
    N_data = 16384
    N_parity = 2048
    
    # Random errors
    error_idx = np.random.choice(N_data, num_errors, replace=False)
    
    # Syndrome s = sum(P[i])
    s = np.zeros(N_parity, dtype=np.uint8)
    for i in error_idx:
        s ^= P[i]
    
    # MLD check: flip if count(s & P[i]) > threshold
    found = []
    for i in range(N_data):
        score = np.sum(np.logical_and(s, P[i]))
        # Weight of P[i] is approx 1024.
        # If score > threshold, flag it.
        if score > threshold:
            found.append(i)
            
    correct = set(found).intersection(set(error_idx))
    fp = set(found).difference(set(error_idx))
    return len(correct), len(fp)

def main():
    P = get_matrix()
    weights = [np.sum(P[i]) for i in range(16384)]
    avg_w = np.mean(weights)
    print(f"Average signature weight: {avg_w:.2f}")
    
    for t in [1, 5, 20]:
        print(f"\nEvaluating MLD for {t} errors:")
        # Try thresholds around avg_w/2
        for th in [avg_w * 0.45, avg_w * 0.5, avg_w * 0.55, avg_w * 0.6]:
            c, fp = test_mld(P, t, th)
            print(f"  Th: {th:4.0f} | Found: {c}/{t} | FP: {fp}")

if __name__ == "__main__":
    main()
