from Go2Py.robot.interface import GO2Real
import time
import numpy as np

from rosbags.rosbag2 import Writer
from rosbags.typesys import Stores, get_typestore
from rosbags.typesys.stores.ros2_foxy import (
    builtin_interfaces__msg__Time as Time,
    sensor_msgs__msg__Imu as Imu,
    std_msgs__msg__Header as Header,
    geometry_msgs__msg__Vector3 as Vector3,
    geometry_msgs__msg__Quaternion as Quaternion,
)


def main():
    robot = GO2Real(mode='lowlevel')
    time.sleep(1)
    soc = robot.getBatteryState()
    imu = robot.getIMU()
    accel = imu['accel']
    gyro = imu['gyro']
    """Iterate over IMAGES and save to output bag."""
    typestore = get_typestore(Stores.ROS2_FOXY)
    with Writer('imu_noise_identification_dataset') as writer:
        conn = writer.add_connection('/go2/imu', Imu.__msgtype__)

        while (soc > 10):
            soc = robot.getBatteryState()
            imu = robot.getIMU()
            a = imu['accel']
            g = imu['gyro']
            accel = Vector3(a[0], a[1], a[2])
            gyro = Vector3(g[0], g[1], g[2])
            q = Quaternion(0, 0, 0, 1)
            timestamp = time.time() * 10**9
            msg = Imu(
                Header(
                    stamp=Time(sec=int(timestamp // 10**9), nanosec=int(timestamp % 10**9)),
                    frame_id='base_link'),
                linear_acceleration=accel,
                angular_velocity=gyro,
                orientation=q,
                orientation_covariance=np.zeros(9),
                linear_acceleration_covariance=np.zeros(9),
                angular_velocity_covariance=np.zeros(9)
            )
            writer.write(conn, timestamp, typestore.serialize_cdr(msg, msg.__msgtype__))
            time.sleep(0.02)
            print(soc)
    robot.close()
    exit()


if __name__ == '__main__':
    main()