#ddrobot_controllers.py
# dew.ninja  Dec 2023
# Controller for differential drive robot

import time
from utime import sleep_ms 
from machine import Pin,Timer, PWM, ADC, DAC, I2C
from neopixel import NeoPixel
import math
import array

# variables that are changed often.
feedback = 1
T = 0.005  # sampling period
cmd_from_VR = True # change to true by 
# ---- code for SSD1306 OLED. Comment out if it is not available 
import ssd1306

# ----- wheel parameters -------
LEFT_WHEEL = 0
RIGHT_WHEEL = 1
WHEEL_RADIUS = 0.08
WHEEL_THICKNESS = 0.04
WHEEL_SEPARATION = 0.45
BASE_LENGTH = 0.65
BASE_WIDTH = 0.55

# range variables
PWMMAX = 1023  # maximum PWM value
PWMMID = 511
PWMMIN = 0
DACMAX = 255
DACMID = 127
DACMIN = 0
ADCMAX = 4095
ADCMID = 2047
ADCMIN = 0
UMAX = 3.3  # limit max controller output to 3.2 vols
UMID = 1.65
UMIN = 0.1 # min controller output

VMAX = 3.3
VMID = 1.65
VMIN = 0.0

WS_LIM = 100.0  # limited wheel speed (rad/s)
w2pwm = PWMMAX/WS_LIM  # convert wheel speed to PWM


# ESP32 Pin assignment 
i2c = I2C(0, scl=Pin(22), sda=Pin(21))

oled_width = 128
oled_height = 64
oled = ssd1306.SSD1306_I2C(oled_width, oled_height, i2c)

def oled_show(oled,s_cmd, s_robot):
    oled.fill(0)
    l1_txt = "L cmd : "+str(round(s_cmd[LEFT_WHEEL],2)) 
    l2_txt = "L robot : "+str(round(s_robot[LEFT_WHEEL],2)) 
    l3_txt = "R cmd : "+str(round(s_cmd[RIGHT_WHEEL],2)) 
    l4_txt = "R robot : "+str(round(s_robot[RIGHT_WHEEL],2)) 
    oled.text(l1_txt,0,0)
    oled.text(l2_txt,0,15)
    oled.text(l3_txt,0,30)
    oled.text(l4_txt,0,45)
    oled.show()

oled_tprev = 0
#oled_showryu(oled,1,2,3)   

# ----------neopixel rings representing motor movement ------------
LW_NeoPixel_Pin = 14
RW_NeoPixel_Pin = 12

NUMPIXELS = 24
LW_pixels = NeoPixel(Pin(LW_NeoPixel_Pin),NUMPIXELS)
RW_pixels = NeoPixel(Pin(RW_NeoPixel_Pin),NUMPIXELS)

rad_per_pixel = 44/(NUMPIXELS*7)

def LW_pixels_display():
    LW_pixels.fill((0,0,0))
    ypixel = (NUMPIXELS - int(wp_robot[LEFT_WHEEL]/rad_per_pixel))%NUMPIXELS # angle of wheel
    if ypixel<0:
        ypixels+= NUMPIXELS
    if ypixel>NUMPIXELS-1:
        ypixel = NUMPIXELS-1
    LW_pixels[ypixel] = (255,0,0) # plant output in red
    LW_pixels.write()

def RW_pixels_display():
    RW_pixels.fill((0,0,0))
    ypixel = (NUMPIXELS - int(wp_robot[RIGHT_WHEEL]/rad_per_pixel))%NUMPIXELS # angle of wheel
    if ypixel<0:
        ypixels+= NUMPIXELS
    if ypixel>NUMPIXELS-1:
        ypixel = NUMPIXELS-1
    RW_pixels[ypixel] = (255,0,0) # plant output in red
    RW_pixels.write()

def update_wheels():  # show both wheel positions
    LW_pixels_display()
    RW_pixels_display()
# ------------------------------------------------

# ----  integrator class -------
class INTEGRATOR:
    def __init__(self, T):
        # states
        self.u0 = 0.0 
        self.u1 = 0.0
        self.y0 = 0.0
        self.y1 = 0.0
        self.T = T
    def reset(self):
        self.u0 = 0.0 
        self.u1 = 0.0
        self.y0 = 0.0
        self.y1 = 0.0   
    def out(self, u):
        self.u1 = self.u0
        self.y1 = self.y0
        self.u0 = u
        y0 = self.y1 + 0.5*self.T*(self.u0+self.u1)
        self.y0 = y0
        return y0
    def sett(self, T):
        self.T = T

# integrator instance for left and right wheels
LW_Int = INTEGRATOR(T)
RW_Int = INTEGRATOR(T)
# -------------- adjust commad using VR ------------------
# --- linear velocity
VCMD_PIN = 34

v_in = ADC(Pin(VCMD_PIN))
v_in.atten(ADC.ATTN_11DB) # range 0 - 3.6 V
v_cmd = 0.0 # linear velocity cmd

WCMD_PIN = 35
w_in = ADC(Pin(WCMD_PIN))
w_in.atten(ADC.ATTN_11DB) # range 0 - 3.6 V
w_cmd = 0.0  # angular velocity cmd

VELMAX = 10   # m/s
WMAX = 10  # rad/s
#adc2vel = VELMAX/ADCMAX

# adc to angular velocity ADCMID->0, ADCMAX->10 rad/s, ADCMIN->-10 rad/s
def adc2w(adcval):  
    return ((2*WMAX)/ADCMAX)*adcval - WMAX

# adc to linear velocity ADCMID->0, ADCMAX->10 m/s, ADCMIN->-10 m/s
def adc2vel(adcval):  
    return ((2*VELMAX)/ADCMAX)*adcval - VELMAX
# ---------------------------------------------
prompt = True  # set prompt off/on so that it won't get in the way.

#
# adc2ws = WS_LIM/ADCMAX

# function to convert adc value to wheel speed
# ADCMAX --> WS_LIM
# ADCMID --> 0
# ADCMIN --> -WS_LIM

def adc2ws(adcval):  
    return ((2*WS_LIM)/ADCMAX)*adcval - WS_LIM


pwm0_out = PWM(Pin(5))  # output to left wheel
pwm0_out.freq(1000) # set PWM frequency to 1 KHz (lower from 5KHz)
pwm1_out = PWM(Pin(4))  # output to right wheel
pwm1_out.freq(1000) # set PWM frequency to 1 KHz (lower from 5KHz)

# dac_out = DAC(Pin(25)) # this doesn't work on Wokwi
dir0_out = Pin(19,Pin.OUT) # direction 0 = CW, 1 = CCW
dir1_out = Pin(18,Pin.OUT) # direction 0 = CW, 1 = CCW
# -----  left wheel velocity. Ch = 0
wr_in = ADC(Pin(39)) # ADC1_CH3. Shown as VN on Wokwi
wr_in.atten(ADC.ATTN_11DB) # range 0 - 3.6 V
# ----- right wheel velocity. Ch = 1 
wl_in = ADC(Pin(32))
wl_in.atten(ADC.ATTN_11DB)

ws_cmd = array.array('f',0.0 for _ in range (2)) # wheel speed commands
ws_robot = array.array('f',0.0 for _ in range (2)) # wheel speed from robot
wp_robot = array.array('f',0.0 for _ in range (2)) # wheel position from robot
# -----------------------------


adc2v = 3.6/ADCMAX  # max of ADC is set to 3.6 v by the package
# v2deg = 100 # convert volt to degree
v2pwm = int(PWMMAX/3.3)
v2dac = int(DACMAX/3.3)


led = Pin(2, Pin.OUT)


T_ms = int(1000*T)
t_data = 0  # time data output


# values for these PID coefficients are computed in PID_update
bi = array.array('f',0.0 for _ in range (2)) 
ad = array.array('f',0.0 for _ in range (2))
bd = array.array('f',0.0 for _ in range (2))
bt = array.array('f',0.0 for _ in range (2))

# controller states
ep0 = array.array('f',0.0 for _ in range (2))  # error for proportional term
e1 = array.array('f',0.0 for _ in range (2))  # true error (for integral term)
e0 = array.array('f',0.0 for _ in range (2))
eus1 = array.array('f',0.0 for _ in range (2)) # error for back calculation term 
eus0 = array.array('f',0.0 for _ in range (2))
ed1 = array.array('f',0.0 for _ in range (2))  # error for derivative term
ed0 = array.array('f',0.0 for _ in range (2))
up0 = array.array('f',0.0 for _ in range (2))  # proportional-term outputs
ui1 = array.array('f',0.0 for _ in range (2))  # integral-term outputs
ui0 = array.array('f',0.0 for _ in range (2))
ud1 = array.array('f',0.0 for _ in range (2))  # derivative-term outputs
ud0 = array.array('f',0.0 for _ in range (2))

u = array.array('f',0.0 for _ in range (2)) # controller output
ulim = array.array('f',0.0 for _ in range (2)) # limited controller output
upos = array.array('f',0.0 for _ in range (2)) # positive controller output
upwm = array.array('i',0 for _ in range (2)) # PWM controller output
udir = array.array('i',0 for _ in range (2)) # DIR controller output

# PID parameters
kp = array.array('f',10.0 for _ in range (2))  # PID gains
ki = array.array('f',0.0 for _ in range (2))
kd = array.array('f',0.0 for _ in range (2))
kt = array.array('f',0.0 for _ in range (2))  # back calculaiton gain
wp = array.array('f',1.0 for _ in range (2))  # proportional weight
wd = array.array('f',1.0 for _ in range (2))  # derivative weight
N = array.array('f',50.0 for _ in range (2))  # derivative filter coefficient




datasize = 200  # number of points in output capture
capture = False  # capture mode 0 = off, 1 = on
capture_flag = False # indicates start/stop of data transfer
data_idx = 0  # data index


timer=Timer(1)



# ------ function to generate pwm/dir outputs from controller output u
def convert2pwmdir(u):
    u_lim = u
    if u > WS_LIM:
        u_lim = WS_LIM         # limit u to WS_LIM
    elif u < -WS_LIM:
        u_lim = -WS_LIM         # limit u to -WS_LIM
    if u_lim >= 0:   # positive u_lim = CW direction
        u_dir = 0
        u_pos = u_lim
    else:
        u_dir = 1
        u_pos = -u_lim
    u_pwm = int(u_pos*w2pwm)
    return u_lim,u_pwm, u_dir

# compute PID coefficients
def PID_update(ch):
    global ad, bd, bi, bt
    bi[ch] = 0.5*T*ki[ch]
    bt[ch] = 0.5*T*kt[ch]
    ad1 = 1+0.5*N[ch]*T
    ad2 = 0.5*N[ch]*T - 1
    ad[ch] = -ad2/ad1
    bd[ch] = kd[ch]*N[ch]/ad1

def PID_update_all():
    PID_update(LEFT_WHEEL)
    PID_update(RIGHT_WHEEL)

# PID controller function
def PID_controller(r,y,ch):
    global e1,e0,ed1,ed0, eus1,eus0, ui1, ui0, ud1, ud0
    
    # state transfer
    e1[ch] = e0[ch]
    ed1[ch] = ed0[ch]
    eus1[ch] = eus0[ch]
    
    ui1[ch] = ui0[ch]
    ud1[ch] = ud0[ch]
    # compute errors for each term
    e0[ch] = r - y
    ep0[ch] = wp[ch]*r - y # weighted proportional error
    ed0[ch] = wd[ch]*r - y # weighted derivative error
    
    up0[ch] = kp[ch]*ep0[ch] # output of P term
    ui0[ch] = ui1[ch] +bi[ch]*(e0[ch]+e1[ch]) + bt[ch]*(eus0[ch]+eus1[ch]) # output of I term
    ud0[ch] = ad[ch]*ud1[ch] +bd[ch]*(ed0[ch] - ed1[ch]) # output of D term
    u0 = up0[ch] + ui0[ch] + ud0[ch]
    #u_lim = limit_and_offset(u)
    return u0

# ------- inverse kinematics function for differential-drive robot------------------
def dd_ikine():
    global v_cmd, w_cmd
    ws_cmd[LEFT_WHEEL] = (1/WHEEL_RADIUS)*(v_cmd+0.5*WHEEL_SEPARATION*w_cmd)
    ws_cmd[RIGHT_WHEEL] = (1/WHEEL_RADIUS)*(v_cmd-0.5*WHEEL_SEPARATION*w_cmd)



# ------------------- timer routine -------------------------------
# ----------------feedback is executed here ---------------
def timer_isr(event):
    global a,b,T,u,ulim, u_states, y_states, capture_flag,t_data, upwm, udir, \
           data_idx,datasize, prompt,  oled_tprev, oled, ws_robot, wp_robot, \
           v_cmd, w_cmd
    led.value(not led.value())

    if cmd_from_VR:
        v_cmd = adc2vel(v_in.read()) # linear velocity command from VR
        w_cmd = adc2w(w_in.read())   # angular velocity command 

    dd_ikine()  # perform inverse kinematics

    # read wheel speed from robot
    wl_adc = wl_in.read()
    wr_adc = wr_in.read()
    #print("WL adc = "+str(wl_adc)+", WR adc = "+str(wr_adc))

    ws_robot[LEFT_WHEEL] = adc2ws(wl_adc)
    ws_robot[RIGHT_WHEEL] = adc2ws(wr_adc)    
    #print(ws_robot)
    # integrate to get position. For display purpose
    wp_robot[LEFT_WHEEL] = LW_Int.out(ws_robot[LEFT_WHEEL])
    wp_robot[RIGHT_WHEEL] = RW_Int.out(ws_robot[RIGHT_WHEEL])
    # print(wp_robot)
    update_wheels() 

    if feedback:

        u[LEFT_WHEEL] = PID_controller(ws_cmd[LEFT_WHEEL],ws_robot[LEFT_WHEEL],LEFT_WHEEL) # left wheel
        u[RIGHT_WHEEL] = PID_controller(ws_cmd[RIGHT_WHEEL],ws_robot[RIGHT_WHEEL],RIGHT_WHEEL) # right wheel

        ulim[LEFT_WHEEL], upwm[LEFT_WHEEL], udir[LEFT_WHEEL] = convert2pwmdir(u[LEFT_WHEEL])
        ulim[RIGHT_WHEEL], upwm[RIGHT_WHEEL], udir[RIGHT_WHEEL] = convert2pwmdir(u[RIGHT_WHEEL])
    else:
        # change controller outputs manually
        upwm[LEFT_WHEEL] = 100
        udir[LEFT_WHEEL] = 1
        upwm[RIGHT_WHEEL] = 100
        udir[RIGHT_WHEEL] = 0        
    pwm0_out.duty(upwm[LEFT_WHEEL])
    dir0_out.value(udir[LEFT_WHEEL])
    pwm1_out.duty(upwm[RIGHT_WHEEL])
    dir1_out.value(udir[RIGHT_WHEEL])

    # --- remove these lines if SSD1306 is not available ---
    oled_tcurrent = time.ticks_ms()
    if (oled_tcurrent - oled_tprev)>300: # display every 300 ms
        oled_tprev = oled_tcurrent
        oled_show(oled,ws_cmd,ws_robot) # change this
    #    pixels_display()  # display angle on NeoPixels
    # ---------------------------------------------

    
    if capture_flag:

        print("[{},{},{},{},{}],".format(round(t_data,2),r,y,ulim,u))   # change this     
        t_data+=T
        data_idx+=1
        if data_idx==datasize:
            data_idx = 0
            t_data = 0
            capture_flag = False  # stop capture data 
            print("])")  # end of np.array([ command
            prompt=True  # prompt on
            # if lsid:
            #     print("LSID ends.")
            #     lsid = False # LSID mode ends

# helper function for cmdInt()
# separate command and parameter
# at "="
def splitCmd(userstr):
    result = userstr.find("=")
    if result == -1:
        noparm = 1
        cmdstr = userstr.strip()
        parmstr = '' # not used
    else:
        noparm = 0
        splitstr = userstr.split("=")
        cmdstr = splitstr[0].strip()
        parmstr = splitstr[1].strip()
    return cmdstr, parmstr, noparm

#command interpreter function
def cmdInt(userstr):
    global ws_cmd,datasize,capture,capture_flag,feedback,T,T_ms, prompt,\
           kp,ki,kd,kt,wp,wd,N, v_cmd, w_cmd, \
           controller,cmd_from_VR
    cmdstr, parmstr, noparm = splitCmd(userstr)

    if cmdstr.lower() == "lwspeed":  # left wheel speed
        if noparm==1: # retrieve current left wheel speed 
            print(ws_cmd[LEFT_WHEEL])
        else:
            lws_tmp = float(parmstr)

            if lws_tmp > WS_LIM:  # limit speed command 
                lws_tmp = WS_LIM
            elif lws_tmp < -WS_LIM: 
                lws_tmp = -WS_LIM
            ws_cmd[LEFT_WHEEL] = lws_tmp
        if capture:  # capture mode on
            capture_flag = True
            print("datamat = np.array([")  # head of data matrix
            prompt=False # turn prompt off             
    elif cmdstr.lower() == "rwspeed":  # right wheel speed
        if noparm==1: # retrieve current right wheel speed 
            print(ws_cmd[RIGHT_WHEEL])
        else:
            rws_tmp = float(parmstr)

            if rws_tmp > WS_LIM:  # limit speed command 
                rws_tmp = WS_LIM
            elif rws_tmp < -WS_LIM: 
                rws_tmp = -WS_LIM
            ws_cmd[RIGHT_WHEEL] = rws_tmp    
        if capture:  # capture mode on
            capture_flag = True
            print("datamat = np.array([")  # head of data matrix
            prompt=False # turn prompt off 
    elif cmdstr.lower() == "vcmd":  # linear velocity command
        if noparm==1: # retrieve current right wheel speed 
            print(v_cmd)
        else:
            tempvar = float(parmstr)

            if tempvar > VELMAX:  # limit speed command 
                tempvar = VELMAX
            elif tempvar < -VELMAX: 
                tempvar = -VELMAX
            v_cmd = tempvar
    elif cmdstr.lower() == "wcmd":  # angular velocity command
        if noparm==1: # retrieve current right wheel speed 
            print(w_cmd)
        else:
            tempvar = float(parmstr)

            if tempvar > WMAX:  # limit speed command 
                tempvar = WMAX
            elif tempvar < -WMAX: 
                tempvar = -WMAX
            w_cmd = tempvar                
    elif cmdstr.lower() == "capture":
        if noparm==1:
            print(capture)
        else:            
            capture = int(parmstr)
            if capture > 1:
                capture = 1
            elif capture < 0:
                capture = 0
            
    elif cmdstr.lower() == "datasize":
        if noparm==1:
            print(datasize)
        else:            
            datasize = int(parmstr)
            if datasize > 2000: # set maximum datasize
                datasize = 2000
            elif datasize < 10: # minimum datasize
                datasize = 10
    elif cmdstr.lower() == "feedback":
        if noparm==1:
            print(feedback)
        else:            
            feedback = int(parmstr)
            if feedback > 1:
                feedback = 1
            elif feedback < 0:
                feedback = 0

    elif cmdstr.lower() == "t":
        if noparm==1:
            print(T)
        else:            
            T = float(parmstr)
            if T > 1: # set maximum T
                T = 1
            elif T < 0.001: # minimum T
                T = 0.001
            T_ms = int(1000*T)
            PID_update_all()
            # recompute coefficients for output feedback controller
            
            timer.init(period=T_ms, mode=Timer.PERIODIC,
                           callback=timer_isr)
    elif cmdstr.lower() == "lkp":
        if noparm==1:
            print(kp[LEFT_WHEEL])
        else:            
            tempvar = float(parmstr)
            if tempvar > 100: # set maximum kp
                tempvar = 100
            elif tempvar < 0: # minimum kp
                tempvar = 0
            kp[LEFT_WHEEL] = tempvar
            PID_update(LEFT_WHEEL)
    elif cmdstr.lower() == "lki":
        if noparm==1:
            print(ki[LEFT_WHEEL])
        else:            
            tempvar = float(parmstr)
            if tempvar > 100: # set maximum
                tempvar = 100
            elif tempvar < 0: # minimum
                tempvar = 0
            ki[LEFT_WHEEL] = tempvar
            PID_update(LEFT_WHEEL)
    elif cmdstr.lower() == "lkd":
        if noparm==1:
            print(kd[LEFT_WHEEL])
        else:            
            tempvar = float(parmstr)
            if tempvar > 100: # set maximum
                tempvar = 100
            elif tempvar < 0: # minimum
                tempvar = 0
            kd[LEFT_WHEEL] = tempvar
            PID_update(LEFT_WHEEL)
    elif cmdstr.lower() == "rkp":
        if noparm==1:
            print(kp[RIGHT_WHEEL])
        else:            
            tempvar = float(parmstr)
            if tempvar > 100: # set maximum kp
                tempvar = 100
            elif tempvar < 0: # minimum kp
                tempvar = 0
            kp[RIGHT_WHEEL] = tempvar
            PID_update(RIGHT_WHEEL)
    elif cmdstr.lower() == "rki":
        if noparm==1:
            print(ki[RIGHT_WHEEL])
        else:            
            tempvar = float(parmstr)
            if tempvar > 100: # set maximum
                tempvar = 100
            elif tempvar < 0: # minimum
                tempvar = 0
            ki[RIGHT_WHEEL] = tempvar
            PID_update(RIGHT_WHEEL)
    elif cmdstr.lower() == "rkd":
        if noparm==1:
            print(kd[RIGHT_WHEEL])
        else:            
            tempvar = float(parmstr)
            if tempvar > 100: # set maximum
                tempvar = 100
            elif tempvar < 0: # minimum
                tempvar = 0
            kd[RIGHT_WHEEL] = tempvar
            PID_update(RIGHT_WHEEL)    
   
    elif cmdstr.lower() == "reset":
        # reset integrator states
        LW_Int.reset()
        RW_Int.reset()
    elif cmdstr.lower() == "vrmode":
        if noparm==1:
            print(cmd_from_VR)
        else:            
            vr_mode = int(parmstr)
            if vr_mode >= 1:
                cmd_from_VR = True
                print("linear and angular vel cmd from VR")
            elif vr_mode <= 0:
                cmd_from_VR = False
                print("VR inactive")

    elif cmdstr.lower() == "vwcmd":  # check linear & angular velocity command
        print("Cmd linear vel : "+str(v_cmd) +" m/s")
        print("Cmd angular vel : "+str(w_cmd) + " rad/s")

    elif cmdstr.lower() == "ws":  # check wheel speed
        print("Cmd : left = "+str(ws_cmd[LEFT_WHEEL]) + ", right = "+str(ws_cmd[RIGHT_WHEEL])+" rad/s")
        print("Robot : left = "+str(ws_robot[LEFT_WHEEL]) + ", right = "+str(ws_robot[RIGHT_WHEEL])+" rad/s")
    elif cmdstr.lower() == "pwmdir":  # check pwm/dir output
        print("PWM output : left = "+str(upwm[LEFT_WHEEL]) + ", right = "+str(upwm[RIGHT_WHEEL]))
        print("DIR output : left = "+str(udir[LEFT_WHEEL]) + ", right = "+str(udir[RIGHT_WHEEL]))
    elif cmdstr.lower() == "wheeladc":  # check wheel ADC vlue
        print("Wheel ADC value : left = "+str(wl_in.read()) + ", right = "+str(wl_in.read()))
       
    else:
        print("Invalid command")    


def user_input():
    if prompt:
        #sys.stdout.write('\nEnter command : ')
        print("\nEnter command : ")
    #user_str = sys.stdin.readline()
    #newline_char = sys.stdin.readline()
    user_str = input()
    cmdInt(user_str)

timer.init(period=T_ms, mode=Timer.PERIODIC,
                           callback=timer_isr)
PID_update_all()

while True:
    user_input()