From 63cc948f660b384496d2e06e0587ec975047d615 Mon Sep 17 00:00:00 2001
From: "chaiwit.p" <chaiwit.p@groupmaker.co.th>
Date: Sat, 17 Sep 2022 14:34:17 +0700
Subject: [PATCH] update support pi3hat

---
 moteus_drive/config/params.yaml                |  1 +
 .../moteus_drive/moteus_control/MoteusDrive.py | 18 ++++++++++++++++--
 moteus_drive/moteus_drive/moteus_node.py       |  9 +++++++--
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/moteus_drive/config/params.yaml b/moteus_drive/config/params.yaml
index d6dc437..10fa52c 100644
--- a/moteus_drive/config/params.yaml
+++ b/moteus_drive/config/params.yaml
@@ -1,6 +1,7 @@
 moteus_drive:
   ros__parameters:
     frame_id: "moteus"
+    connector_device: "pi3hat" #fdcanusb or pi3hat
     rezero_on_startup: false
     moteus_ids:
       - 1
diff --git a/moteus_drive/moteus_drive/moteus_control/MoteusDrive.py b/moteus_drive/moteus_drive/moteus_control/MoteusDrive.py
index c4064c6..514fdeb 100644
--- a/moteus_drive/moteus_drive/moteus_control/MoteusDrive.py
+++ b/moteus_drive/moteus_drive/moteus_control/MoteusDrive.py
@@ -2,12 +2,17 @@
 import math
 from time import sleep
 import moteus
+try:
+    import moteus_pi3hat
+except ImportError:
+    pass
 from rclpy.node import Node
 
 class MoteusDrive(Node):
-    def __init__(self, IDs):
+    def __init__(self, Connector, IDs):
         super().__init__('MoteusDriveCycle')
         self.ids = IDs
+        self.Connector = Connector            
         self.conn = []
         self.state = "stop"
         self.terminate = False
@@ -17,9 +22,18 @@ def __init__(self, IDs):
         self.rezero = False
         self.servo_command = None
         self.feedback_position_device = None
+        self.bus_map = {1: [1], 2: [2]}
     
     async def run(self):
-        transport = moteus.Fdcanusb()
+        for idx, device_id in enumerate(self.ids):
+           self.bus_map[idx+1] = [int(device_id)]
+        
+        # create transport init device
+        if self.Connector == "fdcanusb":
+            transport = moteus.Fdcanusb()
+        elif self.Connector == "pi3hat":
+            transport = moteus_pi3hat.Pi3HatRouter(servo_bus_map=self.bus_map)
+            
         for idx, device_id in enumerate(self.ids):
             self.conn.append(moteus.Controller(id = device_id))
         while True:            
diff --git a/moteus_drive/moteus_drive/moteus_node.py b/moteus_drive/moteus_drive/moteus_node.py
index 6a87fbb..8ce553d 100644
--- a/moteus_drive/moteus_drive/moteus_node.py
+++ b/moteus_drive/moteus_drive/moteus_node.py
@@ -15,6 +15,7 @@ def __init__(self):
         
         # get parameters values
         self.frame_id = self.get_parameter("frame_id").value
+        self.connector_device = self.get_parameter("connector_device").value
         self.rezero_on_startup = self.get_parameter("rezero_on_startup").value
         self.devices = self.get_parameter("moteus_ids").value
         
@@ -25,7 +26,10 @@ def __init__(self):
         self.recv_command = {}
         
         # initialize moteus drive cycle
-        self.moteusDrive = MoteusDrive(self.devices)
+        self.moteusDrive = MoteusDrive(
+            self.connector_device, 
+            self.devices
+        )
         
         # initialize start state drive
         self.drive_rezero()
@@ -48,8 +52,9 @@ def __init__(self):
         
     def declare_param(self):
         self.declare_parameter("frame_id", "moteus_drive", ParameterDescriptor(description="Frame ID"))
+        self.declare_parameter("connector_device", "fdcanusb", ParameterDescriptor(description="Conector device"))
         self.declare_parameter("rezero_on_startup", False, ParameterDescriptor(description="Rezero on startup"))
-        self.declare_parameter("moteus_ids",[1], ParameterDescriptor(description="Moteus IDs"))
+        self.declare_parameter("moteus_ids", [1], ParameterDescriptor(description="Moteus IDs"))
 
     def callback_command(self, msg):
         self.time_command = msg.header.stamp