#include "Behaviors/StateMachine.h"
#include "Behaviors/Services/CameraBehavior.h"
#include "Events/TextMsgEvent.h"
#include <wait.h>

$nodeclass SurfTrain : VisualRoutinesStateNode {

  enum command {
    takeSnapshot, 
    train, 
    listObjects,
    clearData, 
    cont
  };

  // Global variables telling what class the image is and when to finish training
  $provide std::string className;
  
  $provide std::string tempSavePath;
  $provide std::string keypointDataPath;
  $provide std::string trainDataPath;
  $provide std::string objectDataPath;

  $provide std::string coords;

  $nodeclass Initialize : StateNode : doStart {
    $reference SurfTrain::tempSavePath;
    $reference SurfTrain::keypointDataPath;
    $reference SurfTrain::trainDataPath;
    $reference SurfTrain::objectDataPath;
    tempSavePath = "/tmp/image.raw";
    trainDataPath = "surf_data/trainData.dat";
    keypointDataPath = "surf_data/keypointData.dat";
    objectDataPath = "surf_data/objectDatabase.dat";
    postStateCompletion();
  }

  // node that saves the image into a temporary file
  $nodeclass SaveImage : VisualRoutinesStateNode {
		
    virtual void doStart()
    {

      $reference SurfTrain::tempSavePath;

      NEW_SKETCH(camY, uchar, sketchFromRawY());
      int width = camY->getWidth();
      int height = camY->getHeight();
      size_t buffsize = sizeof(char) * height * width;
      char *imgBuf = (char *) malloc(buffsize);
      camY->savePixels(imgBuf,buffsize);
      int factor = 2;
      char *destBuf = resample(imgBuf, width, height, factor);
      ofstream myfile;
      myfile.open(tempSavePath.c_str(), fstream::out 
		  | ios::binary);
      myfile.write(destBuf, buffsize * factor * factor);
      myfile.close();
      cout << "saved to " << tempSavePath << ". w: " 
	   << width << ", h: " << height  << endl;
      postStateCompletion();
    }
		
    char * resample(char * imgBuf, int width, int height, int factor) {
      size_t buffsize = sizeof(char) * height * factor 
	* width * factor;
      char *destBuf = (char *) malloc(buffsize);
      for (int i = 0; i < height; i++) {
	for (int j = 0; j < width; j++) {
	  int pos = xy2pos(i, j, width);
	  char val = imgBuf[pos];
	  int destWidth = (width * factor);
	  int destPos = 
	    xy2pos((i * factor),(j * factor), destWidth);
	  for (int k = 0; k < factor; k++) {
	    for (int l = 0; l < factor; l++) {
	      int resampleRowOffset = k * destWidth;
	      destBuf[destPos + (resampleRowOffset + l)] = val;
	    }					
	  }
	}
      }
      return destBuf;
    }

    int xy2pos(int x, int y, int width) {
      return x * width + y;
    }
		
  }

  // node that calls the train function in Find-Objects
  $nodeclass Fork : VisualRoutinesStateNode : doStart
    {
      $reference SurfTrain::className;
      $reference SurfTrain::tempSavePath;
      $reference SurfTrain::trainDataPath;
      $reference SurfTrain::keypointDataPath;
      $reference SurfTrain::objectDataPath;
      $reference SurfTrain::coords;

      time_t imageid;
      time(&imageid); 
      cout << imageid << endl;
      char idstr [50];
      sprintf(idstr,"c%d",imageid);
      pid_t child_id = fork();
      if ( child_id == 0 )
	{
	  char* tekrootval = getenv("TEKKOTSU_ROOT");
	  std::string const tekkotsuRoot =  "";
	  std::string const maryServerName = "find_object-train";
	  std::string const maryServerPath = tekkotsuRoot + maryServerName;
	  execl(maryServerPath.c_str(), maryServerName.c_str(), idstr,
		tempSavePath.c_str(), trainDataPath.c_str(),
		keypointDataPath.c_str(), coords.c_str(), "1280", "960", NULL);
	  // If we get here, the execlp() failed
	  std::cerr << "ERROR: failed to launch Mary server from " << maryServerPath << std::endl << "Check that TEKKOTSU_ROOT is set properly." << std::endl;
	  _exit(0);
	}  
      else 
	{
	  int status;
	  pid_t done = wait(&status);
	  // If success link image id to class id
	  ofstream myfile;
	  myfile.open(objectDataPath.c_str(), fstream::out | fstream::app);
	  myfile << className.c_str() << "\t" << idstr
		 << "\t" << coords  << endl;
	  myfile.close();
	  postStateCompletion();
	}
    }

  virtual void doStart() {
    cout << "Place the object in front of the camera and " 
	 << "input a class number \"c #\" to take a snapshot. " << endl;
    cout << "Type list to show the list of trained objects. " << endl;
    cout << "Type clear to remove all trained instances. " << endl;
  }


  // node that waits for the user to input what class the image was
  $nodeclass WaitForCommand : StateNode {
    virtual void doStart() {
      cout << "Enter command (c #|list|clear)>" << endl;
      erouter->addListener(this, EventBase::textmsgEGID);
    }
	  
    // upon getting a text message, enter this function
    void doEvent() {
      $reference SurfTrain::className;

      switch ( event->getGeneratorID() ) {
      case EventBase::textmsgEGID: {
	const TextMsgEvent *txtenv = 
	  dynamic_cast<const TextMsgEvent*>(event);
	std::string text = txtenv->getText();
	if (text.compare("list") == 0) {
	  postStateSignal<command>(listObjects);
	} else if (text.compare("clear") == 0) {
	  postStateSignal<command>(clearData);
	} else if (text.find("train") != string::npos) {
	  int pos = text.find(" ");
	  string subs = text.substr(pos + 1);
	  cout << "Class is: " << subs << endl;
	  className = subs;
	  postStateSignal<command>(takeSnapshot);
	} else {
	  cout << "Unexpected command: " << text << endl;
	  postStateSignal<command>(cont);
	}
	break;
      }
      default: {
	cout << "Unexpected event: " << event->getGeneratorID() << endl;
      }		  
      }
    }
  }

  // node that waits for the user to input what class the image was
  $nodeclass ReadCoords : StateNode {
    virtual void doStart() {
      cout << "Enter coordinates top_left_x,top_left_y,bot_right_x,bot_rigth_y>" << endl;
      erouter->addListener(this, EventBase::textmsgEGID);
    }
	  
    // upon getting a text message, enter this function
    void doEvent() {
      $reference SurfTrain::coords;
      switch ( event->getGeneratorID() ) {
      case EventBase::textmsgEGID: {
	const TextMsgEvent *txtenv = 
	  dynamic_cast<const TextMsgEvent*>(event);
	coords = txtenv->getText();
	// Validate input
	postStateCompletion();
	break;
      }
      default: {
	cout << "Unexpected event: " << event->getGeneratorID() << endl;
      }		  
      }
    }
  }

  $nodeclass ListObjects : StateNode : doStart
    {
      $reference SurfTrain::objectDataPath;

      ifstream myfile;
      myfile.open(objectDataPath.c_str());
      if(myfile.is_open())
	{
	  string line;
	  int count = 0;
	  int pos = 0;
	  while(myfile.good())
	    {
	      getline(myfile, line);
	      pos = line.find("\t");
	      string subs = line.substr(0, pos);
	      cout << subs << endl;
	      count++;
	    }
	  myfile.close();
	  if(count == 1 && pos == string::npos)
	    cout << "No objects in database." << endl;
	}
      else
	cout << "No objects in database." << endl;
      postStateCompletion();
    }

  $nodeclass ClearData : StateNode { 

    void clearFile(std::string fname) {
      ofstream myfile;
      myfile.open(fname.c_str(), fstream::out);
      myfile.close();
    }

    virtual void  doStart() {
      $reference SurfTrain::trainDataPath;
      $reference SurfTrain::keypointDataPath;
      $reference SurfTrain::objectDataPath;
      clearFile(trainDataPath);
      clearFile(keypointDataPath);
      clearFile(objectDataPath);
      cout << "All trained instances have been removed." << endl;
      postStateCompletion();
    }
  }

  $nodeclass Finish : StateNode : doStart
    {
      cout << "Done training model." << endl;
      postStateCompletion();
    }

  $nodeclass DisplayImage : VisualRoutinesStateNode : doStart {
    $reference SurfTrain::coords;
    std::istringstream iss(coords);
    std::string values[4];
    std::string token;
    int i = 0;
    while(getline(iss, token, ',')) {
      values[i++] = token;
    }
    int offsetX = atoi(values[0].c_str());
    int offsetY = atoi(values[1].c_str());
    rgb octaves[] = {rgb(127, 127, 255), rgb(255, 127, 127),
		     rgb(127, 255, 127), rgb(255, 255, 127)};
    NEW_SHAPE(scene, GraphicsData, new GraphicsData(camShS));
    std::ifstream fin;
    fin.open("easyKeypointData.txt");
	
    while(fin.is_open() && !fin.eof()) {
      std::string row = "";
      getline(fin, row);
      if (row.compare("") == 0) {
	break;
      }
      std::istringstream iss(row);
      std::string values[7];
      std::string token;
      int i = 0;
      while(getline(iss, token, ',')) {
	values[i++] = token;
      }
      float x = atof(values[0].c_str()) + offsetX;
      float y = atof(values[1].c_str()) + offsetY;
      float size = atof(values[2].c_str())*1.2/9.*2;
      float angle = atof(values[3].c_str());
      float response = atof(values[4].c_str());
      int octave = atoi(values[5].c_str());
      float classId = atof(values[6].c_str());
      std::pair<float,float> p(x/2,y/2);
      scene->add(new GraphicsData::CircleElement(p, size, false, 
						 octaves[octave - 1]));
      float r = size;
      float xp = r * cos(deg2rad(angle)) + x;
      float yp = r * sin(deg2rad(angle)) + y;
      std::pair<float,float> t(xp/2,yp/2);
      scene->add(new GraphicsData::LineElement(p, t, rgb(0,255,0)));
    }
    fin.close();
    postStateCompletion();
  }


  $nodeclass Clear : VisualRoutinesStateNode : doStart{
    camSkS.clear(); camShS.clear();
    postStateCompletion();
  }



  // save an image, wait for user input of what class the image is,
  // run the train function (which saves the features to a file),
  // then ask the user whether he wants to finish or take more images
  $setupmachine{
    Initialize =C=> wait
      coords: ReadCoords
      wait: WaitForCommand
      wait =S<command>(takeSnapshot)=> Clear =C=> SaveImage 
      =C=> coords =C=> Fork =C=> DisplayImage =C=> wait
      wait =S<command>(listObjects)=> ListObjects =C=> wait
      wait =S<command>(clearData)=> ClearData =C=> wait
      wait =S<command>(cont)=> wait
      }
}

REGISTER_BEHAVIOR(SurfTrain);