from numpy import concatenate
from reed_solomon import encode as rs_encode, decode as rs_decode, BITS_PER_BLOCK, MESSAGE_BITS_PER_BLOCK
from viterbi import encode as v_encode, decode as v_decode, trellis
from utils import blockify


def encode(bits):
    rs_encoded_bits = rs_encode(bits)
    rs_encoded_blocks = blockify(rs_encoded_bits, BITS_PER_BLOCK)
    v_encoded_blocks = []
    for block in rs_encoded_blocks:
        v_encoded_blocks.append(v_encode(block))
    return concatenate(v_encoded_blocks)


def decode(bits):
    v_encoded_block_num_bits = BITS_PER_BLOCK * trellis.n // trellis.k
    v_encoded_blocks = blockify(bits, v_encoded_block_num_bits)
    rs_encoded_blocks = []
    for block in v_encoded_blocks:
        rs_encoded_blocks.append(v_decode(block))
    original_blocks = []
    for block in rs_encoded_blocks:
        original_blocks.append(rs_decode(block))
    return concatenate(original_blocks)