#include <Servo.h>       // Include the Servo library to control servo motors
#include <math.h>        // Include math library for trigonometric and square root functions

// === Configuration Constants ===

// Lengths of the robot arm segments in centimeters
const float L1 = 10.0;   // Length from shoulder to elbow
const float L2 = 12.5;   // Length from elbow to wrist
const float L3 = 5.0;    // Length from wrist to gripper (currently not used in calculations)

// Pin assignments for servos
const int SERVO_BASE_PIN     = 3;   // Base rotation servo
const int SERVO_SHOULDER_PIN = 5;   // Shoulder joint servo
const int SERVO_ELBOW_PIN    = 6;   // Elbow joint servo
const int SERVO_WRIST_PIN    = 9;   // Wrist joint servo
const int SERVO_GRIPPER_PIN  = 10;  // Gripper servo

// Servo calibration offsets (used to align physical and logical angles)
const int BASE_OFFSET     = 90;  // Neutral offset for base servo
const int SHOULDER_OFFSET = 90;  // Neutral offset for shoulder servo
const int ELBOW_OFFSET    = 90;  // Neutral offset for elbow servo
const int WRIST_OFFSET    = 90;  // Neutral offset for wrist servo

// Predefined gripper positions (open and closed)
const int GRIPPER_CLOSED = 0;   // Gripper fully closed
const int GRIPPER_OPEN   = 90;  // Gripper fully open

// Safe workspace boundaries in X-Z plane (to avoid overextension)
const float MIN_X = 5.0;
const float MAX_X = 20.0;
const float MIN_Z = 0.0;
const float MAX_Z = 20.0;

// Safe servo angle limits (based on physical constraints of your servos)
const int BASE_MIN_ANGLE     = 0;
const int BASE_MAX_ANGLE     = 180;
const int SHOULDER_MIN_ANGLE = 15;
const int SHOULDER_MAX_ANGLE = 165;
const int ELBOW_MIN_ANGLE    = 10;
const int ELBOW_MAX_ANGLE    = 170;
const int WRIST_MIN_ANGLE    = 0;
const int WRIST_MAX_ANGLE    = 180;

// === Servo Objects ===
// These Servo objects are used to control the corresponding servo motors
Servo baseServo;
Servo shoulderServo;
Servo elbowServo;
Servo wristServo;
Servo gripperServo;

// Track last angles written to servos to avoid unnecessary updates
int lastBaseAngle     = -1;
int lastShoulderAngle = -1;
int lastElbowAngle    = -1;
int lastWristAngle    = -1;

// === Setup Function ===
void setup() {
  // Attach each servo to its respective pin
  baseServo.attach(SERVO_BASE_PIN);
  shoulderServo.attach(SERVO_SHOULDER_PIN);
  elbowServo.attach(SERVO_ELBOW_PIN);
  wristServo.attach(SERVO_WRIST_PIN);
  gripperServo.attach(SERVO_GRIPPER_PIN);

  // Start serial communication for debugging
  Serial.begin(9600);
}

// === Utility Functions ===

// Convert radians to servo angle, add offset, and constrain to [0,180]
int toServoAngle(float angle_rad, int offset) {
  return constrain(offset + (angle_rad * 180.0 / PI), 0, 180);
}

// Open the gripper to predefined open angle
void openGripper() {
  gripperServo.write(GRIPPER_OPEN);
  Serial.println("Gripper Opened");
}

// Close the gripper to predefined closed angle
void closeGripper() {
  gripperServo.write(GRIPPER_CLOSED);
  Serial.println("Gripper Closed");
}

// === Main Move Function ===
// Moves arm to a specified X, Z target with optional override for angles
// Returns true if the target is reachable and safe, false otherwise
bool moveArm(float X, float Z, float base_angle_deg, float gripper_pitch_deg, float overrideElbowRad = -1, float overrideShoulderRad = -1) {

  // === Safety checks ===

  // Ensure the X and Z target is within allowed workspace
  if (X < MIN_X || X > MAX_X || Z < MIN_Z || Z > MAX_Z) {
    Serial.println("ERROR: Target outside safe workspace!");
    return false;
  }

  // Calculate distance from shoulder to target point
  float D = sqrt(X * X + Z * Z);

  // Use cosine law to compute elbow angle
  float cos_angle2 = (D * D - L1 * L1 - L2 * L2) / (2 * L1 * L2);

  // If angle is not within valid range, arm cannot reach
  if (cos_angle2 < -1.0 || cos_angle2 > 1.0) {
    Serial.println("ERROR: Target unreachable (geometry)!");
    return false;
  }

  // === Inverse Kinematics ===

  // Compute θ2 (elbow) using acos, or override if specified
  float theta2 = (overrideElbowRad >= 0) ? overrideElbowRad : acos(cos_angle2);

  // Calculate intermediate values for shoulder angle
  float k1 = L1 + L2 * cos(theta2);
  float k2 = L2 * sin(theta2);
  float theta1 = atan2(Z, X) - atan2(k2, k1);  // Shoulder angle

  // Override shoulder angle if specified
  if (overrideShoulderRad >= 0) {
    theta1 = overrideShoulderRad;
  }

  // Calculate θ3 (wrist) to maintain gripper pitch
  float theta3 = (gripper_pitch_deg * PI / 180.0) - theta1 - theta2;

  // === Convert to Servo Angles ===

  // Convert calculated radian angles to servo degrees with offsets and constraints
  int baseAngle     = constrain(base_angle_deg, BASE_MIN_ANGLE, BASE_MAX_ANGLE);
  int shoulderAngle = constrain(toServoAngle(theta1, SHOULDER_OFFSET), SHOULDER_MIN_ANGLE, SHOULDER_MAX_ANGLE);
  int elbowAngle    = constrain(toServoAngle(theta2, ELBOW_OFFSET), ELBOW_MIN_ANGLE, ELBOW_MAX_ANGLE);
  int wristAngle    = constrain(toServoAngle(theta3, WRIST_OFFSET), WRIST_MIN_ANGLE, WRIST_MAX_ANGLE);

  // === Send angle only if changed ===

  if (lastBaseAngle != baseAngle) {
    baseServo.write(baseAngle);
    lastBaseAngle = baseAngle;
  }
  if (lastShoulderAngle != shoulderAngle) {
    shoulderServo.write(shoulderAngle);
    lastShoulderAngle = shoulderAngle;
  }
  if (lastElbowAngle != elbowAngle) {
    elbowServo.write(elbowAngle);
    lastElbowAngle = elbowAngle;
  }
  if (lastWristAngle != wristAngle) {
    wristServo.write(wristAngle);
    lastWristAngle = wristAngle;
  }

  // === Debug Output ===
  Serial.println("=== Move Arm ===");
  Serial.print("Base Angle: "); Serial.print(baseAngle); Serial.println("°");
  Serial.print("θ1 (Shoulder): "); Serial.print(theta1 * 180.0 / PI); Serial.println("°");
  Serial.print("θ2 (Elbow): "); Serial.print(theta2 * 180.0 / PI); Serial.println("°");
  Serial.print("θ3 (Wrist): "); Serial.print(theta3 * 180.0 / PI); Serial.println("°");
  Serial.print("Gripper Angle: "); Serial.print(gripper_pitch_deg); Serial.println("°");
  Serial.println("-------------------------");

  return true;
}

// === Smooth Movement Function ===
// Moves arm gradually in small steps for smooth motion
void smoothMoveArm(float X, float Z, float base_angle_deg, float gripper_pitch_deg, float overrideElbowRad = -1, float overrideShoulderRad = -1, int steps = 10, int delay_per_step = 50) {
  for (int i = 1; i <= steps; i++) {
    float factor = (float)i / steps;  // Linear interpolation factor
    moveArm(X * factor, Z * factor, base_angle_deg, gripper_pitch_deg, overrideElbowRad, overrideShoulderRad);
    delay(delay_per_step);  // Pause between steps
  }
}

// === Example Pick and Place Task ===
void performPickAndPlace() {
  
  // Move to pickup location (step 1)
  moveArm(10.0, 10.0, 90, -45, PI / 4, 0);
  delay(2500);  // Wait 2.5 seconds

  // Move lower to pick object (step 2)
  moveArm(15.0, 5.0, 90, 105, PI / 3, PI / 2);
  delay(2500);

  // Close the gripper to grab the object
  closeGripper();
  delay(2500);
  
  // Lift arm up with object (step 3)
  moveArm(10.0, 10.0, 180, -45, PI / 4, 0);
  delay(2500);
  
  // Move to place location (step 4)
  moveArm(15.0, 5.0, 180, 105, PI / 3, PI / 2);
  delay(2500);
  
  // Release the object
  openGripper();
  delay(2500);
}

// === Main Loop ===
void loop() {
  performPickAndPlace();  // Repeat pick-and-place sequence
}