// Include necessary libraries
#include <AccelStepper.h>
#include <Servo.h>
#include <SPI.h>
#include <MFRC522.h>
#include <NewPing.h>

// Define pins for components
#define SS_PIN 53         // Slave Select pin for RFID reader
#define US_TRIG_PIN 11   // Ultrasonic sensor trigger pin
#define US_ECHO_PIN 12   // Ultrasonic sensor echo pin
#define X_STEP_PIN 2     // Step pin for X-axis stepper motor driver
#define X_DIR_PIN 5      // Direction pin for X-axis stepper motor driver
#define X_ENABLE_PIN 8   // Enable pin for X-axis stepper motor driver
#define Y1_STEP_PIN 3     // Step pin for first Y-axis stepper motor driver
#define Y1_DIR_PIN 6      // Direction pin for first Y-axis stepper motor driver
#define Y1_ENABLE_PIN 9   // Enable pin for first Y-axis stepper motor driver
#define Y2_STEP_PIN 23    // Step pin for second Y-axis stepper motor driver
#define Y2_DIR_PIN 24     // Direction pin for second Y-axis stepper motor driver
#define Y2_ENABLE_PIN 25  // Enable pin for second Y-axis stepper motor driver
#define SERVO_PIN 7      // Arduino digital pin connected to servo motor

// Define constants
#define RFID_TAG_LENGTH 4         // Length of RFID tag (in bytes)
#define NUM_ROWS 3                // Number of rows of slots
#define NUM_COLS 3                // Number of columns of slots
#define SLOT_WIDTH_CM 16          // Width of each slot in cm
#define SLOT_HEIGHT_CM 16          // Height of each slot in cm
#define X_ORIGIN_CM 8             // X-coordinate origin (leftmost position)
#define Y_ORIGIN_CM 0            // Y-coordinate origin (topmost position, including null space)
#define NULL_SPACE_CM 16          // Length of null space at the beginning of Y-axis
#define SERVO_ROTATION_DELAY 1000 // Delay for servo rotation in milliseconds
#define STEPS_PER_CM_X 100        // Number of steps per cm for X-axis stepper motor
#define STEPS_PER_CM_Y 100        // Number of steps per cm for Y-axis stepper motor
#define STEPS_PER_50_STEPS 50     // Number of steps to move for 50 steps
#define STEPS_PER_100_STEPS 100     // Number of steps to move for 100 steps

// Define instances of necessary objects
AccelStepper stepperX(AccelStepper::DRIVER, X_STEP_PIN, X_DIR_PIN);
AccelStepper stepperY1(AccelStepper::DRIVER, Y1_STEP_PIN, Y1_DIR_PIN);
AccelStepper stepperY2(AccelStepper::DRIVER, Y2_STEP_PIN, Y2_DIR_PIN);
Servo servoX;
MFRC522 rfid(SS_PIN, 9);  // Create MFRC522 instance
NewPing ultrasonicSensor(US_TRIG_PIN, US_ECHO_PIN); // Initialize the ultrasonic sensor

// Define variables and arrays
bool slotOccupied[NUM_ROWS][NUM_COLS] = {{false}};  // Boolean array to track slot occupancy
byte slotRFIDCodes[NUM_ROWS][NUM_COLS][RFID_TAG_LENGTH];  // 3D array to store RFID codes for each slot
long originalX = 0;  // Store the original position of the gantry in X-axis
long originalY1 = 0;  // Store the original position of the first gantry motor in Y-axis
long originalY2 = 0;  // Store the original position of the second gantry motor in Y-axis
float liftAmount = 2.0; // Adjust this value according to your system

void setup() {
  Serial.begin(9600);
  SPI.begin();        // Init SPI bus
  rfid.PCD_Init();    // Init RFID reader

  // Set stepper motor speed and other configurations
  stepperX.setMaxSpeed(100); // Set maximum speed in steps per second
  stepperX.setAcceleration(500); // Set acceleration in steps per second^2
  stepperY1.setMaxSpeed(100);
  stepperY1.setAcceleration(500);
  stepperY2.setMaxSpeed(100);
  stepperY2.setAcceleration(500);

  // Attach servo to pin
  servoX.attach(SERVO_PIN);

  // Initialize ultrasonic sensor
  pinMode(US_TRIG_PIN, OUTPUT);
  pinMode(US_ECHO_PIN, INPUT);


  // Store the original position
  originalX = stepperX.currentPosition();
  originalY1 = stepperY1.currentPosition();
  originalY2 = stepperY2.currentPosition();
}

void loop() {
  // Look for new RFID cards
  if (rfid.PICC_IsNewCardPresent() && rfid.PICC_ReadCardSerial()) {

    // Read the RFID tag
    byte tagData[RFID_TAG_LENGTH];
    for (byte i = 0; i < rfid.uid.size; i++) {
      tagData[i] = rfid.uid.uidByte[i];
    }

    // Check if the tag matches any previously assigned slot
    int slotRow, slotCol;
    if (retrieveSlot(tagData, slotRow, slotCol)) {
      // Retrieve the bag from the assigned slot
      retrieveBagFromSlot(slotRow, slotCol);
    } else {
      // Assign the RFID code to a new slot and start the bag storing process
      if (assignNewSlot(tagData, slotRow, slotCol)) {
        // Successfully assigned to a new slot, start the bag storing process
        startBagStoringProcess(slotRow, slotCol);

      }
    }

    delay(1000); // Delay to avoid reading the same tag multiple times
  }
}
// Function to retrieve the slot based on the stored RFID tag
bool retrieveSlot(byte tagData[], int& slotRow, int& slotCol) {
  for (int i = 0; i < NUM_ROWS; i++) {
    for (int j = 0; j < NUM_COLS; j++) {
      if (slotOccupied[i][j] && compareTags(tagData, slotRFIDCodes[i][j], RFID_TAG_LENGTH)) {
        slotRow = i;
        slotCol = j;
        return true; // Slot found
      }
    }
  }
  return false; // Slot not found
}
// Function to assign the RFID code to the nearest available slot
bool assignNewSlot(byte tagData[], int& slotRow, int& slotCol) {
  // Search for the nearest available slot
  for (int radius = 0; radius < max(NUM_ROWS, NUM_COLS); radius++) {
    for (int i = -radius; i <= radius; i++) {
      for (int j = -radius; j <= radius; j++) {
        int row = NUM_ROWS / 2 + i; // Calculate row index
        int col = NUM_COLS / 2 + j; // Calculate column index

        // Check if the slot is within bounds and available
        if (row >= 0 && row < NUM_ROWS && col >= 0 && col < NUM_COLS && !slotOccupied[row][col]) {
          // Assign the RFID code to the nearest available slot
          for (int k = 0; k < RFID_TAG_LENGTH; k++) {
            slotRFIDCodes[row][col][k] = tagData[k];
          }
          slotOccupied[row][col] = true; // Mark the slot as occupied
          slotRow = row; // Store the assigned row
          slotCol = col; // Store the assigned column
          return true; // Slot successfully assigned
        }
      }
    }
  }
  
  return false; // No available slots found
}

void retrieveBagFromSlot(int slotRow, int slotCol){
  //Calculate the coordinates of the slot
  float slotX = X_ORIGIN_CM + (slotCol * SLOT_WIDTH_CM);
  float slotY = Y_ORIGIN_CM + NULL_SPACE_CM + (slotRow * SLOT_HEIGHT_CM);


  // Convert the slot co ordinates to steps
  long targetX = slotX * STEPS_PER_CM_X;
  long targetY1 = slotY * STEPS_PER_CM_Y;
  long targetY2 = slotY * STEPS_PER_CM_Y;

  // Move to the assigned slot
  stepperX.moveTo(targetX);
  stepperY1.moveTo(targetY1);
  stepperY2.moveTo(targetY2);

  // Wait until all motors reach their target position
  while (stepperX.distanceToGo() != 0 || stepperY1.distanceToGo() != 0 || stepperY2.distanceToGo() != 0) {
    stepperX.run();
    stepperY1.run();
    stepperY2.run();
    delay(1000);
  }

  liftAxis();

  moveToInitialPosition();

}

void startBagStoringProcess(int slotRow, int slotCol) {
  // Calculate the coordinates of the slot
  float slotX = X_ORIGIN_CM + (slotCol * SLOT_WIDTH_CM);
  float slotY = Y_ORIGIN_CM + NULL_SPACE_CM + (slotRow * SLOT_HEIGHT_CM);

  // Convert the slot coordinates to steps
  long targetX = slotX * STEPS_PER_CM_X;
  long targetY1 = slotY * STEPS_PER_CM_Y;
  long targetY2 = slotY * STEPS_PER_CM_Y;

  // Move to the assigned slot
  stepperX.moveTo(targetX);
  stepperY1.moveTo(targetY1);
  stepperY2.moveTo(targetY2);

  // Wait until all motors reach their target position
  while (stepperX.distanceToGo() != 0 || stepperY1.distanceToGo() != 0 || stepperY2.distanceToGo() != 0) {
    stepperX.run();
    stepperY1.run();
    stepperY2.run();
    delay(1000);
  }

  liftAxis(); // Lift the axis

  moveToInitialPosition(); // Move to the initial position to collect bags

  if (checkBagPresence()){; // If the bag was placed on the tray

    // Calculate the new target position after lifting
    targetX = slotX * STEPS_PER_CM_X ; // Move X-axis by 50 steps
    targetY1 = slotY * STEPS_PER_CM_Y + STEPS_PER_100_STEPS; // Keep Y-axis at the same position
    targetY2 = slotY * STEPS_PER_CM_Y + STEPS_PER_100_STEPS; // Keep Y-axis at the same position

    // Move to the new position
    stepperX.moveTo(targetX);
    stepperY1.moveTo(targetY1);
    stepperY2.moveTo(targetY2);

    // Wait until all motors reach their target position
    while (stepperX.distanceToGo() != 0 || stepperY1.distanceToGo() != 0 || stepperY2.distanceToGo() != 0) {
      stepperX.run();
      stepperY1.run();
      stepperY2.run();
      delay(1000);
    }

    lowerYAxis(); // Lower the axis
  }

}

// Function to check the presence of a bag using ultrasonic sensor
bool checkBagPresence() {
  // Trigger ultrasonic sensor to get distance
  long duration = ultrasonicSensor.ping_median(5); // Median filter with 5 samples for accuracy
  float distance = duration / 29.1 / 2; // Convert duration to distance in cm

  // Adjust the threshold according to your setup
  if (distance < 10) { // Adjust this threshold as needed
    return true;
  } else {
    return false;
  }
}

// Function to move the gantry to the initial position
void moveToInitialPosition() {
  stepperX.moveTo(originalX);
  stepperY1.moveTo(originalY1);
  stepperY2.moveTo(originalY2);

  // Wait until all motors reach their target position
  while (stepperX.distanceToGo() != 0 || stepperY1.distanceToGo() != 0 || stepperY2.distanceToGo() != 0) {
    stepperX.run();
    stepperY1.run();
    stepperY2.run();
    delay(1);
  }

}

void liftAxis(){
  // Rotate servo 180 degrees
  rotateServoIn();
  
  // Move Y-axis stepper motor up by 50 steps
  stepperY1.move(STEPS_PER_50_STEPS);
  stepperY2.move(STEPS_PER_50_STEPS);
  stepperY1.runToPosition(); // Wait for the motor to reach the target position
  stepperY2.runToPosition(); // Wait for the motor to reach the target position

  // Rotate servo 180 degrees backwards
  rotateServoOut();

}
// Function to lower the Y-axis by the same amount lifted and rotate servo
void lowerYAxis() {
  //rotate servo 180 degrees
  rotateServoIn();
  // Move Y-axis stepper motor down by 100 steps
  stepperY1.move(-STEPS_PER_100_STEPS);
  stepperY2.move(-STEPS_PER_100_STEPS);
  stepperY1.runToPosition(); // Wait for the motor to reach the target position
  stepperY2.runToPosition(); // Wait for the motor to reach the target position
  // Rotate servo 180 degrees back
  rotateServoOut();
}
// Function to rotate the servo on the X-axis
void rotateServoIn() {
  // Rotate servo 180 degrees
  servoX.write(180);
  delay(SERVO_ROTATION_DELAY); // Adjust delay according to servo speed
}
// Function to rotate the servo on the X-axis
void rotateServoOut() {
  // Rotate servo 180 degrees
  servoX.write(0);
  delay(SERVO_ROTATION_DELAY); // Adjust delay according to servo speed
}
// Function to compare two RFID tags
bool compareTags(byte tag1[], byte tag2[], int length) {
  for (int i = 0; i < length; i++) {
    if (tag1[i] != tag2[i]) {
      return false;
    }
  }
  return true;
}