fw: Tie together all parts for an end-to-end demo

1 files changed, 12 insertions, 2 deletions

diff --git a/controller/fw/tools/reed_solomon.py b/controller/fw/tools/reed_solomon.py
index 9eee6be..c4ca6e4 100644
--- a/controller/fw/tools/reed_solomon.py
+++ b/controller/fw/tools/reed_solomon.py
@@ -64,6 +64,14 @@ def cmdline_func_test(args, print=lambda *args, **kwargs: None, benchmark=False)
             ).decode().replace('0', '.'))
         assert test_data == decoded
 
+def cmdline_func_encode(args, **kwargs):
+    data = np.frombuffer(binascii.unhexlify(args.hex_str), dtype=np.uint8)
+    # Map 8 bit input to 6 bit symbol string
+    data = np.packbits(np.pad(np.unpackbits(data).reshape((-1, 6)), ((0,0),(2, 0))).flatten())
+    encoded = encode(data.tobytes(), nbits=args.bits)
+    print('symbol array:', ', '.join(f'0x{x:02x}' for x in encoded))
+    print('hex string:', binascii.hexlify(encoded).decode())
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     cmd_parser = parser.add_subparsers(required=True)
@@ -75,7 +83,9 @@ if __name__ == '__main__':
     test_parser.add_argument('-s', '--seed', type=int, default=0, help='Random seed')
     test_parser.set_defaults(func=cmdline_func_test)
     enc_parser = cmd_parser.add_parser('encode', help='RS-Encode given hex string')
-    enc_parser.add_argument('hex_str')
+    enc_parser.set_defaults(func=cmdline_func_encode)
+    enc_parser.add_argument('-b', '--bits', type=int, default=8, help='Symbol bit size')
+    enc_parser.add_argument('hex_str', type=str, help='Input data as hex string')
     args = parser.parse_args()
-    args.func(args, print=print)
+    args.func(args)