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

// Create servo objects
Servo servo_waist;
Servo servo_shoulder;
Servo servo_elbow;
Servo servo_wristpitch;
Servo servo_wristroll;
Servo servo_claw;

// Robot arm dimensions (in cm)
const float L1 = 10.0;  // Base to shoulder height
const float L2 = 12.0;  // Upper arm length
const float L3 = 12.0;  // Forearm length
const float L4 = 8.0;   // Wrist to end-effector length

// Box dimensions
const float BOX_SIZE = 3.0;  // 3cm x 3cm x 3cm
const float SQUARE_SIZE = 5.0;  // 5cm x 5cm
const float GRIPPER_OFFSET = 1.5;  // Half of box size for center grip

// Structure to store position and orientation
struct Pose {
    float x, y, z;      // Position
    float roll, pitch, yaw;  // Orientation
};

// Structure to store joint angles
struct JointAngles {
    float theta1;  // Waist
    float theta2;  // Shoulder
    float theta3;  // Elbow
    float theta4;  // Wrist pitch
    float theta5;  // Wrist roll
};

// Structure to store task configuration
struct TaskConfig {
    Pose current_pose;
    JointAngles joint_angles;
    String phase;
    unsigned long timestamp;
};

// Function to calculate inverse kinematics (same as before)
JointAngles inverseKinematics(Pose target) {
    JointAngles angles;
    
    // Calculate waist angle (theta1)
    angles.theta1 = atan2(target.y, target.x);
    
    // Calculate wrist position
    float wrist_x = target.x - L4 * cos(target.pitch) * cos(angles.theta1);
    float wrist_y = target.y - L4 * cos(target.pitch) * sin(angles.theta1);
    float wrist_z = target.z - L4 * sin(target.pitch);
    
    // Calculate distance to wrist
    float r = sqrt(wrist_x*wrist_x + wrist_y*wrist_y);
    float s = wrist_z - L1;
    float D = sqrt(r*r + s*s);
    
    // Calculate shoulder and elbow angles using cosine law
    float cos_theta3 = (D*D - L2*L2 - L3*L3)/(2*L2*L3);
    angles.theta3 = acos(cos_theta3);
    
    float beta = atan2(s, r);
    float alpha = acos((L2*L2 + D*D - L3*L3)/(2*L2*D));
    angles.theta2 = beta + alpha;
    
    // Calculate wrist angles
    angles.theta4 = target.pitch - (angles.theta2 + angles.theta3);
    angles.theta5 = target.roll;
    
    // Convert angles to degrees
    angles.theta1 = degrees(angles.theta1);
    angles.theta2 = degrees(angles.theta2);
    angles.theta3 = degrees(angles.theta3);
    angles.theta4 = degrees(angles.theta4);
    angles.theta5 = degrees(angles.theta5);
    
    return angles;
}

// Function to log configuration
void logConfiguration(TaskConfig config) {
    Serial.println("\n=== " + config.phase + " ===");
    Serial.println("Timestamp: " + String(config.timestamp) + " ms");
    
    Serial.println("\nEnd-Effector Position:");
    Serial.println("X: " + String(config.current_pose.x) + " cm");
    Serial.println("Y: " + String(config.current_pose.y) + " cm");
    Serial.println("Z: " + String(config.current_pose.z) + " cm");
    Serial.println("Roll: " + String(config.current_pose.roll) + " degrees");
    Serial.println("Pitch: " + String(config.current_pose.pitch) + " degrees");
    Serial.println("Yaw: " + String(config.current_pose.yaw) + " degrees");
    
    Serial.println("\nJoint Angles:");
    Serial.println("Waist: " + String(config.joint_angles.theta1) + " degrees");
    Serial.println("Shoulder: " + String(config.joint_angles.theta2) + " degrees");
    Serial.println("Elbow: " + String(config.joint_angles.theta3) + " degrees");
    Serial.println("Wrist Pitch: " + String(config.joint_angles.theta4) + " degrees");
    Serial.println("Wrist Roll: " + String(config.joint_angles.theta5) + " degrees");
    Serial.println("=====================================");
}

// Function to move servos to target angles
void moveToAngles(JointAngles angles) {
    // Map angles to servo values (0-180)
    int servo_waist_val = constrain(map(angles.theta1, -90, 90, 0, 180), 0, 180);
    int servo_shoulder_val = constrain(map(angles.theta2, -90, 90, 0, 180), 0, 180);
    int servo_elbow_val = constrain(map(angles.theta3, -90, 90, 0, 180), 0, 180);
    int servo_wristpitch_val = constrain(map(angles.theta4, -90, 90, 0, 180), 0, 180);
    int servo_wristroll_val = constrain(map(angles.theta5, -90, 90, 0, 180), 0, 180);
    
    // Write values to servos
    servo_waist.write(servo_waist_val);
    servo_shoulder.write(servo_shoulder_val);
    servo_elbow.write(servo_elbow_val);
    servo_wristpitch.write(servo_wristpitch_val);
    servo_wristroll.write(servo_wristroll_val);
}

// Function to perform color box pick and place operation
void pickAndPlaceColorBox(Pose box_pose, Pose target_pose) {
    TaskConfig config;
    const float APPROACH_HEIGHT = 5.0;  // Height above box/target for approach
    
    // 1. Initial position (Home position)
    config.phase = "Initial Configuration";
    config.current_pose = {20.0, 0.0, 20.0, 0.0, -90.0, 0.0};
    config.joint_angles = inverseKinematics(config.current_pose);
    config.timestamp = millis();
    moveToAngles(config.joint_angles);
    logConfiguration(config);
    delay(1000);
    
    // 2. Approach box position
    Pose approach_pose = box_pose;
    approach_pose.z += APPROACH_HEIGHT;
    config.phase = "Box Approach";
    config.current_pose = approach_pose;
    config.joint_angles = inverseKinematics(approach_pose);
    config.timestamp = millis();
    moveToAngles(config.joint_angles);
    logConfiguration(config);
    delay(1000);
    
    // 3. Pick position
    config.phase = "Box Pick Position";
    config.current_pose = box_pose;
    config.joint_angles = inverseKinematics(box_pose);
    config.timestamp = millis();
    moveToAngles(config.joint_angles);
    logConfiguration(config);
    delay(500);
    
    // 4. Close gripper
    servo_claw.write(180);  // Close gripper
    delay(500);
    
    // 5. Lift box
    config.phase = "Box Lift";
    config.current_pose = approach_pose;
    config.joint_angles = inverseKinematics(approach_pose);
    config.timestamp = millis();
    moveToAngles(config.joint_angles);
    logConfiguration(config);
    delay(1000);
    
    // 6. Approach target position
    Pose target_approach = target_pose;
    target_approach.z += APPROACH_HEIGHT;
    config.phase = "Target Approach";
    config.current_pose = target_approach;
    config.joint_angles = inverseKinematics(target_approach);
    config.timestamp = millis();
    moveToAngles(config.joint_angles);
    logConfiguration(config);
    delay(1000);
    
    // 7. Place position
    config.phase = "Box Place Position";
    config.current_pose = target_pose;
    config.joint_angles = inverseKinematics(target_pose);
    config.timestamp = millis();
    moveToAngles(config.joint_angles);
    logConfiguration(config);
    delay(500);
    
    // 8. Open gripper
    servo_claw.write(0);  // Open gripper
    delay(500);
    
    // 9. Final position
    config.phase = "Final Configuration";
    config.current_pose = target_approach;
    config.joint_angles = inverseKinematics(target_approach);
    config.timestamp = millis();
    moveToAngles(config.joint_angles);
    logConfiguration(config);
    delay(1000);
}

void setup() {
    // Attach servos to pins
    servo_waist.attach(22);
    servo_shoulder.attach(23);
    servo_elbow.attach(24);
    servo_wristpitch.attach(25);
    servo_wristroll.attach(26);
    servo_claw.attach(27);
    
    Serial.begin(9600);
    delay(2000);  // Wait for serial connection
    
    // Example coordinates for box pickup and placement
    Pose box_pose = {
        20.0,   // x (cm)
        15.0,   // y (cm)
        0.0,    // z (cm)
        0.0,    // roll (degrees)
        -90.0,  // pitch (degrees)
        0.0     // yaw (degrees)
    };
    
    Pose target_pose = {
        20.0,   // x (cm)
        -15.0,  // y (cm)
        0.0,    // z (cm)
        0.0,    // roll (degrees)
        -90.0,  // pitch (degrees)
        0.0     // yaw (degrees)
    };
    
    // Perform pick and place operation
    pickAndPlaceColorBox(box_pose, target_pose);
}

void loop() {
    // Task runs once in setup()
    delay(1000);
}