import socket
import struct
import math
import time

class ModbusRobot:
    def __init__(self, ip="192.168.18.201", port=502):
        self.ip = ip
        self.port = port

    def send_modbus_tcp(self, register, value):
        transaction_id = 0x01
        protocol_id = 0x00
        length = 0x06
        unit_id = 0x01
        function_code = 0x06
        reg_address = register
        reg_value = value
        
        packet = struct.pack('>HHHBBHH', 
                             transaction_id, 
                             protocol_id, 
                             length, 
                             unit_id, 
                             function_code, 
                             reg_address, 
                             reg_value)
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(2.0)
                s.connect((self.ip, self.port))
                print(f"Sending: {packet.hex(' ')}")
                s.sendall(packet)
                response = s.recv(1024)
                return response
        except Exception as e:
            print(f"Error setting: {e}")

    def send_modbus_tcp_32bit(self, start_register, value):
        transaction_id = 0x01
        protocol_id = 0x00
        unit_id = 0x01
        
        # Преобразование в целое число, если передано float
        val_int = int(value)
        low_word = (val_int >> 16) & 0xFFFF
        high_word = val_int & 0xFFFF
        
        function_code = 0x10
        quantity_of_registers = 2
        byte_count = 4
        length = 11

        packet = struct.pack('>HHHBBHHBHH', 
                             transaction_id, 
                             protocol_id, 
                             length, 
                             unit_id, 
                             function_code, 
                             start_register * 2, 
                             quantity_of_registers, 
                             byte_count,
                             high_word, 
                             low_word)
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(2.0)
                s.connect((self.ip, self.port))
                print(f"Send 32-bit ({val_int}): {packet.hex(' ')}")
                s.sendall(packet)
                response = s.recv(1024)
                print(f"Answer: {response.hex(' ')}")
                return response
        except Exception as e:
            print(f"Error setting: {e}")

    def point_XYZ_ABC(self, X, Y, Z, A, B, C): # Перемещение в точку по координатам и ориентации 
        self.send_modbus_tcp_32bit(start_register=0x08, value=0)
        self.send_modbus_tcp_32bit(start_register=0x02, value=X*1000) 
        self.send_modbus_tcp_32bit(start_register=0x03, value=Y*1000) 
        self.send_modbus_tcp_32bit(start_register=0x04, value=Z*1000) 
        self.send_modbus_tcp_32bit(start_register=0x05, value=A*1000) 
        self.send_modbus_tcp_32bit(start_register=0x06, value=B*1000) 
        self.send_modbus_tcp_32bit(start_register=0x07, value=C*1000)

    def read_modbus_tcp_registers(self, register, count):
        transaction_id = 0x02
        protocol_id = 0x00
        length = 0x06
        unit_id = 0x01
        function_code = 0x04
        
        packet = struct.pack('>HHHBBHH', 
                             transaction_id, 
                             protocol_id, 
                             length, 
                             unit_id, 
                             function_code, 
                             register, 
                             count)
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(2.0)
                s.connect((self.ip, self.port))
                
                print(f"send request (FC04): {packet.hex(' ')}")
                s.sendall(packet)
                response = s.recv(1024)
                
                if len(response) < 9:
                    print("Error: short message")
                    return None
                
                byte_count = response[8]
                register_data = response[9:]
                num_registers = byte_count // 2
                values = struct.unpack('>' + 'H' * num_registers, register_data)
                
                print(f"answer success: {response.hex(' ')}")
                print(f"Values (from {register}): {values}")
                return values
        except Exception as e:
            print(f"Error while reading: {e}")

    def _get_circle_points_in_plane(self, X0, Y0, Z0, R, vector_u, vector_v):
        num_segments = 100
        points = []
        mag_u = math.sqrt(sum(i**2 for i in vector_u))
        mag_v = math.sqrt(sum(i**2 for i in vector_v))
        u = [i/mag_u for i in vector_u]
        v = [i/mag_v for i in vector_v]
        
        for K in range(num_segments + 1):
            angle = (2 * math.pi * K) / num_segments
            cos_a = math.cos(angle)
            sin_a = math.sin(angle)
            
            xk = X0 + R * cos_a * u[0] + R * sin_a * v[0]
            yk = Y0 + R * cos_a * u[1] + R * sin_a * v[1]
            zk = Z0 + R * cos_a * u[2] + R * sin_a * v[2]
            points.append((xk, yk, zk))
        return points

    def run_circle_motion(self, X0, Y0, Z0, R, delay_ms, vector_u, vector_v, repeat_count=0):
        points = self._get_circle_points_in_plane(X0, Y0, Z0, R, vector_u, vector_v)
        iteration = 0
        while True:
            if repeat_count > 0 and iteration >= repeat_count:
                break
            for x, y, z in points:
                self.send_modbus_tcp_32bit(0x02, int(x * 1000))
                self.send_modbus_tcp_32bit(0x03, int(y * 1000))
                self.send_modbus_tcp_32bit(0x04, int(z * 1000))
                self.send_modbus_tcp_32bit(0x08, 0)
                time.sleep(delay_ms / 1000.0)
            iteration += 1
    
    def visualize_robot_model(self, j1=None, j2=None, j3=None, j4=None, j5=None, j6=None):
        import matplotlib.pyplot as plt
        from matplotlib.widgets import Slider, TextBox
        import numpy as np
        H1, OFF_Y, L2, L3_H, L3_V, D6 = 531.5, 185.0, 760.0, 857.5, 220.0, 125.0
        LIMITS = [(-170, 170), (-157, 100), (-76, 98), (-180, 180), (-125, 125), (-720, 720)]

        def get_rot_x(q):
            r = np.radians(q)
            return np.array([[1, 0, 0], [0, np.cos(r), -np.sin(r)], [0, np.sin(r), np.cos(r)]])
        def get_rot_y(q):
            r = np.radians(q)
            return np.array([[np.cos(r), 0, np.sin(r)], [0, 1, 0], [-np.sin(r), 0, np.cos(r)]])
        def get_rot_z(q):
            r = np.radians(q)
            return np.array([[np.cos(r), -np.sin(r), 0], [np.sin(r), np.cos(r), 0], [0, 0, 1]])

        def get_positions(q):
            R1 = get_rot_z(-q[0]); p0 = np.array([0, 0, 0]); p1 = R1 @ np.array([OFF_Y, 0, H1])
            R2 = R1 @ get_rot_y(-q[1]); p2 = p1 + R2 @ np.array([0, 0, L2]); R3 = R2 @ get_rot_y(-q[2])
            v_l3 = np.array([L3_H, 0, L3_V]); p3 = p2 + R3 @ v_l3; alpha = np.arctan2(L3_V, L3_H)
            R_align = get_rot_y(-np.degrees(alpha)); R4_local = R_align @ get_rot_x(q[3]) @ R_align.T
            R5 = R3 @ R4_local @ get_rot_y(-q[4]); p4 = p3 + R5 @ np.array([D6, 0, 0])
            return [p0, p1, p2, p3, p4]

        fig = plt.figure(figsize=(10, 7))
        ax = fig.add_subplot(111, projection='3d')
        ax.view_init(elev=20, azim=-90)

        def draw_robot(q_vals):
            ax.clear()
            pts = get_positions(q_vals)
            x, y, z = zip(*pts)
            ax.plot(x, y, z, 'k-', linewidth=4, solid_capstyle='round')
            p2, p3 = pts[2], pts[3]
            p_end = pts[4]
            ax.scatter(*p_end, color='red', s=30)
            ax.set_xlim([-1500, 1500]); ax.set_ylim([1500, -1500]); ax.set_zlim([0, 2000])
            ax.set_title(f"TCP: X={p_end[1]:.1f}, Y={p_end[0]:.1f}, Z={p_end[2]:.1f}")
            fig.canvas.draw_idle()

        # Если углы переданы -рисуем один раз и выходим
        if all(v is not None for v in [j1, j2, j3, j4, j5, j6]):
            draw_robot([j1, j2, j3, j4, j5, j6])
            plt.show(block=True)
        else:
            # Иначе - интерактивный режим со слайдерами
            plt.subplots_adjust(bottom=0.35)
            sliders = [Slider(plt.axes([0.15, 0.28-(i*0.04), 0.55, 0.02]), f'J{i+1}', 
                       LIMITS[i][0], LIMITS[i][1], valinit=0) for i in range(6)]
            [s.on_changed(lambda val: draw_robot([s.val for s in sliders])) for s in sliders]
            draw_robot([0]*6)
            plt.show()