"""
Streaming 6DoF from QTM with Integrated MATLAB Dynamics Control
"""

import asyncio
import xml.etree.ElementTree as ET
import math
import qtm_rt
from scipy.spatial.transform import Rotation as R
import numpy as np
import time
import socket

# -------------------- Parameter Initialization --------------------
# Constants from MATLAB code
xt = 0
yt = 0
v = 1.0          # Linear velocity
d = 0.75         # Parameter d
# Uncomment to switch r_i and r_o values as needed
# r_i = 0.75
# r_o = 1.75
# r_i = 0.5
# r_o = 1.3
r_i = 0.5        # Inner radius
r_o = 1.40       # Outer radius

kb = d - r_i     # kb based on MATLAB code
ka = r_o - d     # ka based on MATLAB code

k1 = 3           # Control gain k1
k2 = 4           # Control gain k2


# ODE solver options (not directly used in Python but kept for reference)
# opts = {'RelTol':1e-5, 'AbsTol':1e-7}

# UDP Configuration
UDP_IP_ADDRESS = "192.168.0.105"  # Replace with your server IP address
UDP_PORT_NO = 3004
serverSock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
serverSock.bind((UDP_IP_ADDRESS, UDP_PORT_NO))  # Bind to receive initial packet from robot

# -------------------- Helper Functions --------------------
def create_body_index(xml_string):
    """ Extract a name to index dictionary from 6dof settings xml """
    xml = ET.fromstring(xml_string)

    body_to_index = {}
    for index, body in enumerate(xml.findall("*/Body/Name")):
        body_to_index[body.text.strip()] = index

    return body_to_index

def body_enabled_count(xml_string):
    xml = ET.fromstring(xml_string)
    return sum(enabled.text == "true" for enabled in xml.findall("*/Body/Enabled"))

async def main():

    # Connect to qtm
    connection = await qtm_rt.connect("192.168.0.151")

    # Connection failed?
    if connection is None:
        print("Failed to connect")
        return

    # Take control of qtm, context manager will automatically release control after scope end
    async with qtm_rt.TakeControl(connection, "password"):

        realtime = True

        if realtime:
            # Start new realtime
            await connection.new()
        else:
            print("make real time = true")

    # Get 6dof settings from qtm
    xml_string = await connection.get_parameters(parameters=["6d"])
    body_index = create_body_index(xml_string)

    print("{} of {} 6DoF bodies enabled".format(body_enabled_count(xml_string), len(body_index)))
    global counter
    counter = time.time()
    wanted_body = "Khepera104"


    def on_packet(packet):
        info, bodies = packet.get_6d()
        # print(
        #     "Framenumber: {} - Body count: {}".format(
        #         packet.framenumber, info.body_count
        #     )
        # )

    def calculate_control_term(x, theta):
        """
        Calculate the control term based on the current state.
        This function mirrors the MATLAB 'calculate_control_term'.
        """
        x1, y1 = x
        r = math.sqrt(x1**2 + y1**2)  # Distance from origin

        pk = r - d
        q = 1 if pk > 0 else 0

        # Compute px as defined in MATLAB
        px = math.cos(theta) * (x1 - d * math.sin(theta)) + math.sin(theta) * (y1 + d * math.cos(theta))

        # Sigma function (identity in MATLAB)
        c = v * px

        # Compute denominators, handle potential division by zero
        denominator_ka = ka**2 - (r - d)**2
        denominator_kb = kb**2 - (r - d)**2

        # To prevent division by zero, set a small epsilon
        epsilon = 1e-6
        if denominator_ka == 0:
            denominator_ka = epsilon
        if denominator_kb == 0:
            denominator_kb = epsilon

        # Control term as per MATLAB logic
        term_ka = ((r - d) * (v / r)) / (d * k1 * denominator_ka)
        term_kb = ((r - d) * (v / r)) / (d * k1 * denominator_kb)
        control_term = q * ((1 / d) * (v + term_ka + v * c)) + (1 - q) * ((1 / d) * (v + term_kb + v * c))
        return control_term

    def on_packet(packet):
        """ Callback function to handle incoming packets """
        #nonlocal counter
        info, bodies = packet.get_6d()

        if wanted_body is not None and wanted_body in body_index:
            # Extract one specific body
            wanted_index = 6
            position, rotation= bodies[wanted_index]

            x=position[0]/1000
            y=position[1]/1000
            z=position[2]/1000


            rotation = np.array(rotation)
            rotation = rotation.reshape(3,3)
            r = R.from_matrix(rotation)
            rot = r.as_euler('zyx', degrees=True)

            theta=-rot[0]


            theta_r = math.radians(theta)

            # Target positions (as per MATLAB code)
            x_d = xt
            y_d = yt

            # Control parameters
            V = v  # Linear velocity from MATLAB

            # Time management for control update rate
            current_time = time.time()
            elapsed_time = current_time - counter
            if elapsed_time >= 0.05:  # 50 ms control update rate (20 Hz)
                counter = current_time

                # Calculate control_term using the function
                control_term = calculate_control_term([x, y], theta_r)

                # Compute angular velocity w based on control_term
                # Assuming control_term corresponds to dtheta/dt
                w = control_term  # Modify if scaling is needed

                # Construct and send the command
                # Ensure the command format matches your robot's protocol
                command = f"{w}x{V * 1000}".encode()  # Example format: "w_valuexV_value"

                # Send command to robot
                serverSock.sendto(command, addr)
                print(f"Control Term (w): {w:.4f}, Command Sent: {command}")
        else:
            # Optionally handle other bodies or log information
            for idx, (position, rotation) in enumerate(bodies):
                # Retrieve body name from index
                body_names = [name for name, index in body_index.items() if index == idx]
                body_name = body_names[0] if body_names else f"Body_{idx}"
                print(f"{body_name}: Pos: {position} - Rot: {rotation}")

    # Start streaming frames
    await connection.stream_frames(components=["6d"], on_packet=on_packet)
    counter = time.time()

    # Stream for a specified duration (e.g., 150 seconds)
    await asyncio.sleep(150)

    # Stop streaming
    await connection.stream_frames_stop()

# -------------------- Entry Point --------------------
if __name__ == "__main__":
    # Wait to receive the initial packet from the robot to get 'addr'
    print("Waiting to receive initial message from robot to establish UDP communication address...")
    data, addr = serverSock.recvfrom(1024)
    print(f"Received initial message from robot: {data}, addr: {addr}")

    # Run the asynchronous main function with robot_addr
    try:
        asyncio.get_event_loop().run_until_complete(main(addr))
    except KeyboardInterrupt:
        print("Interrupted by user")

    print("Streaming Stopped")

    # Stop the robot by sending a stop command
    stop_command = "0x0".encode()
    serverSock.sendto(stop_command, addr)
    print("Stop command sent to the robot")
 please resolve below error in above code
Waiting to receive initial message from robot to establish UDP communication address...
Received initial message from robot: b'T336.5739TAX-0.0260AXAY-0.0076AYAZ0.9920AZGX1.0230GXGY-1.3728GYGZ-0.7788GZPL7494PLPR9045PRSL0SLSR0SRUA19224UAUB-18755UBUC-12411UCUD2403UDUE17UEIA166IAIB-1225864080IBIC1ICID-1227762128IDIE166IEIF-1227729408IFIG-1225745520IGIH-1095029932IHII-1095029936IIIJ-1227746880IJIK-1227729408IKIL-1322262249IL', addr: ('192.168.0.104', 40833)
Traceback (most recent call last):
  File "/home/mocaplab/khepera-main/libkhepera-2.1/template/Python_Server/asymmetric_barrier_lyapunov function_with_range_rate.py", line 211, in <module>
    asyncio.get_event_loop().run_until_complete(main(addr))
TypeError: main() takes 0 positional arguments but 1 was given