#include <PID_v2.h>                 // SYSTEM PARAMETER - PID_v1 Compatible Library 
#include <LiquidCrystal_I2C.h>      // SYSTEM PARAMETER - ESP32 LCD Compatible Library (By: Robojax)
#include <PubSubClient.h>           // SYSTEM PARAMETER - MQTT Library  (By: Adafruit)

// Motor Parameters and Constants
#define CW          (1)
#define CCW         (-1)
#define STOP        (0)

#define NMOTORS     (2)
#define M0          (0)
#define M1          (1)
#define M2          (2)  // Define M2 and M3 for additional motors
#define M3          (3)

#define ESP         (1)

#ifdef ESP
#define ENC_A1      (16)
#define ENC_A2      (17)
#define PWM_A       (18)
#define INT_A1      (4)
#define INT_A2      (2)

#define ENC_B1      (14)
#define ENC_B2      (15)
#define PWM_B       (19)
#define INT_B1      (25)
#define INT_B2      (26)
#endif

const int enca[] = {ENC_A1, ENC_B1};  // Encoder A pins for each motor
const int encb[] = {ENC_A2, ENC_B2};  // Encoder B pins for each motor
const int pwm[]  = {PWM_A, PWM_B};     // PWM pins for each motor
const int in1[]  = {INT_A1, INT_B1};   // Motor direction pin 1 for each motor
const int in2[]  = {INT_A2, INT_B2};   // Motor direction pin 2 for each motor

#define Umax        (255)    // Max speed motor
#define Debounce    (3)      // Software debounce

#define STOP        (0)
#define START       (1)

#define adc_min     (0)      // Analog control minimum value
#define adc_max     (100)    // Analog control maximum value
#define adc_res_min (0)      // ADC resolution 2^11 bit minimum value
#define adc_res_max (4095)   // ADC resolution 2^11 bit maximum value
#define gain_k_min  (0)      // Analog control minimum gain
#define gain_k_max  (10)     // Analog control maximum gain

#define InputPin    (32)     // Input pin for analog read
#define KiPin       (33)     // Input pin for Ki adjustment
#define KpPin       (34)     // Input pin for Kp adjustment
#define KdPin       (35)     // Input pin for Kd adjustment

#define SWITCH_ON   (12)     // Pin D12 for SWITCH ON
#define SWITCH_DIR  (13)     // Pin D13 for SWITCH OFF

#define BUTTON_START digitalRead(SWITCH_ON)  == 1 // Check if start button is pressed
#define BUTTON_DIR   digitalRead(SWITCH_DIR) == 1 // Check if direction button is pressed

// Global variables
volatile int target[]   = {0, 0, 0, 0};           // Target positions for each motor
volatile int target_f[] = {0, 0, 0, 0};           // Target positions in float for each motor
volatile int posi[]     = {0, 0, 0, 0};           // Current positions for each motor
volatile int pos[]      = {0, 0, 0, 0};           // Position copy for each motor
float positionChange[4] = {0.0, 0.0, 0.0, 0.0};   // Position change array for each motor

long prevT = 0;                                   // Previous time in microseconds
int encoder0Pos = 0, ENC_A1Last = 0, n = 0, Motor_dir = 0; // Encoder and motor variables
int Forward = 1, Backward = -1;                   // Motor direction constants

double Setpoint[] = {0, 0, 0, 0};                 // Setpoints for PID controllers
double Input[]    = {0, 0, 0, 0};                 // Inputs for PID controllers
double Output[]   = {0, 0, 0, 0};                 // Outputs from PID controllers

// PID parameters
double 
Ma_kp = 1.0, 
Ma_ki = 0.000011, 
Ma_kd = 0.00000011,
Mb_kp = 1.0, 
Mb_ki = 0.000011, 
Mb_kd = 0.00000011;

double 
Kp[] = {Ma_kp, Mb_kp}, 
Ki[] = {Ma_ki, Mb_ki}, 
Kd[] = {Ma_kd, Mb_kd};

double 
aggKp  = 1.0, 
aggKi  = 0.000011, 
aggKd  = 0.00000011,
consKp = 0.5, 
consKi = 0.000005, 
consKd = 0.00000005;

// Motor and encoder parameters
float 
t_period = 0.0, 
velocity = 1.0, 
Turn     = 20.0, 
pulses   = 11.0, 
pulsesPerTurn = pulses * Turn, 
Meter = 3.5714,
pulsesPerMeter = pulsesPerTurn * Meter;

int 
rpm[] = {0, 0}, 
outputpwm[] = {0, 0};

volatile unsigned int event = 1;
volatile unsigned long lastime = 0;
volatile char Serial_input_read;

struct command {
    union {
        unsigned short WORD;
        struct {
            unsigned char :8;
            unsigned char ButtonDir :1;
            unsigned char Button :1;
            unsigned char Rotate :1;
            unsigned char Enable :1;
            unsigned char Dir :1;
            unsigned char Reset :1;
            unsigned char Emer :1;
            unsigned char Run :1;
        } BIT;
    } CTRL;
} CMD;

// FreeRTOS variables
TaskHandle_t xTask0h = NULL, xTask1h = NULL, xTask2h = NULL, xTask3h = NULL, xTask5h = NULL, xTask4h = NULL, xTask6h = NULL, xTask7h = NULL;
QueueHandle_t xQueue_y = NULL, xQueue_u = NULL;
bool xQueueOK = LOW;

int T0_ms = 0, T1_ms = 0, T2_ms = 0, T3_ms = 0, T4_ms = 0, T5_ms = 0, T6_ms = 0, T7_ms = 0;
int T0ticks = 0, T1ticks = 0, T2ticks = 0, T3ticks = 0, T4ticks = 0, T5ticks = 0, T6ticks = 0, T7ticks = 0;
const float Tsampling = 0.01;   // Sampling period
const int SampleRate = 10; // Define SampleRate for PID

float Step_rev, Step_kp, Step_ki, Step_kd; // Declare Step_rev, Step_kp, Step_ki, Step_kd

// PID controllers for each motor
PID M0PID(&Input[M0], &Output[M0], &Setpoint[M0], Kp[M0], Ki[M0], Kd[M0], DIRECT);  
PID M1PID(&Input[M1], &Output[M1], &Setpoint[M1], Kp[M1], Ki[M1], Kd[M1], DIRECT);  

LiquidCrystal_I2C lcd(0x3F, 16, 2); // Initialize LCD with I2C address and size

// Ziegler-Nichols parameters (example values)
double L = 0.1; // Dead time
double T = 1.0; // Time constant

void setup() 
{
    Serial.begin(115200);
    pinMode(InputPin, INPUT);
    pinMode(KiPin, INPUT);
    pinMode(KpPin, INPUT);
    pinMode(KdPin, INPUT);
    pinMode(SWITCH_ON, INPUT);
    pinMode(SWITCH_DIR, INPUT);

    for (int k = 0; k < NMOTORS; k++) 
    {
        pinMode(enca[k], INPUT);
        pinMode(encb[k], INPUT);
        pinMode(pwm[k], OUTPUT);
        pinMode(in1[k], OUTPUT);
        pinMode(in2[k], OUTPUT);
    }

    freertos_init(); // Initialize FreeRTOS tasks and queues

    // Attach interrupts for encoders
    attachInterrupt(digitalPinToInterrupt(enca[M0]), readEncoder<M0>, FALLING);
    attachInterrupt(digitalPinToInterrupt(enca[M1]), readEncoder<M1>, FALLING);

    // Calculate PID parameters using Ziegler-Nichols method
    double Kp_zn = 1.2 * T / L;
    double Ki_zn = 2 * L;
    double Kd_zn = 0.5 * L;

    // Update PID controllers
    M0PID.SetTunings(Kp_zn, Ki_zn, Kd_zn);
    M1PID.SetTunings(Kp_zn, Ki_zn, Kd_zn);

    // Set sample time and mode for PID controllers
    M0PID.SetSampleTime(SampleRate);
    M1PID.SetSampleTime(SampleRate);
    M0PID.SetMode(AUTOMATIC);
    M1PID.SetMode(AUTOMATIC);
    M0PID.SetOutputLimits(-255, 255);
    M1PID.SetOutputLimits(-255, 255);

    delay(1000); // Delay to allow initialization

    lcd.init(); // Initialize LCD
    lcd.setBacklight(HIGH);
    lcd.setCursor(0, 0);
    lcd.print("PID Motor");
    lcd.setCursor(0, 1);
    lcd.print("Control");
    delay(1500);
    lcd.clear();

    CMD.CTRL.BIT.Run = START; // Enable motor RUN mode
}

void loop() {}

// Motor control function
void MotorControl(int motor) 
{
    long currT = micros(); // Current time in microseconds
    float deltaT = ((float)(currT - prevT)) / (1.0e6); // Time difference in seconds
    prevT = currT; // Update previous time
    SetTarget(currT / 1.0e6, deltaT, Motor_dir); // Set target position

    noInterrupts(); // Disable interrupts temporarily while reading
    pos[motor] = posi[motor]; // Read the position in an atomic block
    interrupts(); // Enable interrupts

    int pwr, dir;
    Input[motor] = target[motor]; // Set target data from step control
    Setpoint[motor] = posi[motor]; // Modify to fit motor and encoder characteristics

    double gap = abs(Setpoint - Input); // Distance away from setpoint

    if (motor == M0) 
    {
        if (gap < 2) 
        {
            M0PID.SetTunings(consKp, consKi, consKd); // Use conservative tuning parameters
        } 
        else 
        {
            M0PID.SetTunings(aggKp, aggKi, aggKd); // Use aggressive tuning parameters
        }
        M0PID.Compute(); // Calculate new output
    }

    if (motor == M1) 
    {
        if (gap < 2) 
        {
            M1PID.SetTunings(consKp, consKi, consKd); // Use conservative tuning parameters
        } 
        else 
        {
            M1PID.SetTunings(aggKp, aggKi, aggKd); // Use aggressive tuning parameters
        }
        M1PID.Compute(); // Calculate new output
    }

    pwr = (int)fabs(Output[motor]); // Motor power
    if (pwr > Umax) 
    {
        pwr = Umax; // Limit maximum pwm
    } 
    else if (pwr < 0) 
    {
        pwr = 0; // Limit minimum pwm
    }

    dir = 1; // Motor direction positive
    if (Output[motor] < 0) dir = -1; // Motor direction negative
    SetMotor(dir, pwr, pwm[motor], in1[motor], in2[motor]); // Signal the motor
    outputpwm[motor] = pwr; // Output PWM for display
}

// Set target position function
void SetTarget(float t, float deltat, int dir) 
{
    t = fmod(t, 20); // Time is in seconds
    if (t < 5) 
    {
        dir = Motor_dir; // Wait 4 seconds
    } 
    else if (t < 10) 
    {
        dir = Forward; // CW Forward 4 seconds
    } 
    else if (t < 15) 
    {
        dir = Motor_dir; // Wait 4 seconds
    } 
    else 
    {
        dir = Backward; // CCW Backward 4 seconds
    }

    for (int k = 0; k < 4; k++) 
    {
        positionChange[k] = dir * velocity * deltat * pulsesPerMeter; // Position change in encoder pulses
        target_f[k] = target_f[k] + positionChange[k]; // Set target to move robot
    }

    target[M0] = (long)target_f[M0]; // Wheel 1
    target[M1] = (long)-target_f[M1]; // Wheel 2
    target[M2] = (long)target_f[M2]; // Wheel 3
    target[M3] = (long)-target_f[M3]; // Wheel 4
}

// Set motor function
void SetMotor(int dir, int pwmVal, int pwm, int in1, int in2) 
{
    analogWrite(pwm, pwmVal); // Write analog PWM output
    if (dir == CW) 
    {
        digitalWrite(in1, HIGH); // Drive motor CW
        digitalWrite(in2, LOW);
    } 
    else if (dir == CCW) 
    {
        digitalWrite(in1, LOW); // Drive motor CCW
        digitalWrite(in2, HIGH);
    } 
    else 
    {
        digitalWrite(in1, LOW); // Stop motor
        digitalWrite(in2, LOW);
    }
}

// Read encoder function
template <int j>
void readEncoder() 
{
    int pulse = digitalRead(encb[j]);
    if (pulse) 
    {
        posi[j]++;
    } 
    else 
    {
        posi[j]--;
    }
}

// Serial display function
void SerialDisplay() 
{
    for (int k = 0; k < NMOTORS; k++) 
    {
        Serial.print(target[k]);
        Serial.print(" ");
        Serial.print(pos[k]);
        Serial.print(" ");
    }
    Serial.println();
}

// LCD display function
void LCDdisplay() 
{
    lcd.setCursor(0, 0);
    lcd.print("M0:" + String(outputpwm[M0]) + " M1:" + String(outputpwm[M1]));
    lcd.setCursor(0, 1);
    lcd.print("Velocity:" + String(velocity) + "m/s");
}

// Button control function
void ButtonControl() 
{
    if (BUTTON_START) 
    {
        long now = millis();
        if ((now - lastime) > Debounce) 
        {
            lastime = 0;
            CMD.CTRL.BIT.Run ^= 1;
        }
    } else if (BUTTON_DIR) 
    {
        long now = millis();
        if ((now - lastime) > Debounce) 
        {
            lastime = 0;
            CMD.CTRL.BIT.Dir ^= 1;
        }
    } 
    else 
    {
        lastime = 0;
    }
}

// Serial interface function
void SerialInterface(int motor) 
{
    if (Serial.available())
    {
        Serial_input_read = Serial.read();
        Serial.println(Serial_input_read);
        if (!CMD.CTRL.BIT.Run) 
        {
            if (Serial_input_read == '0') 
            {
                Kp[motor] += 0.0001;
                Serial.println(Kd[motor]);
            }
            if (Serial_input_read == '1') 
            {
                Ki[motor] += 0.0001;
                Serial.println(Ki[motor]);
            }
            if (Serial_input_read == '2') 
            {
                Kd[motor] += 0.0001;
                Serial.println(Kd[motor]);
            }
        }
    }
}

// Adjust gain PID function
void AdjusGainPID(int motor) 
{
    if (!CMD.CTRL.BIT.Run) 
    {
        Kp[motor] = Step_kp;
        Ki[motor] = Step_ki;
        Kd[motor] = Step_kd;
    }
}

// Initialize FreeRTOS tasks and queues
void freertos_init() 
{
    xQueueOK = LOW;
    xQueue_y = xQueueCreate(1000, sizeof(float));
    xQueue_u = xQueueCreate(1000, sizeof(float));

    if ((xQueue_y == NULL) | (xQueue_u == NULL)) Serial.println("Error creating the queue");
    else xQueueOK = HIGH;

    T0_ms = 1000 * Tsampling;
    T1_ms = 5;
    T2_ms = 100;
    T3_ms = 1000;
    T4_ms = 5;
    T5_ms = 1000 * Tsampling;
    T6_ms = 1000;
    T7_ms = 1000;

    T0ticks = pdMS_TO_TICKS(T0_ms);
    T1ticks = pdMS_TO_TICKS(T1_ms);
    T2ticks = pdMS_TO_TICKS(T2_ms);
    T3ticks = pdMS_TO_TICKS(T3_ms);
    T4ticks = pdMS_TO_TICKS(T4_ms);
    T5ticks = pdMS_TO_TICKS(T5_ms);
    T6ticks = pdMS_TO_TICKS(T6_ms);
    T7ticks = pdMS_TO_TICKS(T7_ms);

    xTaskCreatePinnedToCore(Task0, "Task0", 10000, NULL, 1, &xTask0h, 0);
    xTaskCreatePinnedToCore(Task1, "Task1", 10000, NULL, 2, &xTask1h, 1);
    xTaskCreatePinnedToCore(Task2, "Task2", 10000, NULL, 2, &xTask2h, 0);
    xTaskCreatePinnedToCore(Task3, "Task3", 10000, NULL, 1, &xTask2h, 1);
    xTaskCreatePinnedToCore(Task4, "Task4", 10000, NULL, 2, &xTask3h, 1);
    xTaskCreatePinnedToCore(Task5, "Task5", 10000, NULL, 1, &xTask5h, 1);
    xTaskCreatePinnedToCore(Task6, "Task6", 10000, NULL, 1, &xTask6h, 1);
    xTaskCreatePinnedToCore(Task7, "Task7", 10000, NULL, 1, &xTask7h, 1);
}

// Task 0 function
void Task0(void* parameter) 
{
    TickType_t xLastWakeTime;
    xLastWakeTime = xTaskGetTickCount();
    for (;;) 
    {
        SerialDisplay(); // Call serial display
        vTaskDelayUntil(&xLastWakeTime, T0ticks); // Set period 
    }
    vTaskDelete(NULL);
}

// Task 1 function
void Task1(void* parameter) 
{
    for (;;) 
    {
        MotorControl(M0); // Correct position
        vTaskDelay(T1ticks); // Set period 
    }
    vTaskDelete(NULL);
}

// Task 2 function
void Task2(void* parameter) 
{
    for (;;) 
    {
        velocity = (Step_rev / 35.0); // Set motor velocity controller
        vTaskDelay(T2ticks); // Set period 
    }
    vTaskDelete(NULL);
}

// Task 3 function
void Task3(void* parameter) 
{
    for (;;) 
    {
        // ButtonControl(); // Uncomment if needed
        vTaskDelay(T3ticks); // Set period 
    }
    vTaskDelete(NULL);
}

// Task 4 function
void Task4(void* parameter) 
{
    for (;;) 
    {
        MotorControl(M1); // Correct position
        vTaskDelay(T4ticks); // Set period 
    }
    vTaskDelete(NULL);
}

// Task 5 function
void Task5(void* parameter) 
{
    TickType_t xLastWakeTime;
    xLastWakeTime = xTaskGetTickCount();
    for (;;) 
    {
        LCDdisplay(); // Update LCD display
        vTaskDelayUntil(&xLastWakeTime, T5ticks); // Set period 
    }
    vTaskDelete(NULL);
}

// Task 6 function
void Task6(void* parameter) 
{
    for (;;) 
    {
        if (!CMD.CTRL.BIT.Dir) 
        {
            SerialInterface(M0); // Handle serial interface for motor 0
        } 
        else 
        {
            SerialInterface(M1); // Handle serial interface for motor 1
        }
        vTaskDelay(T6ticks); // Set period 
    }
    vTaskDelete(NULL);
}

// Task 7 function
void Task7(void* parameter) 
{
    for (;;) 
    {
        if (!CMD.CTRL.BIT.Dir) 
        {
            AdjusGainPID(M0); // Adjust PID gains for motor 0
        } 
        else 
        {
            AdjusGainPID(M1); // Adjust PID gains for motor 1
        }
        vTaskDelay(T7ticks); // Set period 
    }
    vTaskDelete(NULL);
}