'''
Demo of PID control to drive the car in a straight line
  - The PID control uses gyro angle input and adjusts the left / right wheel speeds to keep the angle 0
  v2 : redo the get_gyro function to be a class so as to remove global variables
'''
import machine
from time import sleep_ms, ticks_ms
import pca9685
import motor
from MPU6050 import MPU6050
from Car_lib import accel_calibration, motor_move
from os import listdir

# set up i2c and pca9685
sdaPIN=machine.Pin(13)
sclPIN=machine.Pin(14)
i2c=machine.I2C(0,sda=sdaPIN, scl=sclPIN)

# Initialize the PCA9685 using the default address (0x40).
pwm = pca9685.PCA9685(address=0x5f, i2c = i2c)
# Set the PWM frequency to 60Hz, good for servos.
pwm.freq(60)

# initialize motors
motors = motor.DCMotors(address=0x5f, i2c = i2c)

# initial accelerometer and gyro
mpu = MPU6050()

class GyroSensor:
    def __init__(self, mpu, gzc):
        self.mpu = mpu
        self.gzc = gzc
        self.angZ = 0
        self.gZ = 0
        self.t_old = 0

    def get_gyro(self, reset=False):
        if reset:
            self.angZ = 0
            self.t_old = 0
            return self.angZ

        gyro = self.mpu.read_gyro_data()
        self.gZ = gyro["z"] - self.gzc
        t = ticks_ms()
        
        if self.t_old == 0:
            self.t_old = t
            dt = 20  # Assume 20ms for the first reading
        else:
            dt = t - self.t_old
        
        self.angZ += self.gZ * (dt / 1000)  # box integration
        self.t_old = t

        return self.angZ
    

def drive_straight(gyro_sensor, speed, duration):
    global integral, prev_error, told
    start_time = ticks_ms()
    stflag = True
    while ticks_ms() - start_time < duration:
        # Read gyro angle
        if stflag:
            angle = gyro_sensor.get_gyro(reset = True)
            told = start_time
        else:
            angle = gyro_sensor.get_gyro()
        stflag = False
        
        # Calculate error
        error = Target_angle - angle

        # PID calculation
        integral += error
        derivative = error - prev_error
        correction = (kp * error) + (ki * integral) + (kd * derivative)
        prev_error = error

        # Apply correction to motors
        left_speed = speed - correction
        right_speed = speed + correction

        # Set motor speeds
        motor_move(motors,
                   m1_speed=left_speed,
                   m2_speed=left_speed,
                   m3_speed=right_speed,
                   m4_speed=right_speed
                   )
        # store data
        t = ticks_ms()
        dt = t - told
        told = t
        
        data = {"Time" : t,
                "Target" : Target_angle,
                    "Measure" : angle,
                    "Error" : error,
                    "Correction" : correction,
                    "Left" : left_speed,
                    "Right" : right_speed,
                    "dt"   : dt
                    }
        push = [  str(data[k]) for k in cols ]
        row = ",".join(push)
        file.write(row + "\n")
    
#         print('angle  ',angle)
#         print('error  ',error)
#         print('left_speed  ',left_speed)
#         print('right_speed  ',right_speed)

        sleep_ms(10)

#  ------------   main program   -------------------

# PID constants
'''
To fine-tune the PID gains:
Start with only the proportional gain (set Ki and Kd to zero)
Gradually increase Kp until you observe a fast response without excessive overshoot
Introduce a small Ki to eliminate steady-state error
Add Kd to reduce overshoot and improve stability if needed

'''
kp = 0.5
ki = 0.1
kd = 0.2

kp = 800
ki = 0
kd = 0

kp = 800
ki = 40
kd = 160

# create file to save results
# List all files directory.
print("Root directory: {} \n".format(listdir()))
# Change filename and fileformat here.
filename_template = "PID_logs{}.txt"
# Increment the filename number if the file already exists.
i = 0
while filename_template.format(i) in listdir():
    i += 1
filename = filename_template.format(i)

# open file for writing
print(filename)
file = open(filename, 'w')
cols = ["Time","Measure", "Error", "Correction", "Left", "Right","dt"]
file.write(",".join(cols) + "\n")

# Target angle (straight ahead)
Target_angle = 0

# delay to unplug car
sleep_ms(1000)

# find calibration values
axc, ayc, azc, gzc = accel_calibration(mpu)

# Initialize PID variables
integral = 0
prev_error = 0

# move forward
gyro_sensor = GyroSensor(mpu, gzc)
drive_straight(gyro_sensor, 2000, 10000)

# stop
motor_move(motors)

file.close()