#include "Behaviors/StateMachine.h"
#include "Behaviors/Demos/Navigation/PilotDemo.h"

/** --------------------- Kinematic Constants ----------------------- **/
#define SLOW_SERVO_SPEED  0.5
#define DEFAULT_JOINT_ANG	0
#define LGRIP_CLOSED			0.3
#define RGRIP_CLOSED			-0.3
/** --------------------- World Map Constants ----------------------- **/
#define TAG_HEIGHT	200

$nodeclass SnackBot : PilotDemo {
  $provide Point nextFoodLoc;
	
	/** ----------------------- Localization -------------------------- **/

  /** The map of the robot's environment, used for localizing. Contains all of
      the locations of the april tags.
  **/
  virtual void buildMap() {
    cout << "Building map..." << endl;
		
    // The points which form the playpen polygon
    std::vector<Point> penPoints;		
    penPoints.push_back(Point(-391, -391, 0, allocentric));
    penPoints.push_back(Point(2321, -391, 0, allocentric));
    penPoints.push_back(Point(2321, 376, 0, allocentric));
    penPoints.push_back(Point(1250, 376, 0, allocentric));				
    penPoints.push_back(Point(1250, 1281, 0, allocentric));
    penPoints.push_back(Point(750, 1281, 0, allocentric));
    penPoints.push_back(Point(750, 376, 0, allocentric));
    penPoints.push_back(Point(375, 376, 0, allocentric));
    penPoints.push_back(Point(375, 1281, 0, allocentric));
    penPoints.push_back(Point(-391, 1281, 0, allocentric));
    NEW_SHAPE(worldBounds, PolygonData, new PolygonData(worldShS, penPoints, true));
    worldBounds->setObstacle(false);
    // Position in front of the food
    NEW_SHAPE(foodLocation, PointData, new PointData(worldShS, Point(1000, 0, 0, allocentric)));
		foodLocation->setObstacle(false);
    // Position in front of the cobot
    NEW_SHAPE(cobotLocation, PointData, new PointData(worldShS, Point(0, 0, 0, allocentric)));
		cobotLocation->setObstacle(false);
    // April tag locations
    NEW_SHAPE(tag18, AprilTagData,
	      new AprilTagData(worldShS, Point(1000, 1280, TAG_HEIGHT, allocentric),
			       AprilTags::TagDetection(18)));
    NEW_SHAPE(tag2, AprilTagData,
	      new AprilTagData(worldShS, Point(1000, -390, TAG_HEIGHT, allocentric),
			       AprilTags::TagDetection(2)));
    NEW_SHAPE(tag12, AprilTagData,
	      new AprilTagData(worldShS, Point(550,  375, TAG_HEIGHT, allocentric),
			       AprilTags::TagDetection(12)));
    NEW_SHAPE(tag3, AprilTagData,
	      new AprilTagData(worldShS, Point(-250, -390, TAG_HEIGHT, allocentric),
			       AprilTags::TagDetection(3)));
    NEW_SHAPE(tag0, AprilTagData,
	      new AprilTagData(worldShS, Point(2320, -80, TAG_HEIGHT, allocentric),
			       AprilTags::TagDetection(0)));
    NEW_SHAPE(tag4, AprilTagData,
	      new AprilTagData(worldShS, Point(-375, -80, TAG_HEIGHT, allocentric),
			       AprilTags::TagDetection(4)));
    NEW_SHAPE(tag17, AprilTagData,
	      new AprilTagData(worldShS, Point(1450, 375, TAG_HEIGHT, allocentric),
			       AprilTags::TagDetection(17)));
    NEW_SHAPE(tag1, AprilTagData,
	      new AprilTagData(worldShS, Point(2000, -390, TAG_HEIGHT, allocentric),
			       AprilTags::TagDetection(1)));
    NEW_SHAPE(tag7, AprilTagData,
	      new AprilTagData(worldShS, Point(0, 1280, TAG_HEIGHT, allocentric),
			       AprilTags::TagDetection(7)));
    NEW_SHAPE(tag5, AprilTagData,
	      new AprilTagData(worldShS, Point(-375, 640, TAG_HEIGHT, allocentric),
			       AprilTags::TagDetection(5)));
    NEW_SHAPE(tag14, AprilTagData,
	      new AprilTagData(worldShS, Point(1265, 640, TAG_HEIGHT, allocentric),
			       AprilTags::TagDetection(14)));
  }

  /** State machine which looks for multiple april tags to localize against **/
  $nodeclass LocalizeRobot : StateNode {
		$setupmachine {
      LookUp =C=> Localize =C=>
				LookAtPoint(1000, -250, TAG_HEIGHT) =C=> Localize =C=>
				LookAtPoint(1000, 250, TAG_HEIGHT) =C=> Localize =C=> PostMachineCompletion
		}
  }

	/** ----------------------- Create Movement -------------------------- **/

  /** State machine which moves the robot to where the food is located **/
  $nodeclass MoveRobotToFoodLoc : StateNode {
    $nodeclass MoveToFoodLoc : PilotNode(goToShape) : doStart {
			// Replan if you collide with something
			pilotreq.collisionAction = collisionReplan;
			// Move to to where the food is
      GET_SHAPE(foodLocation, PointData, worldShS);
      pilotreq.targetShape = foodLocation;
    }
    $setupmachine {
      LookUp =C=> RetractArm =C=> MoveToFoodLoc =C=> FaceTowardsY =C=>
				PostMachineCompletion
		}
  }

  /** State machine which moves the robot to where the cobot is located **/
  $nodeclass MoveRobotToCobot : StateNode {
    $nodeclass MoveToCobot : PilotNode(goToShape) : doStart {
			// Move to to where the cobot is
      GET_SHAPE(cobotLocation, PointData, worldShS);
      pilotreq.targetShape = cobotLocation;
			pilotreq.collisionAction = collisionReplan;
    }
    $setupmachine {
      LookUp =C=> RetractArm =C=> MoveToCobot =C=> PostMachineCompletion
		}
  }

	/** Turn the robot towards the Y-axis **/
  $nodeclass FaceTowardsY : PilotNode(walk) : doStart {
    pilotreq.da = deg2rad(90.0) - theAgent->getOrientation();
  }

	/** Given the location of a food object, this state machine moves the
      robot into position so it can grab the object
  **/
  $nodeclass PositionRobotNextToFood : StateNode {
    /** Given the location of a food object, turns the robot towards it **/
    $nodeclass TurnRobotTowardsFood : PilotNode(walk) : doStart {
      $reference SnackBot::nextFoodLoc;
      pilotreq.da = atan(nextFoodLoc.coordY() / nextFoodLoc.coordX());
    }		

    /** Given the location of a food object, moves the robot towards it **/
    $nodeclass MoveRobotTowardsFood : PilotNode(walk) : doStart {
      $reference SnackBot::nextFoodLoc;
      float distance = nextFoodLoc.xyNorm();
      // Move the robot so that the food is 300 mm in front of it
      pilotreq.dx = distance - 300;
    }

    /** Reposition the location of the food object in terms on the new
				robot location.
    **/
    $nodeclass RelocateFoodObject : StateNode : doStart {
      $reference SnackBot::nextFoodLoc;
      nextFoodLoc = Point(300, 0);
    }

    $setupmachine {
			startnode : StateNode
			foodLoc : FindFoodLocation			

      startnode =N=> RetractArm =C=> TurnRobotTowardsFood =C=> MoveRobotTowardsFood
				=C=> LookAtFood =C=> FindFood =C=> foodLoc

			foodLoc =S=> PostMachineCompletion
			foodLoc =F=> RelocateFoodObject
		}
  }

  $nodeclass PositionRobotNearCobot : StateNode {
		$nodeclass TurnRobot : PilotNode(walk) : doStart {
      pilotreq.da = deg2rad(0.0);
    }

		$setupmachine {
			FaceTowardsY =C=> PostMachineCompletion
		}
	}

	/** ----------------------- Camera Movement -------------------------- **/

  /** Reposition the head so that the camera is looking at the specified point **/
  $nodeclass LookAtPoint(float x, float y, float z) : HeadPointerNode : doStart {
    // Slow the head down
    for (unsigned int i = 0; i < NumHeadJoints; i++) {
			getMC()->setMaxSpeed(i, 0.5);
    }
		// Set the pan value
  	getMC()->lookAtPoint(x,y,z);
  }

  /** Point the head towards the ground **/
  $nodeclass LookAtGround : HeadPointerNode : doStart {
    // Slow the head down
    for (unsigned int i = 0; i < NumHeadJoints; i++) {
			getMC()->setMaxSpeed(i, SLOW_SERVO_SPEED);
    }
		// Look at the ground
    getMC()->lookAtPoint(500, 0, 0);
	}

	/** Reposition the head so that the camera is just about horizontal
			and looking straight ahead
	**/
  $nodeclass LookUp : HeadPointerNode : doStart {
    // Slow the head down
    for (unsigned int i = 0; i < NumHeadJoints; i++) {
			getMC()->setMaxSpeed(i, 0.5);
    }
		// Set the pan value
  	getMC()->setJointValue(0, 0);
		// Set the tilt value
  	getMC()->setJointValue(1, -0.3);			
  }

	/** Reposition the head so that the camera is pointing towards the food
	**/
  $nodeclass LookAtFood : HeadPointerNode : doStart {
    // Slow the head down
    for (unsigned int i = 0; i < NumHeadJoints; i++) {
			getMC()->setMaxSpeed(i, 0.5);
    }
		// Set the pan value
  	getMC()->lookAtPoint(300, 0, 0);	
  }

	/** ----------------------- Arm Movement -------------------------- **/

	/** Moves the arm to a retracted position with its gripper closed **/	
  $nodeclass RetractArm : StateNode {
		$nodeclass Pose1 : PostureNode : doStart {
   		// Slow the arm down, set default joint values
   	  for (unsigned int i = 0; i < NumArmJoints; i++) {
      	getMC()->setMaxSpeed(ArmBaseOffset + i, SLOW_SERVO_SPEED);
    	}
    	// Retract the arm
    	getMC()->setOutputCmd(ArmShoulderOffset, 2.5);
    	getMC()->setOutputCmd(ArmElbowOffset, -2.5);
    	getMC()->setOutputCmd(ArmWristOffset, 2.2);
			// Close the gripper	
    	getMC()->setOutputCmd(LeftFingerOffset, LGRIP_CLOSED);
    	getMC()->setOutputCmd(RightFingerOffset, RGRIP_CLOSED);
  	}		

		$nodeclass Pose2 : PostureNode : doStart {
   		// Slow the arm down, set default joint values
   	  for (unsigned int i = 0; i < NumArmJoints; i++) {
      	getMC()->setMaxSpeed(ArmBaseOffset + i, SLOW_SERVO_SPEED);
     	 	getMC()->setOutputCmd(ArmBaseOffset + i, DEFAULT_JOINT_ANG);
    	}
    	// Turn the arm
    	getMC()->setOutputCmd(ArmBaseOffset, 1.5);
    	getMC()->setOutputCmd(ArmShoulderOffset, 2.5);
    	getMC()->setOutputCmd(ArmElbowOffset, -2.5);
    	getMC()->setOutputCmd(ArmWristOffset, 2.2);
			// Close the gripper	
    	getMC()->setOutputCmd(LeftFingerOffset, LGRIP_CLOSED);
    	getMC()->setOutputCmd(RightFingerOffset, RGRIP_CLOSED);
  	}

		$setupmachine {
			Pose1 =C=> Pose2 =C=> PostMachineCompletion
		}
	}

  /** Move the arm to grab the food object. First postions the arm above the food
      with the correct orientation. Then lowers the arm to slightly above the food
      object. Then grips the food object **/ 
  $nodeclass GrabFood : DynamicMotionSequenceNode : doStart {
    $reference SnackBot::nextFoodLoc;
    // Creates a prep and target position for the arm. The prep position moves the
    // arm above the food and points it downward. The target position is the same,
    // except that it is close enough to the food to actually grip it.
    float xPos = nextFoodLoc.coordX();
    float yPos = nextFoodLoc.coordY();
    float zPrepPos = 200;
    float zTgtPos = 30;
    fmat::Column<3> prepPos = fmat::pack(xPos, yPos, zPrepPos);
    fmat::Column<3> targetPos = fmat::pack(xPos, yPos, zTgtPos);
    fmat::Column<3> offset = fmat::pack(0, -50, 0);
    fmat::Quaternion oriTgt = fmat::Quaternion::aboutY(-M_PI/2);		
    fmat::Quaternion oriOffset = fmat::Quaternion::identity();
    // Create a posture engine for computing the postures
    PostureEngine* pe = new PostureEngine();
    // Solve for the prep position
    bool solved = pe->solveLink(prepPos, oriTgt, GripperFrameOffset, offset, oriOffset);

    cout << "Prep position solved: " << solved << endl;

    pe->setOutputCmd(LeftFingerOffset, -1);
    pe->setOutputCmd(RightFingerOffset, 1);
    // Add this to the motion sequence
    getMC()->advanceTime(3000);
    getMC()->setPose(*pe);
    // Solve for the target position
    solved = pe->solveLink(targetPos, oriTgt, GripperFrameOffset, offset, oriOffset);

    cout << "Target position solved: " << solved << endl;

    pe->setOutputCmd(LeftFingerOffset, -1);
    pe->setOutputCmd(RightFingerOffset, 1);	
    // Add this to the motion sequence
    getMC()->advanceTime(2000);
    getMC()->setPose(*pe);
    // Close the fingers
    pe->setOutputCmd(LeftFingerOffset, 0.3);
    pe->setOutputCmd(RightFingerOffset, -0.3);
    // Add this to the motion sequence
    getMC()->advanceTime(1000);
    getMC()->setPose(*pe);
    // Move back to the prep position with the fingers closed
    solved = pe->solveLink(prepPos, oriTgt, GripperFrameOffset, offset, oriOffset);

    cout << "Pick up position solved: " << solved << endl;

    pe->setOutputCmd(LeftFingerOffset, 0.3);
    pe->setOutputCmd(RightFingerOffset, -0.3);
    // Add this to the motion sequence
    getMC()->advanceTime(1500);
    getMC()->setPose(*pe);
    // Free memory
    delete(pe);
  }

  $nodeclass GiveFoodToCobot : StateNode {

    $nodeclass GripFood : PostureNode : doStart {
      // Slow the arm down
      for (unsigned int i = 0; i < NumArmJoints; i++) {
				getMC()->setMaxSpeed(ArmBaseOffset + i, 0.5);
      }
      // Set joint values
      getMC()->setOutputCmd(LeftFingerOffset, 0.3);
      getMC()->setOutputCmd(RightFingerOffset, -0.3);
    }

		$nodeclass Prep1 : PostureNode : doStart {
      // Slow the arm down, set default joint values
      for (unsigned int i = 0; i < NumArmJoints; i++) {
				getMC()->setMaxSpeed(ArmBaseOffset + i, 0.5);
				getMC()->setOutputCmd(ArmBaseOffset + i, 0);
      }
      // Set joint values
  	  getMC()->setOutputCmd(ArmBaseOffset, 1.5);
    	getMC()->setOutputCmd(ArmShoulderOffset, 2.5);
    	getMC()->setOutputCmd(ArmElbowOffset, -2.5);
    	getMC()->setOutputCmd(ArmWristOffset, 0);
      getMC()->setOutputCmd(WristRotateOffset, -1.5);
      getMC()->setOutputCmd(LeftFingerOffset, 0.3);
      getMC()->setOutputCmd(RightFingerOffset, -0.3);
    }

		$nodeclass Prep2 : PostureNode : doStart {
      // Slow the arm down, set default joint values
      for (unsigned int i = 0; i < NumArmJoints; i++) {
				getMC()->setMaxSpeed(ArmBaseOffset + i, 0.5);
				getMC()->setOutputCmd(ArmBaseOffset + i, 0);
      }
      // Set joint values
      getMC()->setOutputCmd(ArmBaseOffset, 1);
      getMC()->setOutputCmd(ArmShoulderOffset, 2.5);
      getMC()->setOutputCmd(ArmElbowOffset, -1.25);
      getMC()->setOutputCmd(ArmWristOffset, 0);
      getMC()->setOutputCmd(WristRotateOffset, -1.5);
      getMC()->setOutputCmd(LeftFingerOffset, 0.3);
      getMC()->setOutputCmd(RightFingerOffset, -0.3);
    }

    $nodeclass Prep3 : PostureNode : doStart {
      // Slow the arm down, set default joint values
      for (unsigned int i = 0; i < NumArmJoints; i++) {
				getMC()->setMaxSpeed(ArmBaseOffset + i, 0.5);
				getMC()->setOutputCmd(ArmBaseOffset + i, 0);
      }
      // Set joint values
      getMC()->setOutputCmd(ArmBaseOffset, 0);
      getMC()->setOutputCmd(ArmShoulderOffset, 1.45);
      getMC()->setOutputCmd(ArmElbowOffset, 0.2);
      getMC()->setOutputCmd(ArmWristOffset, 0.3);
      getMC()->setOutputCmd(WristRotateOffset, -1.5);
      getMC()->setOutputCmd(LeftFingerOffset, 0.3);
      getMC()->setOutputCmd(RightFingerOffset, -0.3);
    }

    $nodeclass TossFood : DynamicMotionSequenceNode : doStart {
      // Create a new posture engine for calculating joint positions
      PostureEngine* pe = new PostureEngine();
      // Set the initial position
      getMC()->setOutputCmd(ArmBaseOffset, 0);
      getMC()->setOutputCmd(ArmShoulderOffset, 1.45);
      getMC()->setOutputCmd(ArmElbowOffset, 0.2);
      getMC()->setOutputCmd(ArmWristOffset, 0.3);
      getMC()->setOutputCmd(WristRotateOffset, -1.5);
      getMC()->setOutputCmd(LeftFingerOffset, 0.3);
      getMC()->setOutputCmd(RightFingerOffset, -0.3);
      getMC()->setPose(*pe);
      // Set the first position
      getMC()->setOutputCmd(ArmShoulderOffset, 1.4);
      getMC()->setOutputCmd(ArmElbowOffset, 0);
      getMC()->setOutputCmd(ArmWristOffset, 0);
      // Move to the first position
      getMC()->advanceTime(50);
      getMC()->setPose(*pe);
      // Set the second position
      getMC()->setOutputCmd(ArmShoulderOffset, 1.35);
      getMC()->setOutputCmd(ArmElbowOffset, -0.1);
      getMC()->setOutputCmd(ArmWristOffset, -0.1);
      // Move to the second position
      getMC()->advanceTime(110);
      getMC()->setPose(*pe);		
      // Set the third position
      getMC()->setOutputCmd(LeftFingerOffset, -0.25);
      getMC()->setOutputCmd(RightFingerOffset, 0.25);
      // Move to the third position
      getMC()->advanceTime(160);
      getMC()->setPose(*pe);
      // Set the fourth position			
      getMC()->setOutputCmd(LeftFingerOffset, -0.5);
      getMC()->setOutputCmd(RightFingerOffset, 0.5);
      // Move to the fourth position
      getMC()->advanceTime(100);
      getMC()->setPose(*pe);
      // Free memory
      delete(pe);		
    }
	
    $setupmachine {
      GripFood =C=> RetractArm =C=> Prep1 =C=> Prep2 =C=> Prep3 =T(5000)=> TossFood =T(1000)=> RetractArm
		}
  }

	/** ----------------------- Vision -------------------------- **/
	
  $nodeclass FindFood : MapBuilderNode(MapBuilderRequest::localMap) : doStart {
    mapreq.addObjectColor(blobDataType, "blue");
		// Filter out noise
		mapreq.addMinBlobArea("blue", 50);
  }

  /** State machine which visually searches the food area and builds a map
      of the location of food objects
  **/
  $nodeclass SweepHeadForFood : VisualRoutinesStateNode {
    $setupmachine {
      LookAtGround =C=> FindFood =C=> PostMachineCompletion
		}
  }

	/** Gets the position of a food object that the robot wants to grab. **/
  $nodeclass FindFoodLocation : VisualRoutinesStateNode : doStart {
    // Get all of the potential food blobs		
    NEW_SHAPEVEC(blob_shapes, BlobData, select_type<BlobData>(localShS));
    // Check that the vector isn't empty, if so report failure
    if (blob_shapes.size() == 0) {
      cout << "Didn't find any food objects" << endl;
      postStateFailure();
    } else {
      // Pick the biggest blob in the vector and get its location
      Shape<BlobData> blob = blob_shapes[0];
      float biggestArea = blob->getArea();
      for (unsigned int i = 1; i < blob_shapes.size(); ++i) {
				// Get the distance of the next blob
				float nextArea = blob_shapes[i]->getArea();
				// Compare it to the largest known area
				if (nextArea > biggestArea) {
	  			biggestArea = nextArea;
	  			blob = blob_shapes[i];
				}
			}
    	$reference SnackBot::nextFoodLoc;
    	nextFoodLoc = blob->getCentroid();

	    cout << "Food located at: " << nextFoodLoc << endl;

  	  postStateSuccess();

			NEW_SHAPE(foodLocationPointer, PointData, new PointData(localShS, nextFoodLoc));
 	 	}
  }

	/** ----------------------- State Machine -------------------------- **/	

  /** Performs start operations for the robot **/
  $nodeclass StartNode : StateNode : doStart {
		particleFilter->resetFilter();
  }

  $setupmachine {
		startdemo : StartNode
		foodLoc : FindFoodLocation


    startdemo =N=> ResetArm =C=> LocalizeRobot =C=> MoveRobotToFoodLoc =C=>
			SweepHeadForFood =C=> foodLoc

    foodLoc =F=> StateNode
    foodLoc =S=> PositionRobotNextToFood =C=>
			GrabFood =C=> RetractArm =C=> ResetGripper =C=> RetractArm =C=>
			MoveRobotToFoodLoc =C=> LocalizeRobot =C=> MoveRobotToCobot =C=>
			PositionRobotNearCobot =C=> GiveFoodToCobot
	}
}

REGISTER_BEHAVIOR(SnackBot);