import odrive
from odrive.enums import *
import time
import socket

class Odrive:
    def __init__(self) -> None:
        self.drive = odrive.find_any()
        time.sleep(1)
        self.init_odrive()
        self.set_zero()
        self.set_closed_loop()

    def init_odrive(self) -> None:
        self.drive.axis0.controller.config.input_mode = 1
        self.drive.axis1.controller.config.input_mode = 1
        self.drive.axis0.motor.config.current_lim = 70
        self.drive.axis1.motor.config.current_lim = 70
        self.drive.axis0.motor.config.current_lim_margin = 30
        self.drive.axis1.motor.config.current_lim_margin = 30

    def set_zero(self) -> None:
        self.drive.axis0.encoder.set_linear_count(0)
        self.drive.axis1.encoder.set_linear_count(0)

    def set_idle(self) -> None:
        self.drive.axis0.requested_state = AXIS_STATE_IDLE
        self.drive.axis1.requested_state = AXIS_STATE_IDLE

    def set_closed_loop(self) -> None:
        self.drive.axis0.requested_state = AXIS_STATE_CLOSED_LOOP_CONTROL
        self.drive.axis1.requested_state = AXIS_STATE_CLOSED_LOOP_CONTROL
        time.sleep(1)

    def goto_raw(self, mot_0, mot_1) -> None:
        self.drive.axis0.controller.input_pos = mot_0
        self.drive.axis1.controller.input_pos = mot_1

class Server:
    def __init__(self) -> None:
        self.armed = False
        self.x_offset = 0  #TODO: figure out what these are...
        self.x_multiplier = 1  #TODO: figure out what these are...
        self.y_offset = 0  #TODO: figure out what these are...
        self.y_multiplier = 1  #TODO: figure out what these are...
        self.sock = socket.socket(socket.AF_INET,socket.SOCK_DGRAM)
        self.udp_host = socket.gethostname()
        self.udp_port = 59200
        self.sock.bind((self.udp_host, self.udp_port))
        self.init_odrive()
        self.control_loop()

    def init_odrive(self):
        self.drive = Odrive()

    def control_loop(self):
        while True:
            data, _ = self.sock.recvfrom(1024)
            print(f"Received Message: {data}, type: {type(data)}")
            if self.armed:
                angle_0, angle_1 = self.calculate_angles(data[0], data[1])
                self.drive.goto_raw(angle_0, angle_1)

    def calculate_angles(self, target_pos_x, target_pos_y):
        x = (target_pos_x + self.x_offset) * self.x_multiplier
        y = (target_pos_y + self.y_offset) * self.y_multiplier
        pi = 3.142
        pulley_diameter = 10.0  # cm
        distance_per_revolution = pi * pulley_diameter

        target_angle_0 = (y / distance_per_revolution) - (x / distance_per_revolution)
        target_angle_1 = (y / distance_per_revolution) + (x / distance_per_revolution)

        return target_angle_0, target_angle_1