"""
This script demonstrates the usage of the sick_scan_api to interact with a SICK lidar device,
process point cloud data, and visualize it using matplotlib.
Functions:
    pySickScanCartesianPointCloudMsgToXYZ(pointcloud_msg):
        Converts a SickScanCartesianPointCloudMsg to 3D arrays (x, y, z coordinates).
    pySickScanCartesianPointCloudMsgCallback(api_handle, pointcloud_msg):
        Callback function for processing cartesian point cloud messages.
Usage:
    The script initializes the lidar device, registers a callback for point cloud messages, and enters
    a main loop to plot the point cloud data in real-time.
Note:
    Ensure that the sick_scan_api library is correctly loaded from the specified paths.
"""

import numpy as np
import os
import sys
from sick_scan_api import *
import matplotlib.pyplot as plt

# global settings
class ApiTestSettings:
    def __init__(self):
        self.plot_figure = None
        self.plot_axes = None
        self.plot_points_x = []
        self.plot_points_y = []
        
# Convert a SickScanCartesianPointCloudMsg to 3D arrays
def pySickScanCartesianPointCloudMsgToXYZ(pointcloud_msg):
    # get pointcloud fields
    num_fields = pointcloud_msg.fields.size
    msg_fields_buffer = pointcloud_msg.fields.buffer
    field_offset_x = -1
    field_offset_y = -1
    field_offset_z = -1
    for n in range(num_fields):
        field_name = ctypesCharArrayToString(msg_fields_buffer[n].name)
        field_offset = msg_fields_buffer[n].offset
        if field_name == "x":
            field_offset_x = msg_fields_buffer[n].offset
        elif field_name == "y":
            field_offset_y = msg_fields_buffer[n].offset
        elif field_name == "z":
            field_offset_z = msg_fields_buffer[n].offset
    # Extract x,y,z
    cloud_data_buffer_len = (pointcloud_msg.row_step * pointcloud_msg.height) # length of polar cloud data in byte
    assert(pointcloud_msg.data.size == cloud_data_buffer_len and field_offset_x >= 0 and field_offset_y >= 0 and field_offset_z >= 0)
    cloud_data_buffer = bytearray(cloud_data_buffer_len)
    for n in range(cloud_data_buffer_len):
        cloud_data_buffer[n] = pointcloud_msg.data.buffer[n]
    points_x = np.zeros(pointcloud_msg.width * pointcloud_msg.height, dtype = np.float32)
    points_y = np.zeros(pointcloud_msg.width * pointcloud_msg.height, dtype = np.float32)
    points_z = np.zeros(pointcloud_msg.width * pointcloud_msg.height, dtype = np.float32)
    point_idx = 0
    for row_idx in range(pointcloud_msg.height):
        for col_idx in range(pointcloud_msg.width):
            # Get lidar point in polar coordinates (range, azimuth and elevation)
            pointcloud_offset = row_idx * pointcloud_msg.row_step + col_idx * pointcloud_msg.point_step
            points_x[point_idx] = np.frombuffer(cloud_data_buffer, dtype = np.float32, count = 1, offset = pointcloud_offset + field_offset_x)[0]
            points_y[point_idx] = np.frombuffer(cloud_data_buffer, dtype = np.float32, count = 1, offset = pointcloud_offset + field_offset_y)[0]
            points_z[point_idx] = np.frombuffer(cloud_data_buffer, dtype = np.float32, count = 1, offset = pointcloud_offset + field_offset_z)[0]
            point_idx = point_idx + 1
    return points_x, points_y, points_z

# Callback for cartesian pointcloud messages
def pySickScanCartesianPointCloudMsgCallback(api_handle, pointcloud_msg):
    pointcloud_msg = pointcloud_msg.contents # dereference msg pointer (pointcloud_msg = pointcloud_msg[0])
    print("pySickScanCartesianPointCloudMsgCallback (ROS-{}): api_handle={}, {}x{} pointcloud, {} echo(s), segment {}, {} fields, frame_id \"{}\", topic {}, timestamp {}.{:06d}".format(
        0, api_handle, pointcloud_msg.width, pointcloud_msg.height, pointcloud_msg.num_echos , pointcloud_msg.segment_idx, pointcloud_msg.fields.size, 
        pointcloud_msg.header.frame_id, pointcloud_msg.topic, pointcloud_msg.header.timestamp_sec, pointcloud_msg.header.timestamp_nsec))
    global api_test_settings
    api_test_settings.plot_points_x, api_test_settings.plot_points_y, api_test_settings.plot_points_z = pySickScanCartesianPointCloudMsgToXYZ(pointcloud_msg)

# Python usage example for sick_scan_api
if __name__ == "__main__":
    # Configuration and commandline arguments
    api_test_settings = ApiTestSettings()
    cli_args = " ".join(sys.argv[1:])

    # Load sick_scan_library
    if os.name == 'nt': # Load windows dll
        sick_scan_library = SickScanApiLoadLibrary(["build/Debug/", "build_win64/Debug/", "src/build/Debug/", "src/build_win64/Debug/", "src/sick_scan_xd/build/Debug/", "src/sick_scan_xd/build_win64/Debug/", "./", "../"], "sick_scan_xd_shared_lib.dll")
    else: # Load linux so
        sick_scan_library = SickScanApiLoadLibrary(["build/", "build_linux/", "src/build/", "src/build_linux/", "src/sick_scan_xd/build/", "src/sick_scan_xd/build_linux/", "./", "../"], "libsick_scan_xd_shared_lib.so")
    
    api_handle = SickScanApiCreate(sick_scan_library) # Create a sick_scan instance

    # Initialize lidar by launchfile
    print("sick_scan_xd_api_test.py: initializing lidar, commandline arguments = \"{}\"".format(cli_args))
    SickScanApiInitByLaunchfile(sick_scan_library, api_handle, cli_args)

    # Register a callback for PointCloud messages
    cartesian_pointcloud_callback = SickScanPointCloudMsgCallback(pySickScanCartesianPointCloudMsgCallback)
    SickScanApiRegisterCartesianPointCloudMsg(sick_scan_library, api_handle, cartesian_pointcloud_callback)

    # Create figure and axes
    api_test_settings.plot_figure = plt.figure(figsize=(5, 8))  # Optional: fixed figure size
    api_test_settings.plot_axes = plt.axes()

    # Set fixed ranges for x and y axes
    api_test_settings.plot_axes.set_xlim(0, 5)
    api_test_settings.plot_axes.set_ylim(0, 5)

    # Ensure equal scaling for both axes
    api_test_settings.plot_axes.set_aspect('equal')

    # Run main loop
    while True:
        try:
            print("sick_scan_xd_api_test.py plotting pointcloud")
            plot_points_x = np.copy(api_test_settings.plot_points_x)
            plot_points_y = np.copy(api_test_settings.plot_points_y)
            api_test_settings.plot_axes.clear()
            api_test_settings.plot_axes.scatter(plot_points_x, plot_points_y,  marker='o', linestyle='-', s = 2)
            plt.draw()
            plt.pause(1.0)
        except:
            break

    # Stop lidar, close connection and api handle
    print("sick_scan_xd_api_test.py finishing...")
    SickScanApiDeregisterCartesianPointCloudMsg(sick_scan_library, api_handle, cartesian_pointcloud_callback)
    SickScanApiClose(sick_scan_library, api_handle)
    SickScanApiRelease(sick_scan_library, api_handle)
    SickScanApiUnloadLibrary(sick_scan_library)
    print("sick_scan_xd_api_test.py finished.")