#include "Behaviors/StateMachine.h"

using namespace std;

$nodeclass HW5_Student : VisualRoutinesStateNode {
  $provide LocalizationParticle agent();
  $provide vector<Landmark> lmks();
  $provide int numLmks(0);
  $provide int curLmk(0);
  $provide int bestParticle(0);
  static const int numParticles = 10;
  
  $nodeclass Emaze : VisualRoutinesStateNode : doStart {
    cout << endl;
    cout << "----------" << endl;
    cout << endl;

    cout << "Building map..." << endl;
		
    // world bounds
		
    vector<Point> pts;
    pts.push_back(Point(   -457.2,      457.2));
    pts.push_back(Point(   1981.2,      457.2));
    pts.push_back(Point(   1981.2,    -1524.0));
    pts.push_back(Point(   1219.2,    -1524.0));
    pts.push_back(Point(   1219.2,     -355.6));
    pts.push_back(Point(   1143.0,     -355.6));
    pts.push_back(Point(   1143.0,    -1524.0));
    pts.push_back(Point(    381.0,    -1524.0));
    pts.push_back(Point(    381.0,     -355.6));
    pts.push_back(Point(    304.8,     -355.6));
    pts.push_back(Point(    304.8,    -1524.0));
    pts.push_back(Point(   -457.2,    -1524.0));
		
    NEW_SHAPE(worldBounds, PolygonData, new PolygonData(worldShS, pts, true));
		
    // add landmarks
		
    NEW_SHAPE(tag0, AprilTagData,
	      new AprilTagData(worldShS, Point(-456.2,    50.8, 177.8),
			       AprilTags::TagDetection(0)));
		
    NEW_SHAPE(tag1, AprilTagData,
	      new AprilTagData(worldShS, Point( -76.2,   456.2, 177.8),
			       AprilTags::TagDetection(1)));

    NEW_SHAPE(tag2, AprilTagData,
	      new AprilTagData(worldShS, Point( 762.0,   456.2, 177.8),
			       AprilTags::TagDetection(2)));

    NEW_SHAPE(tag3, AprilTagData,
	      new AprilTagData(worldShS, Point(1600.2,   456.2, 177.8),
			       AprilTags::TagDetection(3)));

    NEW_SHAPE(tag4, AprilTagData,
	      new AprilTagData(worldShS, Point(1980.2,    50.8, 177.8),
			       AprilTags::TagDetection(4)));
		
    NEW_SHAPE(tag5, AprilTagData,
	      new AprilTagData(worldShS, Point(1980.2,  -508.0, 177.8),
			       AprilTags::TagDetection(5)));
		
    NEW_SHAPE(tag6, AprilTagData,
	      new AprilTagData(worldShS, Point(1980.2, -1371.6, 177.8),
			       AprilTags::TagDetection(6)));
		
    NEW_SHAPE(tag7, AprilTagData,
	      new AprilTagData(worldShS, Point(1600.2, -1523.0, 177.8),
			       AprilTags::TagDetection(7)));

    NEW_SHAPE(tag8, AprilTagData,
	      new AprilTagData(worldShS, Point(1220.2, -1371.6, 177.8),
			       AprilTags::TagDetection(8)));

    NEW_SHAPE(tag9, AprilTagData,
	      new AprilTagData(worldShS, Point(1220.2,  -508.0, 177.8),
			       AprilTags::TagDetection(9)));

    NEW_SHAPE(tag10, AprilTagData,
	      new AprilTagData(worldShS, Point(1142.0,  -508.0, 177.8),
			       AprilTags::TagDetection(10)));

    NEW_SHAPE(tag11, AprilTagData,
	      new AprilTagData(worldShS, Point(1142.0, -1371.6, 177.8),
			       AprilTags::TagDetection(11)));

    NEW_SHAPE(tag12, AprilTagData,
	      new AprilTagData(worldShS, Point( 762.0, -1523.0, 177.8),
			       AprilTags::TagDetection(12)));

    NEW_SHAPE(tag13, AprilTagData,
	      new AprilTagData(worldShS, Point( 382.0, -1371.6, 177.8),
			       AprilTags::TagDetection(13)));

    NEW_SHAPE(tag14, AprilTagData,
	      new AprilTagData(worldShS, Point( 382.0,  -508.0, 177.8),
			       AprilTags::TagDetection(14)));

    NEW_SHAPE(tag15, AprilTagData,
	      new AprilTagData(worldShS, Point( 303.8,  -508.0, 177.8),
			       AprilTags::TagDetection(15)));

    NEW_SHAPE(tag16, AprilTagData,
	      new AprilTagData(worldShS, Point( 303.8, -1371.6, 177.8),
			       AprilTags::TagDetection(16)));

    NEW_SHAPE(tag17, AprilTagData,
	      new AprilTagData(worldShS, Point( -76.2, -1523.0, 177.8),
			       AprilTags::TagDetection(17)));

    NEW_SHAPE(tag18, AprilTagData,
	      new AprilTagData(worldShS, Point(-456.2, -1371.6, 177.8),
			       AprilTags::TagDetection(18)));

    NEW_SHAPE(tag19, AprilTagData,
	      new AprilTagData(worldShS, Point(-456.2,  -508.0, 177.8),
			       AprilTags::TagDetection(19)));

    pilot->setWorldBounds(worldBounds);
  };
	
  $nodeclass StartStudent : StateNode : doStart {
    cout << endl;
    cout << "----------" << endl;
    cout << endl;

    cout << "Welcome to Particle Filter Bingo: Student Mode" << endl;
    cout << endl;
    cout << "This is a tool to teach you how particle filters work. Your teacher has chosen the position and direction of the robot in the grid and will tell you the approximate locations of some landmarks that the robot can see. Your job will be to use this information to figure out the position and orientation of the robot using particle filters. Since particle filters are randomized, there is a luck factor, and the student who gets closest to the position of the robot will be the winner." << endl;
  }

  $nodeclass RequestHexCode : StateNode : doStart {
    cout << endl;
    cout << "----------" << endl;
    cout << endl;

    cout << "Enter the hexadecimal code that your teacher gives you." << endl;
  }

  $nodeclass ParseHexCode : StateNode : doStart {
    $reference HW5_Student::agent;

    const TextMsgEvent *txtev = dynamic_cast<const TextMsgEvent*>(event);
    if ( txtev == NULL ) {
      cout << "Error: Null text event" << endl;
      postStateFailure(); // Failure retries
      return;
    }

    cout << "Parsing hex code... ";
    std::istringstream is(txtev->getText());

    union {
      struct {
        float x, y, th;
      } fval;
      struct {
        int a, b, c;
      } ival;
    } hexConverter;

    // Read in the 3 hexadecimal integers, and convert to floats
    char hexCode[25];
    is.getline(hexCode, 25);
    cout << hexCode << endl;

    if (strlen(hexCode)!=24) {
      cout << "Error: Invalid hex code" << endl;
      postStateFailure();
      return;
    }

    // Each group of 8 hex chars corresponds with the bit level representation of a float
    // Hopefully the students don't figure this out :-)
    sscanf(hexCode, "%8x%8x%8x", &hexConverter.ival.a, &hexConverter.ival.b, &hexConverter.ival.c);

    /*
    cout << "X: " << hexConverter.fval.x << endl;
    cout << "Y: " << hexConverter.fval.y << endl;
    cout << "Th: " << hexConverter.fval.th << endl;
    */

    // Convert the ints to floats and store
    agent.x = hexConverter.fval.x;
    agent.y = hexConverter.fval.y;
    agent.theta = hexConverter.fval.th;

    postStateSuccess(); // Success continues
  }

  $nodeclass RequestNumLandmarks : StateNode : doStart {
    cout << endl;
    cout << "----------" << endl;
    cout << endl;

    cout << "Enter the number of landmarks that your teacher will give you." << endl;
  }

  $nodeclass ParseNumLandmarks : StateNode : doStart {
    $reference HW5_Student::numLmks;

    const TextMsgEvent *txtev = dynamic_cast<const TextMsgEvent*>(event);
    if ( txtev == NULL ) {
      cout << "Error: Null text event" << endl;
      postStateFailure(); // Failure retries
      return;
    }

    cout << "Parsing number of landmarks... ";
    std::istringstream is(txtev->getText());
    is >> numLmks;
    cout << numLmks << endl;

    if (numLmks > 0 && numLmks <= 4) {
      postStateSuccess(); // Success continues
    } else {
      cout << "Error: Invalid number of landmarks." << endl;
      postStateFailure();
    }
  }

  $nodeclass GenerateParticles : StateNode : doStart {
    $reference HW5_Student::curLmk;
    
    particleFilter->resizeParticles(HW5_Student::numParticles);
    particleFilter->resetFilter();
    particleFilter->displayParticles(HW5_Student::numParticles);

    curLmk = 0;
  }

  $nodeclass ShowScore : StateNode : doStart {
    $reference HW5_Student::lmks;
    vector<LocalizationParticle> particles = particleFilter->getParticles();

    int bestParticle = -1;
    float bestScore = 1.0e100;

    cout << endl;
    cout << lmks.size() << " landmarks" << endl;
    
    cout << "Particle |    x    |    y    |  theta ";
    for (vector<Landmark>::iterator it = lmks.begin(); it != lmks.end(); ++it) {
      cout << " |  Lmk " << setw(2) << (*it).id;
    }
    cout << " | Total Score" << endl;

    cout << "--------------------------------------";
    for (vector<Landmark>::iterator it = lmks.begin(); it != lmks.end(); ++it) {
      cout << "---" << "--------";
    }
    cout << "--------------" << endl;
    
    for (int i = 0; i < HW5_Student::numParticles; i++) {
      LocalizationParticle &pt = particles[i];

      cout << setw(8) << (i+1) << " | ";
      cout << setw(7) << fixed << setprecision(1) << pt.x << " | ";
      cout << setw(7) << fixed << setprecision(1) << pt.y << " | ";
      cout << setw(7) << fixed << setprecision(3) << pt.theta;

      float errSum = 0;
      for (vector<Landmark>::iterator it = lmks.begin(); it != lmks.end(); ++it) {
        Landmark &lmk = *it;
        float error = lmk.computeError(pt);
        errSum += error;

        cout << " | " << fixed << setw(7) << setprecision(1) << error;
      }

      cout << " | " << fixed << setw(11) << setprecision(1) << errSum << endl;

      if (errSum < bestScore) {
        bestScore = errSum;
        bestParticle = i;
      }
    }

    cout << endl;
    cout << "Best particle is " << bestParticle+1  << " with score " << bestScore << endl;

    particleFilter->setBestIndex(bestParticle);

    cout << "Refresh the world shape space to see where your best guess is so far." << endl;
  }

  $nodeclass StudentTurn : StateNode : doStart {
    $reference HW5_Student::numLmks;
    $reference HW5_Student::curLmk;

    curLmk++;
    
    if (curLmk<=numLmks) {
      cout << endl;
      cout << "----------" << endl;
      cout << endl;

      cout << "Round " << curLmk << ":" << endl;
      cout << "Enter the landmark number, distance, and theta separated by spaces for the landmark that your teacher provides you." << endl;

      postStateSuccess();
    } else {
      postStateFailure(); // We've gone through all the landmarks
    }
  }
  
  class Landmark {
  public:
    int id;
    float theta, dist;
    Point center;

    Landmark(const int _id, const float _theta, const float _dist) : id(_id), theta(_theta), dist(_dist), center() {
      // Find the matching April Tag
      NEW_SHAPEVEC(tags, AprilTagData, select_type<AprilTagData>(worldShS));
      SHAPEVEC_ITERATE(tags, AprilTagData, t) {
        if (t->getTagID()==id) center = t->getCentroid();
      } END_ITERATE;
    }

    float computeError(LocalizationParticle &testpt) {
      Point pt(testpt.x, testpt.y);
      Point pt2lmk = center - pt;
      float testDist = pt2lmk.xyNorm();
      float testTheta = pt2lmk.atanYX() - testpt.theta;

      float errDist = testDist - dist;
      float errTheta = testTheta - theta;

      // sum squared error over 1000 for smaller numbers
      return (errDist*errDist + errTheta*errTheta)/1000;
    }
  };

  $nodeclass ParseLandmark : StateNode : doStart {
    $reference HW5_Student::lmks;

    const TextMsgEvent *txtev = dynamic_cast<const TextMsgEvent*>(event);
    if ( txtev == NULL ) {
      cout << "Error: Null text event" << endl;
      postStateFailure(); // Failure retries
      return;
    }

    cout << "Parsing landmark..." << endl;
    std::istringstream is(txtev->getText());

    int id;
    float dist, theta;
    is >> id >> dist >> theta;

    cout << "Landmark #" << id << endl;
    cout << "Distance: " << dist << endl;
    cout << "Theta: " << theta << endl;

    lmks.push_back(Landmark(id, theta, dist));

    postStateSuccess(); // Success continues
  }

  $nodeclass Done : StateNode : doStart {
    $reference HW5_Student::agent;
    $reference HW5_Student::bestParticle;

    cout << endl;
    cout << "----------" << endl;
    cout << endl;

    cout << "The actual position of the agent is:" << endl;
    cout << "X: " << agent.x << endl;
    cout << "Y: " << agent.y << endl;
    cout << "Theta: " << agent.theta << endl;
    
    cout << endl;

    LocalizationParticle &best = particleFilter->getParticles()[bestParticle];
    cout << "Your best particle was:" << endl;
    cout << "X: " << best.x << endl;
    cout << "Y: " << best.y << endl;
    cout << "Theta: " << best.theta << endl;

    cout << endl;

    cout << "This means the overall error for your particle filter is: " << agent.sumSqErr(best) << endl;

    cout << endl;

    cout << "Report your error and best particle to your teacher and they will pick the student who has the best particle in the class." << endl;

    cout << endl;

    cout << "Do you understand the power of particle filters now? You only made " << numParticles << " guesses as to the location of the robot, and that led to an okay estimate. But when you have lots of guesses by combining all the particles in your class, you can make an even better estimate and get pretty close." << endl;
    cout << "Also, note that each landmark your teacher gives you helps you make a better estimate of which particle is the closest. After just one landmark, your best particle probably was very wrong after you added in all the landmarks." << endl;

    stop();
  }

  $setupmachine {
    startnode: StartStudent =N=> Emaze =N=> reqHC
  
    // Get the hex code
    reqHC: RequestHexCode =TM=> parseHC =S=> reqNumLmks
    parseHC: ParseHexCode =F=> reqHC

    // Get the number of landmarks
    reqNumLmks: RequestNumLandmarks =TM=> parseNumLmks
    parseNumLmks: ParseNumLandmarks =F=> reqNumLmks

    // Initialize the particles
    parseNumLmks =S=> GenerateParticles =N=> ShowScore =N=> turn

    // Request one landmark at a time
    turn: StudentTurn =TM=> parseLmk =S=> ShowScore =N=> turn
    parseLmk: ParseLandmark =F=> StateNode =TM=> parseLmk
    
    // After read all landmarks, report results
    turn =F=> Done
  }
}

REGISTER_BEHAVIOR(HW5_Student);