From be73dd8b5700571886d88febe2f5cc8463bcc254 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=AF=9B=E4=BA=9A=E7=90=9B?= <maoyachen55@gmail.com>
Date: Wed, 21 Jun 2023 11:35:54 +0800
Subject: [PATCH] add pybind support

---
 .gitignore         |  6 ++++++
 Makefile           | 22 +++++++++++++++-------
 viterbi.py         | 40 ++++++++++++++++++++++++++++++++++++++++
 viterbi_python.cpp | 13 +++++++++++++
 4 files changed, 74 insertions(+), 7 deletions(-)
 create mode 100644 .gitignore
 create mode 100644 viterbi.py
 create mode 100644 viterbi_python.cpp

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..6d9e35c
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,6 @@
+*.o
+.DS_Store
+.vscode
+*.so
+viterbi_main
+viterbi_test
\ No newline at end of file
diff --git a/Makefile b/Makefile
index 0dc5138..d922cdc 100644
--- a/Makefile
+++ b/Makefile
@@ -5,8 +5,10 @@
 
 CXX = g++
 
-BINS = viterbi_main viterbi_test
-SRCS = viterbi.cpp viterbi_main.cpp viterbi_test.cpp
+BINS     = viterbi_main viterbi_test $(PYLIB)
+SRCS     = viterbi.cpp viterbi_main.cpp viterbi_test.cpp
+PYLIB    = viterbicodec$(shell python3-config --extension-suffix)
+CPPFLAGS = $(shell pybind11-config --includes) -std=c++11 
 
 all: $(BINS)
 
@@ -16,20 +18,26 @@ clean:
 test: viterbi_test
 	./viterbi_test
 
+$(PYLIB): viterbi.cpp viterbi_python.cpp
+	$(CXX) $(CPPFLAGS) -undefined dynamic_lookup -shared -fPIC $^ -o $@
+
 viterbi.o: viterbi.cpp viterbi.h
-	$(CXX) $(CPPFLAGS) $(CXXFLAGS) -c $<
+	$(CXX) $(CPPFLAGS) -c $<
 
 viterbi_main.o: viterbi_main.cpp viterbi.h
-	$(CXX) $(CPPFLAGS) $(CXXFLAGS) -c $<
+	$(CXX) $(CPPFLAGS) -c $<
 
 viterbi_main: viterbi_main.o viterbi.o
-	$(CXX) $(CPPFLAGS) $(CXXFLAGS) $^ -o $@
+	$(CXX) $(CPPFLAGS) $^ -o $@
 
 viterbi_test.o: viterbi_test.cpp viterbi.h
-	$(CXX) $(CPPFLAGS) $(CXXFLAGS) -c $<
+	$(CXX) $(CPPFLAGS) -c $<
 
 viterbi_test: viterbi_test.o viterbi.o
-	$(CXX) $(CPPFLAGS) $(CXXFLAGS) $^ -o $@
+	$(CXX) $(CPPFLAGS) $^ -o $@
+
+%.o: %.cpp
+	$(CXX) $(CPPFLAGS) -c $<
 
 .PHONY: all clean test
 
diff --git a/viterbi.py b/viterbi.py
new file mode 100644
index 0000000..c9d7544
--- /dev/null
+++ b/viterbi.py
@@ -0,0 +1,40 @@
+from viterbicodec import ViterbiCodec, reverse_bits
+
+all = ["Viterbi"]
+
+
+def list2str(bits):
+    return "".join(map(lambda x: "1" if x > 0 else "0", bits))
+
+
+def str2list(bits):
+    return list(map(lambda x: 1 if x == "1" else 0, bits))
+
+
+class Viterbi(ViterbiCodec):
+    def __init__(self, constraint, polynomials):
+        for i in range(len(polynomials)):
+            polynomials[i] = reverse_bits(constraint, polynomials[i])
+        self.constraint = constraint
+        self.polynomials = polynomials
+        ViterbiCodec.__init__(self, constraint, polynomials)
+
+    def encode(self, bits):
+        bits = list2str(bits)
+        output = ViterbiCodec.encode(self, bits)[: -2 * (self.constraint - 1)]
+        return str2list(output)
+
+    def decode(self, bits):
+        bits = list2str(bits)
+        bits += "0" * (2 * (self.constraint - 1))
+        output = ViterbiCodec.decode(self, bits)
+        return str2list(output)
+
+
+if __name__ == "__main__":
+    bits = [1] * 100
+    codec = Viterbi(7, [0o133, 0o171])
+    a = codec.encode(bits)
+    print(a)
+    b = codec.decode(a)
+    print(b)
diff --git a/viterbi_python.cpp b/viterbi_python.cpp
new file mode 100644
index 0000000..abb8c2e
--- /dev/null
+++ b/viterbi_python.cpp
@@ -0,0 +1,13 @@
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "viterbi.h"
+
+PYBIND11_MODULE(viterbicodec, m) {
+  pybind11::class_<ViterbiCodec>(m, "ViterbiCodec")
+      .def(pybind11::init<int, const std::vector<int>&>())
+      .def("encode", &ViterbiCodec::Encode)
+      .def("decode", &ViterbiCodec::Decode);
+
+  m.def("reverse_bits", &ReverseBits);
+}
\ No newline at end of file