From da28b3ad82b63bfdfe462f4666f4530a13b06fa0 Mon Sep 17 00:00:00 2001
From: Nicolas Hiillos <nicolashiillos@gmail.com>
Date: Thu, 2 Mar 2023 18:10:54 +0200
Subject: [PATCH] Add lite6 path plotting script

---
 scripts/plot_lite6_csv/plot_lite6_csv.py      | 82 +++++++++++++++++++
 scripts/plot_lite6_csv/shell.nix              | 11 +++
 src/lite6_controller/src/lite6_controller.cpp |  4 +-
 3 files changed, 95 insertions(+), 2 deletions(-)
 create mode 100644 scripts/plot_lite6_csv/plot_lite6_csv.py
 create mode 100644 scripts/plot_lite6_csv/shell.nix

diff --git a/scripts/plot_lite6_csv/plot_lite6_csv.py b/scripts/plot_lite6_csv/plot_lite6_csv.py
new file mode 100644
index 0000000..6fb79f4
--- /dev/null
+++ b/scripts/plot_lite6_csv/plot_lite6_csv.py
@@ -0,0 +1,82 @@
+# 3D Heatmap in Python using matplotlib
+  
+# to make plot interactive 
+#%matplotlib
+  
+# importing required libraries
+from mpl_toolkits.mplot3d import Axes3D
+import matplotlib.pyplot as plt
+import numpy as np
+from pylab import *
+
+import sys
+  
+print('arg:', sys.argv)
+if len(sys.argv) <= 1:
+    print('Give file path as arg')
+    exit()
+
+filepath = sys.argv[1]
+f = open(filepath, "r")
+data = f.read()
+
+# creating a dataset
+x_succ = []
+y_succ = []
+z_succ = []
+
+x_fail = []
+y_fail = []
+z_fail = []
+
+for l in data.split("\n"):
+    d = l.split(",")
+    if len(d) == 0:
+        continue
+    s = d[0]
+    for p in d[1:]:
+        p = p.split(" ")
+        if len(p) < 3:
+            continue
+        if s == 'success':
+           x_succ.append(float(p[0])) 
+           y_succ.append(float(p[1]))
+           z_succ.append(float(p[2]))
+        else:
+           x_fail.append(float(p[0]))
+           y_fail.append(float(p[1]))
+           z_fail.append(float(p[2]))
+
+x_succ = np.array(x_succ)
+y_succ = np.array(y_succ)
+z_succ = np.array(z_succ)
+
+x_fail = np.array(x_fail)
+y_fail = np.array(y_fail)
+z_fail = np.array(z_fail)
+  
+# creating figures
+fig = plt.figure(figsize=(10, 10))
+ax = fig.add_subplot(111, projection='3d')
+
+# setting color bar
+color_map1 = cm.ScalarMappable(cmap=cm.Greens_r)
+color_map1.set_array([x_succ + y_succ + z_succ])
+  
+color_map2 = cm.ScalarMappable(cmap=cm.Reds_r)
+color_map2.set_array([x_fail + y_fail + z_fail])
+
+# creating the heatmap
+img1 = ax.scatter(x_succ, y_succ, z_succ, marker='s',
+                 s=2, color='green')
+  
+img2 = ax.scatter(x_fail, y_fail, z_fail, marker='s',
+                 s=0.1, color='red')
+# adding title and labels
+ax.set_title("3D Heatmap")
+ax.set_xlabel('X-axis')
+ax.set_ylabel('Y-axis')
+ax.set_zlabel('Z-axis')
+  
+# displaying plot
+plt.show()
diff --git a/scripts/plot_lite6_csv/shell.nix b/scripts/plot_lite6_csv/shell.nix
new file mode 100644
index 0000000..1e7f4f2
--- /dev/null
+++ b/scripts/plot_lite6_csv/shell.nix
@@ -0,0 +1,11 @@
+# shell.nix
+{ pkgs ? import <nixpkgs> {} }:
+let
+  my-python-packages = p: with p; [
+    matplotlib
+    numpy
+    # other python packages
+  ];
+  my-python = pkgs.python3.withPackages my-python-packages;
+in my-python.env
+
diff --git a/src/lite6_controller/src/lite6_controller.cpp b/src/lite6_controller/src/lite6_controller.cpp
index 30cd64d..e27d96e 100644
--- a/src/lite6_controller/src/lite6_controller.cpp
+++ b/src/lite6_controller/src/lite6_controller.cpp
@@ -384,7 +384,7 @@ public:
       move_group.execute(ts[0]);
 
       status = status + "," + pointsToString(&goal->motion.path,0,0,0);
-      //appendLineToFile("OUTPUT.csv", status);
+      appendLineToFile("OUTPUT.csv", status);
 
       result->result = "success";
       goal_handle->succeed(result);
@@ -394,7 +394,7 @@ public:
 
     status = "failure";
     status = status + "," + pointsToString(&goal->motion.path,0,0,0);
-    //appendLineToFile("OUTPUT.csv", status);
+    appendLineToFile("OUTPUT.csv", status);
 
     RCLCPP_ERROR(this->get_logger(), "Planner failed to return trajectory in time");
     result->result = "failure";