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

$nodeclass SnackBot : PilotDemo {
	$provide Point nextFoodLoc;

	/** TODO:
				1) Use vision to identify and categorize food objects.
				2) Make a node that moves the robot to a position where it can grab
					 the food.
				3) Be able to identify the Cobot's location and move the robot to a
					 location in front of the Cobot where it can put the food in the
					 basket.
				4) Put the food in the Cobot's basket (may have to toss)
	**/

	/** Identify all of the potential food objects of the requested type.
			Currently this node assumes that food objects are green squares.
			Eventually we want to use vision to identify and distinguish between
			food objects.
	**/

	/** Move the arm to a retracted position so that the arm doesn't bump into
			anything as the arm moves
	**/
	$nodeclass RetractArm : 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.9);
		getMC()->setOutputCmd(ArmElbowOffset, -1.9);
		getMC()->setOutputCmd(ArmWristOffset, 1.9);	
	}

	$nodeclass FaceTowardsY : PilotNode(walk) : doStart {
		pilotreq.da = deg2rad(90.0) - theAgent->getOrientation();
	}

	$nodeclass FindFood : MapBuilderNode(MapBuilderRequest::localMap) : doStart {
		mapreq.addObjectColor(blobDataType, "blue");
	}

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

	/** Performs start operations for the robot
	**/
	$nodeclass StartNode : StateNode : doStart {
	}

	/** State machine which moves the robot to where the food is located
	**/
	$nodeclass MoveRobotToFoodLoc : StateNode {
		$nodeclass MoveToFoodLoc : PilotNode(localize) : doStart {
			GET_SHAPE(tag12, AprilTagData, worldShS);	
			//GET_SHAPE(foodLocation, PointData, worldShS);
			//pilotreq.targetShape = foodLocation;
			pilotreq.landmarks.clear();
			pilotreq.landmarks.push_back(tag12);
		}

		$setupmachine {
			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 {	
			GET_SHAPE(cobotLocation, PointData, worldShS);
			pilotreq.targetShape = cobotLocation;
		}

		$setupmachine {
			RetractArm =C=> MoveToCobot =C=> FaceTowardsY =C=> PostMachineCompletion
		}
	}

	/** State machine which visually searches the food area and builds a map
			of the location of food objects
	**/
	$nodeclass SweepHeadForFood : VisualRoutinesStateNode {

		/** Eventually, this node should sweep the head from a horizontal
				position to the ground position, building a map of food objects.
				For now, it just moves the head to a single position **/
		$nodeclass LookAtFood : HeadPointerNode : doStart {
			// Slow the head down
			for (unsigned int i = 0; i < NumHeadJoints; i++) {
				getMC()->setMaxSpeed(i, 0.5);
			}	
   	 getMC()->lookAtPoint(500, 0, 0);
  	}

		$setupmachine {
			LookAtFood =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));
		// Filter out noise		w
		SHAPEVEC_ITERATE(blob_shapes, BlobData, myblob)
			if (myblob->getArea() < 10) {
				camShS.deleteShape(myblob);
			}
		END_ITERATE;
		// 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 a the closest blob in the vector and get its location
			Shape<BlobData> blob = blob_shapes[0];
			float closestDist = blob->getCentroid().xyNorm();
			for (unsigned int i = 1; i < blob_shapes.size(); ++i) {
				// Get the distance of the next blob
				float nextDist = blob_shapes[i]->getCentroid().xyNorm();
				// Compare it to the current closest known distance
				if (nextDist < closestDist) {
					closestDist = nextDist;
					blob = blob_shapes[i];
				}
			}
			$reference SnackBot::nextFoodLoc;
			nextFoodLoc = blob->getCentroid();

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

			postStateSuccess();
		}
	}

	/** 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 camera to look at the food object
		**/
		$nodeclass LookAtFood : HeadPointerNode : doStart {
			// Slow the head down
			for (unsigned int i = 0; i < NumHeadJoints; i++) {
				getMC()->setMaxSpeed(i, 0.5);
			}	
   	 getMC()->lookAtPoint(300, 0, 0);
  	}

		/** 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 {
				TurnRobotTowardsFood =C=> MoveRobotTowardsFood
				  =C=> LookAtFood =C=> RelocateFoodObject =N=> 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> noOffset = fmat::pack(0, 0, 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, noOffset, 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, noOffset, 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.2);
		pe->setOutputCmd(RightFingerOffset, -0.2);
		// Add this to the motion sequence
		getMC()->advanceTime(1000);
		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 PrepToss : 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.75);
			getMC()->setOutputCmd(ArmShoulderOffset, 1.6);
			getMC()->setOutputCmd(ArmElbowOffset, 0.3);
			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.75);
			getMC()->setOutputCmd(ArmShoulderOffset, 1.6);
			getMC()->setOutputCmd(ArmElbowOffset, 0.3);
			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(ArmElbowOffset, 0);
			getMC()->setOutputCmd(ArmWristOffset, 0.2);
			// Move to the first position
			getMC()->advanceTime(50);
			getMC()->setPose(*pe);
			// Set the second position
			getMC()->setOutputCmd(ArmWristOffset, 0);
			getMC()->setOutputCmd(LeftFingerOffset, 0);
			getMC()->setOutputCmd(RightFingerOffset, 0);
			// Move to the second position
			getMC()->advanceTime(50);
			getMC()->setPose(*pe);		
			// Set the third position
			getMC()->setOutputCmd(LeftFingerOffset, 0);
			getMC()->setOutputCmd(RightFingerOffset, 0);
			// Move to the third position
			getMC()->advanceTime(100);
			getMC()->setPose(*pe);
			// Set the third position
			getMC()->setOutputCmd(ArmWristOffset, -0.1);			
			getMC()->setOutputCmd(LeftFingerOffset, -0.5);
			getMC()->setOutputCmd(RightFingerOffset, 0.5);
			// Move to the third position
			getMC()->advanceTime(100);
			getMC()->setPose(*pe);
			// Free memory
			delete(pe);		
		}
	
		$setupmachine {
			GripFood =TM("aim")=> PrepToss =TM("fire")=> TossFood
		}
	}

	$setupmachine {
		startdemo : StartNode
		foodLoc : FindFoodLocation

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

		foodLoc =F=> StateNode
		foodLoc =S=> PositionRobotNextToFood =C=> GrabFood =C=> MoveRobotToCobot
			=C=> GiveFoodToCobot
	}
}

REGISTER_BEHAVIOR(SnackBot);