from ctypes import *
from pyDXL import DXLProtocolV2
class TIMU(Structure):
  _pack_ = 1
  _fields_ = [
     ('acc', c_float * 3), ('gyro', c_float * 3), ('mag', c_float * 3),
    ('lia', c_float * 3), ('rv', c_float * 5), ('grv', c_float * 3),
    ('temp', c_int16) ]

dx = DXLProtocolV2('/dev/ttyUSB0', 1000000)
r = dx.Read(200, 108, sizeof(TIMU))
if r:
  imu = TIMU.from_buffer_copy(r)

#!/usr/bin/python3

# matplotlibによる3D表示
# クォータニオンによる回転処理
# PySide6を使ったGUI

from PySide6.QtCore import Qt, Slot, QObject
from PySide6.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QHBoxLayout, QWidget, QComboBox, QLabel, QPushButton, QMessageBox)

from matplotlib.figure import Figure
from matplotlib.backends.backend_qt5agg import  FigureCanvasQTAgg, NavigationToolbar2QT as NavigationToolbar
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

import re
from ctypes import *
import numpy as np
from serial.tools.list_ports import comports
from scipy.spatial.transform import Rotation

from pyDXL import DXLProtocolV2

# コントロールテーブルのうちIMUの部分
class TIMU(Structure):
  _pack_ = 1
  _fields_ = [
    ('acc', c_float * 3), ('gyro', c_float * 3), ('mag', c_float * 3), ('lia', c_float * 3),
    ('rv', c_float * 5), ('grv', c_float * 3), ('temp', c_int16)
  ]

# PySide6のGUIのクラス
class MainWindow(QMainWindow):
  # コンストラクタ
  def __init__(self, *args, **kwargs):
    super(MainWindow, self).__init__(*args, **kwargs)
    self.dx = None

    # タイトル設定
    self.setWindowTitle('DXMIO')
    # PORT SELECT AREA
    layout_port = self.create_area_port()
    # matplotlibのグラフのクラスオブジェクト
    layout_plot = self.create_area_matplot()
    # GUIレイアウト作成
    layout = QVBoxLayout()
    # PORT AREA
    layout.addLayout(layout_port)
    # MATPLOT AREA
    layout.addWidget(layout_plot)
    # matplotlibのツールバーを作成
    toolbar = NavigationToolbar(layout_plot, self)
    layout.addWidget(toolbar)

    # GUI画面全体のレイアウトを作成
    widget = QWidget()
    widget.setLayout(layout)
    self.setCentralWidget(widget)

    # アニメーション開始
    self.ani = FuncAnimation(self.fig, self.plot, interval=10, blit=False, cache_frame_data=False)

  def sort_key_num(self, port):
    match = re.search(r'\d+', port.device)
    if match:
      return (port.device[:match.start()], int(match.group()))
    return (port.device, 0)

  # PORT選択エリア
  def create_area_port(self):
    # LABEL
    label = QLabel('COM PORT')
    # Port名Combobox
    self.combo = QComboBox()
    self.combo.setFixedWidth(100)
    # 検出したポート名を設定
    ports = sorted(comports(), key=self.sort_key_num)
    for port in ports:
      self.combo.addItem(port.device)
    # Openボタン
    self.button_open = QPushButton('OPEN')
    self.button_open.setFixedWidth(100)
    self.button_open.clicked.connect(self.open)
    # Closeボタン
    self.button_close = QPushButton('CLOSE')
    self.button_close.setFixedWidth(100)
    self.button_close.clicked.connect(self.close)
    self.button_close.setEnabled(False)
    # レイアウトに設定
    layout = QHBoxLayout()
    layout.addWidget(label, 0)
    layout.addWidget(self.combo, 1)
    layout.addWidget(self.button_open, 2)
    layout.addWidget(self.button_close, 3)
    # QHBoxLayoutの左寄せ用
    layout.addStretch()
    return layout

  # MATPLOTエリア
  def create_area_matplot(self):
    self.fig = Figure(figsize=(8, 8), dpi=100) 
    # MATPLOTの軸設定
    self.ax = self.fig.add_subplot(111, projection='3d')
    self.ax.set_aspect('equal')
    self.ax.set_proj_type('persp', focal_length=.4)
    graph_canvas = FigureCanvasQTAgg(self.fig)
    return graph_canvas

  # メッセージBOX表示
  def message_box(self, msg):
    msgbox = QMessageBox(self)
    msgbox.setWindowTitle('ERROR')
    msgbox.setText(msg)
    msgbox.setIcon(QMessageBox.Icon.Critical)
    msgbox.setDefaultButton(QMessageBox.Ok)
    msgbox.exec()

  # OPENボタン Callback
  def open(self):
    try:
      portname = self.combo.currentText()
      port = '\\\\.\\' + portname
      self.dx = DXLProtocolV2(port, 1000000)
      self.combo.setEnabled(False)
      self.button_open.setEnabled(False)
      self.button_close.setEnabled(True)
    except:
      self.dx = None
      self.combo.setEnabled(True)
      self.button_open.setEnabled(True)
      self.button_close.setEnabled(False)
      # エラーメッセージ
      self.message_box(portname + 'がOpenできません')

  # CLOSEボタン Callback
  def close(self):
    del self.dx
    self.dx = None
    self.combo.setEnabled(True)
    self.button_open.setEnabled(True)
    self.button_close.setEnabled(False)

  # 表示更新
  def update_plt(self, acc, grv, quat):
    # ボックスの頂点座標を定義 (中心を原点とする)
    # Vertex index: 0:(-1,-2,-.5), 1:(1,-2,-.5), ...
    vertices = np.array([ [-1, -2, -.2], [1, -2, -.2], [1, 2, -.2], [-1, 2, -.2], [-1, -2, .2], [1, -2, .2], [1, 2, .2], [-1, 2, .2] ]) * 0.5

    # 面を構成する頂点のインデックス
    faces = [
      [0, 1, 2, 3], # Bottom
      [4, 5, 6, 7], # Top
      [0, 1, 5, 4], # Front
      [2, 3, 7, 6], # Back
      [0, 3, 7, 4], # Left
      [1, 2, 6, 5]  # Right
    ]

    # 現在のプロットを消去し軸ラベルと範囲を設定
    self.ax.cla()
    self.ax.set(xlim=(-1, 1), ylim=(-1, 1), zlim=(-1, 1))
    self.ax.update({'xlabel':'x', 'ylabel':'y', 'zlabel':'z'})
    np.set_printoptions(suppress=True, precision=4)

    # ボックスをプロット
    rotated_vertices = Rotation.from_quat(quat).apply(vertices)
    # 回転した頂点を用いて面を形成
    rotated_faces = [rotated_vertices[face] for face in faces]
    self.ax.add_collection3d(Poly3DCollection(rotated_faces, facecolors='cyan', linewidths=1, edgecolors='r', alpha=.4))

    # 加速度を矢印でプロット
    rotated_acc = Rotation.from_quat(quat).apply(acc)
    self.ax.quiver(0, 0, 0, rotated_acc[0], rotated_acc[1], rotated_acc[2], color='red', length=0.1, linewidth=2,arrow_length_ratio=0.1)

    # 重力を矢印でプロット
    rotated_grv = Rotation.from_quat(quat).apply(grv)
    self.ax.quiver(0, 0, 0, rotated_grv[0], rotated_grv[1], rotated_grv[2], color='blue', length=0.1, linewidth=2,arrow_length_ratio=0.1)

    self.ax.set_title(
      'Quaternion: '+', '.join(f'{x:7.3f}' for x in quat)
      +'\nAcceleation: '+', '.join(f'{x:7.3f}' for x in acc)
      +'\nGravitiy: '+', '.join(f'{x:7.3f}' for x in grv)
    )

  # DXMIOからのIMUデータで画面更新
  def plot(self, data):
    if self.dx is not None:
      try:
        r = self.dx.Read(200, 108, sizeof(TIMU))
        if r:
          imu = TIMU.from_buffer_copy(r)
          self.update_plt(np.ctypeslib.as_array(imu.lia), np.ctypeslib.as_array(imu.grv), np.ctypeslib.as_array(imu.rv)[:4])
      except:
        pass

if __name__ == '__main__':
  # Qtアプリケーションの作成
  app = QApplication()

  # フォームを作成して表示
  form = MainWindow()
  form.show()

  # 画面表示のためのループ
  app.exec()

import time
from pyDXL import *
dx = DXLProtocolV2('/dev/ttyAMA0', 57600, 0.05)
ID = 1
dx.Write8(ID, 512, 0)
if input('Clear? (y/N) ') == 'y':
  dx.TxPacket(ID, DXLProtocolV2.INST_CLEAR, (0x01, 0x44, 0x58, 0x4c, 0x22))
  dx.RxPacket(timeout = 5.0)
  time.sleep(1)
dx.Write8(ID, 33, 1)
dx.Write8(ID, 512, 1)
for i in range(-550000, 550000, 550000 >> 4):
  dx.Write32(ID, 528, i)
  time.sleep(0.5)
dx.Write8(ID, 512, 0)