#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)); **/
    penPoints.push_back(Point(-425, -425, 0, allocentric));
    penPoints.push_back(Point(2321, -425, 0, allocentric));
    penPoints.push_back(Point(2321, 425, 0, allocentric));
    penPoints.push_back(Point(1350, 425, 0, allocentric));				
    penPoints.push_back(Point(1350, 1281, 0, allocentric));
    penPoints.push_back(Point(650, 1281, 0, allocentric));
    penPoints.push_back(Point(650, 425, 0, allocentric));
    penPoints.push_back(Point(475, 425, 0, allocentric));
    penPoints.push_back(Point(475, 1281, 0, allocentric));
    penPoints.push_back(Point(-425, 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)));
  }

	/** 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 {
      RetractArm =C=> TurnRobotTowardsFood =C=> MoveRobotTowardsFood
				=C=> LookAtFood =C=> FindFood =C=> FindFoodLocation =N=> PostMachineCompletion
		}
  }	

	/** ------------------------ Arm Control ------------------------ **/

	/** Moves the arm to a retracted position with its gripper closed **/	
  $nodeclass RetractArm : 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);
    }
    // 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);
  }

	/** ----------------------- 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
		}
  }

	/** ----------------------- Kinematics -------------------------- **/

  /** 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);
  }

    /** 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, 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);			
  }

  /** 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);
  }

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

	/** 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
    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 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 Machines -------------------------- **/	

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

  /** State machine which looks for multiple april tags to localize against **/
  $nodeclass LocalizeRobot : StateNode {
    $nodeclass MyLocalize : PilotNode(localize) : doStart {
			// Set the landmarks to use for localization
      GET_SHAPE(tag18, AprilTagData, worldShS);	
      GET_SHAPE(tag2, AprilTagData, worldShS);	
			GET_SHAPE(tag12, AprilTagData, worldShS);	
			GET_SHAPE(tag3, AprilTagData, worldShS);	
			GET_SHAPE(tag7, AprilTagData, worldShS);	
			GET_SHAPE(tag4, AprilTagData, worldShS);	
			GET_SHAPE(tag17, AprilTagData, worldShS);	
			GET_SHAPE(tag1, AprilTagData, worldShS);				
      pilotreq.landmarks.clear();
      pilotreq.landmarks.push_back(tag18);
      pilotreq.landmarks.push_back(tag2);
      pilotreq.landmarks.push_back(tag12);
      pilotreq.landmarks.push_back(tag3);
      pilotreq.landmarks.push_back(tag7);
      pilotreq.landmarks.push_back(tag4);
      pilotreq.landmarks.push_back(tag17);
      pilotreq.landmarks.push_back(tag1);
			// Tell the robot which kind of april tags we're using
      MapBuilderRequest* lex = new MapBuilderRequest(MapBuilderRequest::worldMap);
      lex->setAprilTagFamily();
      pilotreq.landmarkExtractor = lex;
    }

/**    $setupmachine {
      LookUp =C=> Localize =C=>
				LookAtPoint(1000, -250, TAG_HEIGHT) =C=> Localize =C=>
				LookAtPoint(1000, -500, TAG_HEIGHT) =C=> Localize =C=>				
				LookAtPoint(1000, -1000, TAG_HEIGHT) =C=> Localize =C=>				
				LookAtPoint(0, -1000, TAG_HEIGHT) =C=> Localize =C=>
				LookAtPoint(1000, 250, TAG_HEIGHT) =C=> Localize =C=>
				LookAtPoint(1000, 500, TAG_HEIGHT) =C=> Localize =C=>
				LookAtPoint(1000, 1000, TAG_HEIGHT) =C=> Localize =C=>
				LookAtPoint(0, 1000, TAG_HEIGHT) =C=> Localize =C=> PostMachineCompletion
		} **/

		$setupmachine {
      LookUp =C=> Localize =C=>
				LookAtPoint(1000, -250, TAG_HEIGHT) =C=> Localize =C=>
				LookAtPoint(1000, 250, TAG_HEIGHT) =C=> Localize =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=> FaceTowardsY =C=> PostMachineCompletion
		}
  }

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

  $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, 0);
    	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, 0);
      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 =TM("ready")=> RetractArm =TM("aim")=> Prep1 =C=> Prep2 =C=> Prep3 =TM("fire")=> TossFood =TM("done")=> RetractArm
		}
  }

	/** State machine which moves the robot onto the x-axis **/
	$nodeclass MoveRobotToXAxis : StateNode {
		/** Backs the robot up onto the x-axis **/
    $nodeclass BackToX : PilotNode(walk) : doStart {
      pilotreq.dx = -theAgent->getCentroid().coordY();
    }

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

  $nodeclass PositionRobotNearCobot : StateNode {
		
		$nodeclass FindRedBlueBlobs : MapBuilderNode(MapBuilderRequest::localMap) : doStart {
    	mapreq.addObjectColor(blobDataType, "blue");
    	mapreq.addObjectColor(blobDataType, "red");			
  	}

		$setupmachine {
			LookUp =C=> FindRedBlueBlobs
		}
	}

  $setupmachine {
		startdemo : StartNode
		foodLoc : FindFoodLocation

//    startdemo =N=> LocalizeRobot

    startdemo =N=> LocalizeRobot =TM("go")=> MoveRobotToFoodLoc =C=>
			SweepHeadForFood =C=> foodLoc

    foodLoc =F=> StateNode
    foodLoc =TM("go")=> PositionRobotNextToFood =TM("go")=>
			GrabFood =TM("go")=> MoveRobotToXAxis =C=> LocalizeRobot =TM("go")=> 
			MoveRobotToCobot =C=> PositionRobotNearCobot =TM("go")=> GiveFoodToCobot

//		startdemo =N=> SweepHeadForFood =C=> FindFoodLocation =S=>
//			PositionRobotNextToFood =C=> GrabFood

//		startdemo =N=> GiveFoodToCobot
	}
}

REGISTER_BEHAVIOR(SnackBot);