#include <Servo.h>
#include <math.h>

// === Configuration Constants ===

// Arm link lengths in cm
const float L1 = 10.0;  // Shoulder to elbow
const float L2 = 10.0;  // Elbow to wrist
const float L3 = 5.0;   // Wrist to gripper (unused)

// Servo pins
const int SERVO_BASE_PIN     = 3;
const int SERVO_SHOULDER_PIN = 5;
const int SERVO_ELBOW_PIN    = 6;
const int SERVO_WRIST_PIN    = 9;
const int SERVO_GRIPPER_PIN  = 10;

// Servo calibration offsets
const int BASE_OFFSET     = 90;
const int SHOULDER_OFFSET = 90;
const int ELBOW_OFFSET    = 90;
const int WRIST_OFFSET    = 90;

// Gripper positions
const int GRIPPER_CLOSED = 0;
const int GRIPPER_OPEN   = 90;

// Safe workspace boundaries (adjust for your robot)
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 (adjust for your robot)
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 ===
Servo baseServo;
Servo shoulderServo;
Servo elbowServo;
Servo wristServo;
Servo gripperServo;

// Last written angles (optimization)
int lastBaseAngle     = -1;
int lastShoulderAngle = -1;
int lastElbowAngle    = -1;
int lastWristAngle    = -1;

// === Setup ===
void setup() {
  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);

  Serial.begin(9600);
}

// === Utility Functions ===
int toServoAngle(float angle_rad, int offset) {
  return constrain(offset + (angle_rad * 180.0 / PI), 0, 180);
}

void openGripper() {
  gripperServo.write(GRIPPER_OPEN);
  Serial.println("Gripper Opened");
}

void closeGripper() {
  gripperServo.write(GRIPPER_CLOSED);
  Serial.println("Gripper Closed");
}

// === Main Move Function ===
// Returns true if target 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 ===

  // Workspace boundary check
  if (X < MIN_X || X > MAX_X || Z < MIN_Z || Z > MAX_Z) {
    Serial.println("ERROR: Target outside safe workspace!");
    return false;
  }

  // Inverse kinematics check
  float D = sqrt(X * X + Z * Z);
  float cos_angle2 = (D * D - L1 * L1 - L2 * L2) / (2 * L1 * L2);

  if (cos_angle2 < -1.0 || cos_angle2 > 1.0) {
    Serial.println("ERROR: Target unreachable (geometry)!");
    return false;
  }

  // === Inverse kinematics ===

  float theta2 = (overrideElbowRad >= 0) ? overrideElbowRad : acos(cos_angle2);

  float k1 = L1 + L2 * cos(theta2);
  float k2 = L2 * sin(theta2);
  float theta1 = atan2(Z, X) - atan2(k2, k1);

  if (overrideShoulderRad >= 0) {
    theta1 = overrideShoulderRad;
  }

  float theta3 = (gripper_pitch_deg * PI / 180.0) - theta1 - theta2;

  // === Convert to servo angles ===

  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);

  // === Write to servos if angle 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 Move (optional improvement) ===
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;
    moveArm(X * factor, Z * factor, base_angle_deg, gripper_pitch_deg, overrideElbowRad, overrideShoulderRad);
    delay(delay_per_step);
  }
}

// === Example Sequence ===
void performPickAndPlace() {
  
  // Move 1
  moveArm(10.0, 10.0, 90, 0, PI / 3, 0);
  delay(2500);
  
  // Move 2
  moveArm(15.0, 5.0, 90, 10, 0, PI / 4);
  delay(2500);
  
  // Close gripper
  closeGripper();
  delay(2500);
  
  // Move 3
  moveArm(10.0, 10.0, 90, 0, PI / 3, 0);
  delay(2500);
  
  // Move 4
  moveArm(10.0, 10.0, 180, 0, PI / 3, 0);
  delay(2500);
  
  // Move 5
  moveArm(15.0, 5.0, 180, 10, 0, PI / 4);
  delay(2500);
  
  // Open gripper
  openGripper();
  delay(2500);
}

// === Loop ===
void loop() {
  performPickAndPlace();
}