import numpy as np
import sys

ALPHABET = "abcdefghijklmnopqrstuvwxyz" \
           "ABCDEFGHIJKLMNOPQRSTUVWXYZ" \
           "0123456789 ."

CHAR_TO_IDX = {c: i for i, c in enumerate(ALPHABET)}

G = 10.0
ENERGY_LIMIT = 2000.0
TEXT_LENGTH = 40
ALPHABET_LENGTH = len(ALPHABET)

assert ALPHABET_LENGTH == 64

def pair_to_index(a: str, b: str) -> int:
    i1 = CHAR_TO_IDX[a]
    i2 = CHAR_TO_IDX[b]
    return (i1 << 6) + i2

def hadamard(r: int) -> np.ndarray:
    if r == 0:
        return np.array([[1.]], dtype=np.float32)
    H = hadamard(r - 1)
    return np.block([[H, H], [H, -H]])

def Br(r: int) -> np.ndarray:
    H = hadamard(r)
    return np.vstack([H, -H])

def make_codebook(r: int = 11, num_blocks: int = TEXT_LENGTH // 2, Eb: float = ENERGY_LIMIT):
    B = Br(r)
    C = np.hstack((B, B)).astype(np.float32)
    n = C.shape[1]
    alpha = Eb / (num_blocks * 2 * n)
    C *= np.sqrt(alpha, dtype=C.dtype)
    return C, alpha

def encode_message(msg: str, C: np.ndarray) -> np.ndarray:
    if len(msg) != TEXT_LENGTH:
        raise ValueError("Message must be exactly 40 characters.")

    idxs = [CHAR_TO_IDX[c] for c in msg]
    pair_idxs = [(idxs[i] << 6) | idxs[i+1] for i in range(0, TEXT_LENGTH, 2)]

    signal = np.repeat(C[pair_idxs].ravel(), 2).astype(np.float32)

    if not np.isclose(signal.dot(signal), ENERGY_LIMIT, atol=1e-3):
        raise RuntimeError("energy check failed")

    return signal

if __name__ == "__main__":
    if len(sys.argv) > 1:
        msg = sys.argv[1]
    else:
        msg = input(f"Enter a message ({TEXT_LENGTH} characters): ").strip()
    if len(msg) != TEXT_LENGTH:
        print(f"Message must be exactly {TEXT_LENGTH} characters.")
        sys.exit(1)
    C, _ = make_codebook()
    signal = encode_message(msg, C)
    np.savetxt("input.txt", signal)
    print("Signal written to input.txt")