#include <Servo.h>

#define LOAD_SENSOR_PIN A0
#define BUTTON_PIN 7  // Pin for the push button

Servo servo1, servo2, servo3;
Servo gripperServo;  // Gripper servo

bool buttonPressed = false;
int currentState = 0;
int lastState = 0;

const int homeAngles[] = {45, 45, 45};        // Home position angles for servo1, servo2, servo3
const int place1Angles[] = {90, 90, 90};
const int place2Angles[] = {120, 120, 120};
const int place3Angles[] = {150, 150, 150};
const int place4Angles[] = {180, 180, 180};
const int pickPlaceAngles[] = {0, 0, 0}; // Picking place angles

void setup() {
    servo1.attach(3);  // Attach servo to PWM pin 3
    servo2.attach(5);  // Attach servo to PWM pin 5
    servo3.attach(6);  // Attach servo to PWM pin 6
    gripperServo.attach(10);  // Attach gripper servo to PWM pin 10

    pinMode(BUTTON_PIN, INPUT);
    lastState = digitalRead(BUTTON_PIN);
  
    Serial.begin(9600);
    
    // Move arm to home position during setup
    moveArmToHomePosition();
}

void loop() {
    currentState = digitalRead(BUTTON_PIN);
  
    if (currentState == HIGH && lastState == LOW) {
        buttonPressed = true;
    } else {
        buttonPressed = false;
    }
  
    lastState = currentState;
  
    if (buttonPressed) {
        performPickAndPlace();
    }
}

void performPickAndPlace() {
    float weight = readLoadSensor();
    int targetPlace = determineTargetPlace(weight);
    moveArmToPickObject(targetPlace);
    delay(1000);  // Delay after picking
    
    moveArmToDropPlace(targetPlace);
    delay(1000);  // Delay after dropping
    
    moveArmToHomePosition(); // Move back to home position
    
    // Reset the buttonPressed flag
    buttonPressed = false;
}

float readLoadSensor() {
    int sensorValue = analogRead(LOAD_SENSOR_PIN);
    float weight = map(sensorValue, 0, 1023, 0, 600);  // Map the sensor value to weight range 0g - 600g
    return weight;
}

int determineTargetPlace(float weight) {
    // Placeholder implementation for determining the target place based on weight
    // This is just a sample logic, you need to adjust this to your requirements
    if (weight >= 0 && weight <= 150) {
        return 1;
    } else if (weight > 150 && weight <= 300) {
        return 2;
    } else if (weight > 300 && weight <= 450) {
        return 3;
    } else if (weight > 450 && weight <= 600) {
        return 4;
    }
}

void moveArmToPickObject(int targetPlace) {
    // Move to picking place angles
    setServoAngles(pickPlaceAngles);
    delay(1000);
    
    gripObject();  // Close the gripper
    delay(1000);
}

void moveArmToDropPlace(int targetPlace) {
    // Move to drop place angles
    if (targetPlace == 1) {
        setServoAngles(place1Angles);
    } else if (targetPlace == 2) {
        setServoAngles(place2Angles);
    } else if (targetPlace == 3) {
        setServoAngles(place3Angles);
    } else if (targetPlace == 4) {
        setServoAngles(place4Angles);
    }
    delay(1000);
    
    releaseObject();  // Open the gripper
    delay(1000);
}

void setServoAngles(int angles[]) {
    servo1.write(angles[0]);
    servo2.write(angles[1]);
    servo3.write(angles[2]);
}

void gripObject() {
    // Close the gripper
    gripperServo.write(90);
    delay(1000);  // Delay for gripping
}

void releaseObject() {
    // Open the gripper
    gripperServo.write(0);
    delay(1000);  // Delay for releasing
}

void moveArmToHomePosition() {
    // Move to home position angles for servo1, servo2, servo3
    setServoAngles(homeAngles);
    delay(1000);
}