import machine
from machine import Pin, ADC
from neopixel import NeoPixel
from motion_controller.accel_stepper import AccelStepper
from utime import sleep_us, sleep_ms, ticks_ms

"""
CL57T(V4.0) Closed-Loop Stepper Driver
10
Figure 10: Sequence chart of control signals
Remark:
a) t1: ENA must be ahead of DIR by at least 200ms. Usually, ENA+ and ENA- are NC (not connected). See
“Connector P1 Configurations” for more information.
b) t2: DIR must be ahead of PUL effective edge by 2us to ensure correct direction;
c) t3: Pulse width not less than 1us;
d) t4: Low level width not less than 1us;
e) Duty cycle of PUL signal is recommended 50%"""


class MotionControllerBase(AccelStepper):
    REV = 1
    FWD = 0
    def __init__(self, step_pin: str, dir_pin: str, steps_per_revolution: float):
        self.step_pin = Pin(step_pin, Pin.OUT, value=0)
        self.dir_pin = Pin(dir_pin, Pin.OUT, value=0)
        self._last_dir  = self.REV

        self.steps_per_revolution = steps_per_revolution

        super().__init__(self.fwd, self.rev)

    def one_step(self):
        self.step_pin.value(1)
        sleep_us(1)
        self.step_pin.value(0)

    def fwd(self):
        if self._last_dir != self.FWD:
            self.dir_pin.value(1)
            sleep_us(2)
            self._last_dir = self.FWD
        self.one_step()

    def rev(self):
        if self._last_dir != self.REV:
            self.dir_pin.value(0)
            sleep_us(2)
            self._last_dir = self.REV
        self.one_step()


class ThrustController(MotionControllerBase):
    """"""
    HOMING_BACKOFF = 100

    def __init__(self, step_pin: str, dir_pin: str, steps_per_revolution: float, steps_per_mm: float, homing_pin: str):
        self.homing_pin = Pin(homing_pin, Pin.IN, Pin.PULL_UP)
        super().__init__(step_pin, dir_pin, steps_per_revolution)

        self.steps_per_mm = steps_per_mm

        self.set_max_speed(self.mm_to_steps(200))
        self.set_acceleration(self.mm_to_steps(400))
        self.set_speed(0)

    def mm_to_steps(self, mm: float) -> int:
        # 2mm pitch * 17 teeth per rev = 34mm/rev

        return int(mm * self.steps_per_mm)

    def home(self):
        print("Homing")
        self.set_max_speed(self.mm_to_steps(10))
        self.move(self.mm_to_steps(-480))

        while self.homing_pin.value():
            self.run()
        print("Homing complete")

        self.set_current_position(0)
        self.run_to_new_position(self.HOMING_BACKOFF)

        self.set_max_speed(self.mm_to_steps(800))

    def move_to_mm(self, mm: float):
        self.move_to(self.mm_to_steps(mm))

    def set_min_position(self, min_position: float):
        pass

    def set_max_position(self, max_position: float):
        pass

    def set_max_thrust_distance(self, max_distance: float):
        pass

    def set_max_thrust_per_second(self, thrust_per_second: float):
        pass


class RotationController(MotionControllerBase):
    """"""

    def __init__(self, step_pin: str, dir_pin: str, steps_per_revolution: float):
        super().__init__(step_pin, dir_pin, steps_per_revolution)
        self.set_max_speed(self.revolutions_per_second_to_steps_per_second(2.2))
        self.set_acceleration(self.revolutions_per_second_to_steps_per_second(100))

    def revolutions_per_second_to_steps_per_second(self, rps: float) -> int:
        return int(self.steps_per_revolution * rps)

    def set_rps(self, rps: float):
        return self.set_speed(self.revolutions_per_second_to_steps_per_second(rps))


def set_color(neopixel, color):
    for pixel in range(3):
        neopixel[pixel] = color
        neopixel.write()


# To Do
# Fix jitter when manual control
# Increase speed
# Calibrate distances (why are they wrong?)
# Remove and tighten spur gear
# Fix speed for rotation on knob
# Set knob centers
# Hook up 5V power

def machine_main():
    thrust = ThrustController("D10", "D9", 3200, 3200 / (2 * 17), "D21")
    rotation = RotationController("D8", "D7", 320000)

    neopixel_pin = Pin("D6", Pin.OUT)  # Pin number for v1 of the above DevKitC, use pin 38 for v1.1
    np = NeoPixel(neopixel_pin, 3)  # "1" = one RGB led on the "led bus"

    on_pin = Pin("D20", Pin.IN, Pin.PULL_UP)
    off_pin = Pin("D19", Pin.IN, Pin.PULL_UP)

    set_color(np, (255, 0, 0))
    thrust.home()
    set_color(np, (255, 255, 0))

    uint_max = 65535.0
    slider_off_offset = 800.0
    uint_remaining_range = uint_max - slider_off_offset

    MAX_THRUST_MM = 120.0
    MIN_THRUST_MM = 20.0
    THRUST_MM_DELTA = MAX_THRUST_MM - MIN_THRUST_MM
    MAX_TPS = 5
    MAX_TPS_MM_PER_SECOND = THRUST_MM_DELTA * MAX_TPS
    THRUST_MM_PER_SECOND_ACCEL = 2000
    thrust.set_max_speed(thrust.mm_to_steps(5*MAX_TPS_MM_PER_SECOND))
    thrust.set_acceleration(thrust.mm_to_steps(THRUST_MM_PER_SECOND_ACCEL))

    min_limit_pin = ADC(Pin("A1", Pin.IN))
    max_limit_pin = ADC(Pin("A2", Pin.IN))
    thrust_per_second_pin = ADC(Pin("A4", Pin.IN))
    rotation_revolutions_per_second_pin = ADC(Pin("A3", Pin.IN))

    max_rotations_per_second = 2.0
    rotation_deadband = 2000.0
    lower_rotation_enable = (uint_max / 2.0) - rotation_deadband
    upper_rotation_enable = (uint_max / 2.0) + rotation_deadband
    uint_rotation_remaining_range = lower_rotation_enable

    MANUAL_CONTROL = 0
    TPS_CONTROL = 1
    last_state = MANUAL_CONTROL

    EXTENDED = 0
    RETRACTED = 1
    last_move = EXTENDED

    last_thrust_value = 0
    last_speed_percentage = 0

    last_ticks_ms = ticks_ms()
    running = True
    last_running = False
    while True:
        print_debug = (ticks_ms() - last_ticks_ms) > 1000
        if not off_pin.value():
            running = False
        elif not on_pin.value():
            running = True

        min_thrust_val = min_limit_pin.read_u16()
        max_thrust_val = max_limit_pin.read_u16()
        thrust_per_second_val = thrust_per_second_pin.read_u16()
        rotation_revolutions_per_second_value = rotation_revolutions_per_second_pin.read_u16()

        min_limit_percentage = (min_thrust_val - slider_off_offset) / uint_remaining_range
        max_limit_percentage = (max_thrust_val - slider_off_offset) / uint_remaining_range
        if print_debug:
            print(f"{min_thrust_val}:{max_thrust_val}|{thrust_per_second_val}:{rotation_revolutions_per_second_value}")

        if running:
            # Thrust Depth Control
            if min_thrust_val < slider_off_offset:
                if last_state == TPS_CONTROL:
                    last_thrust_value = 0
                    last_state = MANUAL_CONTROL

                if abs(last_thrust_value - max_thrust_val) > 1000:
                    final_mm = (THRUST_MM_DELTA * max_limit_percentage) + MIN_THRUST_MM
                    thrust.move_to_mm(final_mm)
                    last_thrust_value = final_mm
            else:
                slider_set_min_mm = (min_limit_percentage * THRUST_MM_DELTA) + MIN_THRUST_MM
                slider_set_max_mm = (max_limit_percentage * THRUST_MM_DELTA) + MIN_THRUST_MM
                slider_set_max_mm = max(slider_set_max_mm, slider_set_min_mm)

                if last_state == MANUAL_CONTROL:
                    last_move = EXTENDED
                    last_state = TPS_CONTROL
                # if print_debug:
                #     print(f"slider_set_min_mm mm: {slider_set_min_mm}|slider_set_max_mm: {slider_set_max_mm}|last_state: {last_state}")

                if not thrust.is_running():
                    new_endpoint_mm = slider_set_min_mm if last_move == EXTENDED else slider_set_max_mm
                    last_move = RETRACTED if last_move == EXTENDED else EXTENDED
                    if print_debug:
                        print(f"New mm: {new_endpoint_mm}|last_move: {last_move}")
                    thrust.move_to_mm(new_endpoint_mm)

            # Thrust speed control

            # Rotation speed control
            if rotation_revolutions_per_second_value < lower_rotation_enable:
                speed_percentage = (
                                                   lower_rotation_enable - rotation_revolutions_per_second_value
                                           ) / uint_rotation_remaining_range
                speed = speed_percentage * max_rotations_per_second

                if abs(last_speed_percentage - speed) < 0.01:
                    rotation.set_rps(speed)
                    last_speed_percentage = speed_percentage

                if print_debug:
                    print(f"{speed}")

            elif rotation_revolutions_per_second_value > upper_rotation_enable:
                speed_percentage = (
                                                   rotation_revolutions_per_second_value - upper_rotation_enable
                                           ) / uint_rotation_remaining_range
                speed = -speed_percentage * max_rotations_per_second

                if abs(last_speed_percentage - speed) < 0.01:
                    rotation.set_rps(speed)
                    last_speed_percentage = speed_percentage

                if print_debug:
                    print(f"{speed}")

            # if last_rotation_sp
            else:
                if last_speed_percentage != 0:
                    rotation.set_speed(0)
                    last_speed_percentage = 0

            if running != last_running:
                last_thrust_value = 0
                set_color(np, (0, 255, 0))
                last_running = running

        else:
            if running != last_running:
                set_color(np, (255, 255, 0))
                last_running = running
                thrust.move_to_mm(MIN_THRUST_MM)
                rotation.set_speed(0)

        thrust.run()
        rotation.run_speed()

        if print_debug:
            last_ticks_ms = ticks_ms()


machine_main()
