#!/usr/bin/env python3

import numpy as np
import sionna as sn
import tensorflow as tf

# Based on https://nvlabs.github.io/sionna/examples/Sionna_tutorial_part1.html#Forward-Error-Correction-(FEC)

k = 12
n = 20

binary_source = sn.utils.BinarySource()

BATCH_SIZE = 8

u = binary_source([BATCH_SIZE, k])
print("Input bits are: \n", u.numpy())

encoder = sn.fec.ldpc.LDPC5GEncoder(k, n)
c = encoder(u)
print("Encoded bits are: \n", c.numpy())

ber = 0.05
flipped_bits = int(n * ber)
flipped_positions = np.zeros((BATCH_SIZE, flipped_bits), dtype=int)
for i in range(BATCH_SIZE):
  flipped_positions[i] = np.random.permutation(n)[:flipped_bits]
print(f"Applying BER {ber:0.2f} - {flipped_bits} Bits.")
print("Bit error positions (0-indexed):\n", flipped_positions)

e = c.numpy()
for i in range(e.shape[0]):
  for j in flipped_positions[i]:
    e[i][j] = 1.0 if e[i][j] == 0.0 else 0.0
    
e = tf.convert_to_tensor(e, np.float32)

print("Codeword with errors: \n", e.numpy())

decoder = sn.fec.ldpc.LDPC5GDecoder(encoder,
                                    hard_out=True,
                                    return_infobits=True,
                                    num_iter=20,
                                    cn_type="minsum")

w = decoder(e)
print("Decoded bits are: \n", w.numpy())

for i in range(u.shape[0]):
    errors_corrected = 'Success.' if (u[i].numpy() == w[i].numpy()).all() else 'Failed!'
    print(f"Corrected all errors in word {i + 1}: {errors_corrected}")