import numpy as np
from encoder import G, ALPHABET, make_codebook

def decode_blocks_with_state(Y: np.ndarray, C: np.ndarray):
    N = Y.size
    assert N % 2 == 0
    assert N <= 1_000_000

    n_cw = C.shape[1] * 2
    nb = Y.size // n_cw
    assert nb * n_cw == Y.size, "length mismatch between Y and codebook"

    Y_blocks = Y.reshape(nb, n_cw)
    Y_even = Y_blocks[:, ::2].astype(np.float32)
    Y_odd  = Y_blocks[:, 1::2].astype(np.float32)
    C = C.astype(np.float32)

    sqrtg = np.sqrt(G)
    s1 = sqrtg * (Y_even @ C.T) + (Y_odd @ C.T)
    s2 = (Y_even @ C.T) + sqrtg * (Y_odd @ C.T)

    best_if_s1 = np.argmax(s1, axis=1)
    best_if_s2 = np.argmax(s2, axis=1)

    tot1 = np.sum(np.max(s1, axis=1))
    tot2 = np.sum(np.max(s2, axis=1))

    state = 1 if tot1 >= tot2 else 2
    indices = best_if_s1 if state == 1 else best_if_s2

    msg_chars = []
    for idx in indices:
        first  = ALPHABET[idx >> 6]
        second = ALPHABET[idx & 0x3F]
        msg_chars.append(first)
        msg_chars.append(second)

    decoded = ''.join(msg_chars)
    return decoded, state

def decode_blocks(Y: np.ndarray, C: np.ndarray) -> str:
    return decode_blocks_with_state(Y, C)[0]

if __name__ == "__main__":
    Y = np.loadtxt("output.txt")
    C, _ = make_codebook()
    decoded, state = decode_blocks_with_state(Y, C)
    print(f"Decoded message: {decoded}")
    print(f"Detected channel state: {state}")