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

      MapBuilderRequest* lex = new MapBuilderRequest(MapBuilderRequest::worldMap);
      lex->setAprilTagFamily();
      pilotreq.landmarkExtractor = lex;

    }

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