"""Hamming(7,4) encoding/correction and BSC capacity. Python 3.10+, stdlib."""
from __future__ import annotations
from itertools import product
import math


G = ((1,0,0,0,1,1,0), (0,1,0,0,1,0,1),
     (0,0,1,0,0,1,1), (0,0,0,1,1,1,1))
H = ((1,1,0,1,1,0,0), (1,0,1,1,0,1,0), (0,1,1,1,0,0,1))


def bits(values, size):
    """Pre: values is finite. Post: validated binary tuple of specified size."""
    values = tuple(values)
    if len(values) != size or any(x not in (0, 1) for x in values):
        raise ValueError(f"expected {size} binary values")
    return values


def encode(message):
    """Pre: four bits. Post: systematic seven-bit Hamming codeword."""
    message = bits(message, 4)
    return tuple(sum(message[i]*G[i][j] for i in range(4)) % 2 for j in range(7))


def correct(received):
    """Pre: seven bits. Post: closest codeword; original guaranteed only for <=1 flip."""
    received = list(bits(received, 7))
    syndrome = tuple(sum(row[i]*received[i] for i in range(7)) % 2 for row in H)
    if syndrome != (0, 0, 0):
        columns = [tuple(H[row][column] for row in range(3)) for column in range(7)]
        received[columns.index(syndrome)] ^= 1
    return tuple(received)


def capacity(p):
    """Pre: BSC crossover p in [0,1]. Post: capacity in bits/channel use."""
    if not math.isfinite(p) or not 0 <= p <= 1:
        raise ValueError("p must lie in [0,1]")
    h = -sum(x*math.log2(x) for x in (p, 1-p) if x)
    return 1-h


def demo():
    sent = encode((0, 1, 1, 0))
    received = list(sent)
    received[2] ^= 1
    print("Sent:", sent)
    print("Received:", tuple(received))
    print("Corrected:", correct(received))
    double = list(sent)
    double[0] ^= 1
    double[1] ^= 1
    print("Two flips recover original:", correct(double) == sent)
    print("BSC(.1) capacity:", round(capacity(.1), 6), "bits/use")
    print("Rate of Hamming(7,4):", 4/7)


def self_test():
    count = 0
    # Invariant: preceding messages/errors recover the original. Bound: remaining 128 cases.
    for message in product((0, 1), repeat=4):
        sent = encode(message)
        check(correct(sent), sent)
        count += 1
        for position in range(7):
            damaged = list(sent)
            damaged[position] ^= 1
            check(correct(damaged), sent)
            count += 1
    check(capacity(0), 1.0)
    check(capacity(1), 1.0)
    check(capacity(.5), 0.0)
    return count + 3

def check(actual, expected):
    """Require an expected result; this check also runs under python -O."""
    if actual != expected:
        raise AssertionError(f"expected {expected!r}, got {actual!r}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        print(f"{self_test()} checks passed")
    else:
        demo()