#!/usr/bin/python3
# Turn a P3-35 slowly in current mode
# Dependencies:
#   python3 -m pip install canopen
#
# Usage:
#   ./turnp3.py <can_device> <can_id> <phase_current_amps> <rad_per_sec>

# Python Standard Library
import sys
import platform
import time      
import threading
import signal

# pypi modules
import canopen

# Declare a CAN mutex
mutex = threading.Lock()

def do_every_thread(period,f,*args):
    def gen_tick(): # Generator for driftless ticks
        t = time.time()
        while True:
            t += period
            yield max(t - time.time(),0)
    g = gen_tick()
    while True:
        time.sleep(next(g)) # Generate the next sleep duration, then sleep
        f(*args) # Call the registered functions with its arguments

def do_every(period,f,*args):
    # Spin off a new thread to execute the given function periodically
    thd = threading.Thread(target=do_every_thread, args=(period,f,*args))
    thd.daemon = True
    thd.start()

def set_angle(node, start, inc, max):
    # theta in encoder cts, between 0 and cts_per_elec_cyc
    try: set_angle.theta = (set_angle.theta + inc) % max
    except: set_angle.theta = start

    # theta_e needs to be scaled to int16_t and centered on zero (+/- 32k) before being written
    if(set_angle.theta >= max / 2):
        angle = int(((set_angle.theta - max) << 16) / max)
    else:
        angle = int((set_angle.theta << 16) / max)

    # Write to commanded_theta_e
    mutex.acquire()
    try: node.sdo.download(0x2262,0,angle.to_bytes(length=2,byteorder='little',signed=True))
    except: print("angle={}".format(angle)) # Unable to convert angle
    finally: mutex.release()

def show_data(node):
    try: show_data.t = time.monotonic() - show_data.start
    except: 
        show_data.start = time.monotonic()
        show_data.t = 0.0

    mutex.acquire()
    try:
        temp = int.from_bytes(node.sdo.upload(0x2202,0), byteorder='little', signed=False)
        therm = int.from_bytes(node.sdo.upload(0x2209,0), byteorder='little', signed=False)
        #theta = int.from_bytes(node.sdo.upload(0x2262,0), byteorder='little', signed=True)
        encoder_raw = int.from_bytes(node.sdo.upload(0x2264,0), byteorder='little', signed=False)
        encoder_pos = int.from_bytes(node.sdo.upload(0x2240,0), byteorder='little', signed=True)
    finally: mutex.release()
    print("{:.2f}, {}, {}, {}, {}".format(show_data.t, temp, therm, encoder_raw, encoder_pos))

def quit(signo, _frame):
    print("Interrupted by signal(%d), shutting down" % signo)
    wait_for_ctrl_c.evt.set() # Interrupts the wait

def wait_for_ctrl_c():
    # Sleep until interrupted by Ctrl-C
    system = platform.system()      # Determine the operating system
    if system == "Windows": # Ctrl-C generates a KeyboardInterrupt exception
        while True:
            try:
                time.sleep(1)
            except KeyboardInterrupt:
                pass  # This is not a failure
                break # Just exit the while loop now
    elif system == "Linux": # Ctrl-C generates a SIGINT signal
        # Set up a new interruptable event
        wait_for_ctrl_c.evt = threading.Event()
        # Catch these process signals
        for sig in ('TERM', 'INT'):
            signal.signal(getattr(signal, 'SIG'+sig), quit)

        wait_for_ctrl_c.evt.wait()

def turn(can_device, can_id, phase_current, rad_per_sec):
    network = canopen.Network() # Start with an empty network
    system = platform.system()  # Determine the operating system
    print("Connecting to {0} in {1}...".format(can_device, system))
    try: 
        if system == "Windows":
            network.connect(bustype='pcan', 
                            channel='PCAN_USBBUS'+str(int(can_device[-1:])+1), 
                            bitrate=1000000)
        elif system == "Linux":
            network.connect(bustype='socketcan', 
                            channel=can_device, 
                            bitrate=1000000)
    except:
        return "Failed to connect."
  
    print("Connection succeeded, adding CANopen node id {}...".format(can_id))
    node = network.add_node(can_id, None)
  
    # Reboot the node, start fresh
    node.nmt.send_command(0x81)
    time.sleep(1)
    
    # Read encoder resolution 2383,23
    cts_per_rev = int.from_bytes(node.sdo.upload(0x2383,23), byteorder='little', signed=False)
    # Read pole pairs 2383,2
    pole_pairs = int.from_bytes(node.sdo.upload(0x2383,2), byteorder='little', signed=False)

    cts_per_elec_cycle = round(cts_per_rev / pole_pairs)

    # Show commutation / control / feedback configuration
    revCommutation = int.from_bytes(node.sdo.upload(0x2383,22), byteorder='little', signed=False)
    revFeedback = int.from_bytes(node.sdo.upload(0x2383,33), byteorder='little', signed=False)
    revControl = int.from_bytes(node.sdo.upload(0x2383,3), byteorder='little', signed=False)
    print("Reversal settings: commutation[{}], feedback[{}], control[{}]".format(revCommutation, revFeedback, revControl))
    
    # Calculate a torque command (-1000 to +1000 of max_motor_torque) from the given phase_current (A)
    # Read torque constant (mNm/A) 2383,12
    kt = int.from_bytes(node.sdo.upload(0x2383,12), byteorder='little', signed=False)
    # Read max_motor_torque 2383,13 (peak rated), max_app_torque (2383,26)
    max_torque = int.from_bytes(node.sdo.upload(0x2383,26), byteorder='little', signed=False)
    trq_cmd = int(phase_current * kt * 1000 / max_torque)

    # Set overwrite_theta_e 2280,3
    node.sdo.download(0x2280,3,b'\x01')

    # Set encoder_pos 2240,0
    node.sdo.download(0x2240,0,b"\x00\x00\x00\x00")

    # Set torque mode 6060,0 = 10
    node.sdo.download(0x6060,0,b'\x0A')

    # Set the torque 6071,0
    print("Torque cmd: {}".format(trq_cmd))
    node.sdo.download(0x6071,0,trq_cmd.to_bytes(length=2,byteorder='little'))

    # Calculate the optimal update rate and increment amount
    # cts/rev / 6.28 rad/rev * rad/s * s/cycle = cts/cycle
    # Search periods from 0.0100 and 0.0200 (from 100 Hz down to 50 Hz)
    #   for the period and increment that yields the smallest error
    best_error = 1.0
    for i in range(100,200):
        cts_flt = cts_per_rev / 6.283185 * rad_per_sec * i / 10000.0
        error = abs(round(cts_flt) - cts_flt)
        if(error < best_error):
            best_error = error
            best_period = i / 10000.0
            best_inc = round(cts_flt)

    # Loops:
    #   Write new theta_e
    print("Incrementing {} encoder cts at {:.1f} Hz to achieve {:.2f} rad/s.".format(best_inc, 1/best_period, best_inc / best_period * 6.283185 / cts_per_rev))
    do_every(best_period, set_angle, node, 0, best_inc, cts_per_elec_cycle)

    print("\nPress Ctrl-C to Exit\n")

    #   Periodically show time, temp, and therm
    print("Time(s), PuckTemp(C), MotorTherm(C), EncoderRaw(cts), EncoderPos(cts)")
    do_every(0.5, show_data, node)

    wait_for_ctrl_c()

    # Set Mode to Idle (0)
    print("Setting mode to IDLE...")
    node.sdo.download(0x6060,0,b'\x00')
    
    # Disconnect from CAN
    print("Disconnecting CAN network...")
    network.disconnect()

    return "Done."

if __name__ == "__main__":
    # Read the command-line arguments
    can_device = sys.argv[1]  # Typically 'can0' for both Linux & Windows
    can_id = int(sys.argv[2]) # 1-127
    phase_current = float(sys.argv[3]) # Phase current in Amps
    rad_per_sec = float(sys.argv[4]) # Rotational speed

    if(rad_per_sec < 0.2):
        print("rad_per_sec must be >= 0.2")
        exit(0)

    # Run
    result = turn(can_device, can_id, phase_current, rad_per_sec)

    print(result)
