#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include <iostream>
#include <algorithm>
#include <functional>
#include <sys/stat.h>
#include <fstream>
#include <sstream>
#include <iterator>
#include <ctime>
#include <iomanip>

// needed for _mkdir
#if defined(_WIN32)
#include <direct.h>
#endif

#include "masking.h"
#include "stickletrack.h"
#include "tracking.h"

double frametime = 0;
int framenum = 0;

prefs Prefs;
props Props;

normalprefs normalPrefs;

static VideoCapture capture;

static bool stopvideo=false;
static bool leftbuttondown=false;
static Point2f startClickPos;
static int ilatetagintime=-1, iearlytagintime=-1;
static int startClickFrame;
static bool stoppedbefore;
static bool dooneframe=false;
static vector<tag> tags;
static int tagsselected=0;
static int nearestTags[2];
static vector<tag> **tagintime;
static int tagintimeidx=0;
static int moresleep=0;
static int gotoframe=0;
static ofstream tanjaLogFile;

struct comargs {
  int backBufferSize;
  char *videoFileName;
  int maxOutputFPS;
  double rescalingFactor;
};

static comargs comArgs;

void genBaseDir() { 
#if defined(_WIN32)
  Props.basedir = getenv("HOMEDRIVE");
  Props.basedir += getenv("HOMEPATH");
#else
  Props.basedir = getenv("HOME");
#endif
  
  Props.basedir += "/.stickletrack";

  cout << "Using " << Props.basedir << " for data output." << endl;
}

bool isWindowClosed (const char* name) {
  if ( cvGetWindowHandle(name) == NULL )
    return true;
  else
    return false;
}

void beforeExit() {
  cout << "Exitting ..." << endl;

#if defined(_WIN32)
  _mkdir(Props.basedir.c_str());
  _mkdir((Props.basedir+"/prefs").c_str());
#else 
  mkdir( Props.basedir.c_str(), 0777 );
  mkdir( (Props.basedir+"/prefs").c_str(), 0777 );
#endif

  cout << "Saving preferences ..." << endl;
  int scissorpointsize = Prefs.scissorpoints->size();
  ofstream prefFile( (Props.basedir+"/prefs/"+Props.videohash+".prf").c_str(), ios::binary );
  prefFile.write((char*)&Prefs, sizeof(Prefs));
  prefFile.write((char*)&scissorpointsize, sizeof(int));
  
  for (int ipoint = 0; ipoint < scissorpointsize; ipoint++)
    prefFile.write( (char*)&(*Prefs.scissorpoints)[ipoint], sizeof(Point) );

  prefFile.close();
}

void computeNormalizedPrefs() {
  normalPrefs.contours_minshortaxis = Prefs.contours_minshortaxis * Props.diagonal / 1000.0;
  normalPrefs.contours_maxshortaxis = Prefs.contours_maxshortaxis * Props.diagonal / 1000.0;
  normalPrefs.contours_minlongaxis = Prefs.contours_minlongaxis * Props.diagonal / 500.0;
  normalPrefs.contours_maxlongaxis = Prefs.contours_maxlongaxis * Props.diagonal / 500.0;

  normalPrefs.contours_minarea = pow(Prefs.contours_minarea,2) * Props.width * Props.height / 1000000.0;
  normalPrefs.contours_maxarea = pow(Prefs.contours_maxarea,2) * Props.width * Props.height / 1000000.0;
  
  normalPrefs.contours_mincircum = Prefs.contours_mincircum * Props.diagonal / 500.0;
  normalPrefs.contours_maxcircum = Prefs.contours_maxcircum * Props.diagonal / 500.0;
  
  normalPrefs.contours_maxspeed = Prefs.contours_maxspeed * Props.diagonal / 100.0;

  normalPrefs.contours_maxrot = 10 * Prefs.contours_maxrot;

  normalPrefs.halfdecay = 10.0 * Prefs.halfdecay / Props.fps;
}

void trackbarCallbackUpdateNormPrefs (int trackpos, void *userdata) {
  computeNormalizedPrefs();
}

void drawTimes(Mat& mContours) {
  stringstream alles;

  alles.precision(2);

  alles << framenum << " : " << fixed << (framenum)/(double)Props.fps << "s : +" << (double)moresleep/Props.fps << "ms";

  string text = alles.str();
  int fontFace = FONT_HERSHEY_DUPLEX;

  double fontScale = Props.width / 600.0;
  int thickness = Props.width / 250.0;

  int baseline=0;
  Size textSize = getTextSize(text, fontFace,
			      fontScale, thickness, &baseline);

  Point textOrg(Props.width*0.01, Props.width*0.01 + textSize.height);

  rectangle(mContours, Point(0, 0),
	    Point(Props.width, Props.width*0.02 + textSize.height),
	    Scalar(0,0,255), -1);

  putText(mContours, text, textOrg, fontFace, fontScale,
	  Scalar::all(255), thickness, 8);
}

void readCreatePrefs() {
  ifstream prefFile( (Props.basedir+"/prefs/"+Props.videohash+".prf").c_str(), ios::binary );
  if ( prefFile.is_open() ) {
    int scissorpointsize;

    prefFile.read((char*)&Prefs, sizeof(Prefs));
    prefFile.read((char*)&scissorpointsize, sizeof(int));

    Prefs.scissorpoints = new vector<Point>;
    for (int ipoint = 0; ipoint < scissorpointsize; ipoint++) {
      Point tmpPoint;
      prefFile.read((char*)&tmpPoint, sizeof(Point));
      Prefs.scissorpoints->push_back(tmpPoint);
    }

    prefFile.close();

    cout << "Loaded saved preferences." << endl;
  }
  else {
    Prefs.scissorpoints = new vector<Point>;
  }
}

void loadDefaultPrefs () {
  Prefs.halfdecay = 100;
  Prefs.forethreshold = 0;
  Prefs.mincolor[0] = 0; Prefs.mincolor[1] = 0; Prefs.mincolor[2] = 0;
  Prefs.maxcolor[0] = 255; Prefs.maxcolor[1] = 255; Prefs.maxcolor[2] = 255;
  Prefs.manyfish = 3;
  Prefs.contours_minshortaxis = 0;
  Prefs.contours_maxshortaxis = 100;
  Prefs.contours_minlongaxis = 0;
  Prefs.contours_maxlongaxis = 100;
  Prefs.contours_minarea = 0;
  Prefs.contours_maxarea = 100;
  Prefs.contours_mincircum = 0;
  Prefs.contours_maxcircum = 100;
  Prefs.contours_maxspeed = 100;
  Prefs.contours_maxrot = 100;
}

unsigned long frameHash(const Mat& frame) {
  unsigned long hash = 5381;
  int c;
  unsigned char *str = (unsigned char*)frame.data;

  while (c = *str++) {
    hash = ((hash << 5) + hash) + c;
  }

  return hash;
}

int timeWarp(int x, int y) {
  int horiMove = startClickPos.x - x;
  int wannago = startClickFrame + 10*horiMove*Props.fps/Props.width;

  int maxgo = min(ilatetagintime, Props.totframes);
  int mingo = max(1, ilatetagintime - comArgs.backBufferSize*Props.fps + 2);

  if (wannago > maxgo)
    return maxgo;
  else if (wannago < mingo)
    return mingo;
  else
    return wannago;
}

int rotatingTime( int whereyouwant ) {
  if ( whereyouwant < 0 )
    return comArgs.backBufferSize*Props.fps + whereyouwant;
  else
    return whereyouwant;
}

void mouseTracking(int evt, int x, int y, int flags, void* param){
  if ( stopvideo && evt == CV_EVENT_LBUTTONDBLCLK ) {
    tagsselected++;

    double mindist = -1;
    for (int itag = 0; itag < tags.size(); itag++) {
      double dist = pow(tags[itag].x - x, 2) + pow(tags[itag].y - y, 2);
      if ( dist < mindist || mindist == -1 ) {
	nearestTags[tagsselected-1] = itag;
	mindist = dist;
      }
    }

    if (tagsselected == 2) {
      if ( nearestTags[0] < tags.size() && nearestTags[0] >= 0
	   && nearestTags[1] < tags.size() && nearestTags[1] >= 0
	   && nearestTags[0] != nearestTags[1] ) {

	tag tmp;

	tmp = tags[nearestTags[0]];
	tags[nearestTags[0]] = tags[nearestTags[1]];
	tags[nearestTags[1]] = tmp;
	tags[nearestTags[1]].color = tags[nearestTags[0]].color;
	tags[nearestTags[0]].color = tmp.color;

	for (int itagintime = framenum-1; itagintime <= ilatetagintime; itagintime++) {
	  tmp = (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[0]];

	  (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[0]] = 
	    (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[1]];

	  (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[1]] = tmp;

	  (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[1]].color = 
	    (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[0]].color;

	  (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[0]].color = tmp.color;
	} 
      }
      tagsselected = 0;

      gotoframe = framenum;
      dooneframe = true;
    }
  }

  if ( evt == CV_EVENT_RBUTTONUP ) {
    stopvideo = (stopvideo+1)%2;
  }

  if (evt==CV_EVENT_LBUTTONDOWN) {
    leftbuttondown = true;
    startClickPos = Point2f(x,y);
    startClickFrame = framenum;
    
    if (stopvideo)
      stoppedbefore = true;
    else
      stoppedbefore = false;
    
    stopvideo = true;
  }

  if (evt==CV_EVENT_LBUTTONUP) {
    leftbuttondown = false;
    
    int vertiMove = startClickPos.y - y;

    if ( ! stoppedbefore )
      stopvideo = false;
 }

  if ( leftbuttondown && startClickPos.x - x != 0  && ! (startClickPos.x > Props.width*5.0/6.0 && x > Props.width*5.0/6.0) ) {
    gotoframe = timeWarp(x, y);
    dooneframe = true;
  }

  if ( leftbuttondown && startClickPos.y - y != 0) {
    int vertiMove = startClickPos.y - y;

    if ( startClickPos.x > Props.width*5.0/6.0 && x > Props.width*5.0/6.0 ) {
      moresleep -= vertiMove*10000/Props.height;
      if ( moresleep < 0 )
	moresleep = 0;
      
      drawTimes(tracking_getFrame());
    }    
  }
}

void writeTanjaLog (int writeframenum, vector<tag> *tagstowrite) {
  stringstream outline;

  if ( writeframenum%(Props.fps/comArgs.maxOutputFPS) == 0 ) {
    if ( tagstowrite->size() >= Prefs.manyfish ) {
      outline << (double)writeframenum/Props.fps;
      for ( int itag = 0; itag < Prefs.manyfish; itag++) {
	outline << "," << (*tagstowrite)[itag].x << "," << (*tagstowrite)[itag].y << "," << (double)writeframenum/Props.fps-(*tagstowrite)[itag].lastseen;
      }
      outline << endl;
      tanjaLogFile.write( outline.str().c_str(), outline.str().size() );
      tanjaLogFile.flush();
    }
  }
}

void openTanjaLog () {
  stringstream tanjalogfilename;

#if defined(_WIN32)
  _mkdir(Props.basedir.c_str());
  _mkdir((Props.basedir+"/logs").c_str());
#else 
  mkdir( Props.basedir.c_str(), 0777 );
  mkdir( (Props.basedir+"/logs").c_str(), 0777 );
#endif

#if defined(_WIN32)
  _mkdir(Props.basedir.c_str());
  _mkdir((Props.basedir+"/logs/"+Props.videohash).c_str());
#else 
  mkdir( Props.basedir.c_str(), 0777 );
  mkdir( (Props.basedir+"/logs/"+Props.videohash).c_str(), 0777 );
#endif

  tanjalogfilename << Props.basedir << "/logs/"  << Props.videohash << "/" << time(0) << ".dat";
  tanjaLogFile.open( tanjalogfilename.str().c_str() );

  if ( ! tanjaLogFile.is_open() ) {
    cerr << "Could not open tanjalogfile " << tanjalogfilename.str() << "!" << endl;
    exit(1);
  }

  tanjaLogFile.write("# STICKLETRACK - TANJALOG\n", 26);
  tanjaLogFile.write("# format: <time in s>,<tag_1 xpos>,<tag_1 ypos>,<tag_1 losttime>,...,<tag_n xpos>,<tag_n ypos>,<tag_n losttime>\n", 112);

  tanjaLogFile.flush();
}

int process(VideoCapture& capture) {
  int n = 0;
  char filename[200];
  bool pleaseExit = false;

  namedWindow("stickletrack_original", CV_WINDOW_KEEPRATIO);
    
  Mat frame, origframe, combinedmask;

  Props.fps = capture.get(CV_CAP_PROP_FPS);
  Props.width = capture.get(CV_CAP_PROP_FRAME_WIDTH)*comArgs.rescalingFactor;
  Props.height = capture.get(CV_CAP_PROP_FRAME_HEIGHT)*comArgs.rescalingFactor;
  Props.totframes = capture.get(CV_CAP_PROP_FRAME_COUNT);
  Props.diagonal = sqrt( pow(Props.width, 2) + pow(Props.height, 2) );

  if ( Props.width == 0 || Props.height == 0 || Props.fps == 0 || Props.totframes == 0 ) {
    cerr << "Something got wrong while reading video-file info!" << endl;
    cerr << "Width: " << Props.width << endl;
    cerr << "Height: " << Props.height << endl;
    cerr << "FPS: " << Props.fps << endl;
    cerr << "Total frames: " << Props.totframes << endl;
    exit(1);
  }

  Mat frameintime[comArgs.backBufferSize*Props.fps];

  tagintime = new vector<tag>*[comArgs.backBufferSize*Props.fps];

  for (int iback = 0; iback < comArgs.backBufferSize*Props.fps; iback++)
    tagintime[iback] = new vector<tag>;

  capture >> frame;

  if (frame.empty())
    exit(1);

  stringstream tmphash;
  tmphash << frameHash(frame);
  Props.videohash = tmphash.str();

  loadDefaultPrefs();
  readCreatePrefs();

  computeNormalizedPrefs();

  openTanjaLog();

  masking_init();
  tracking_init(&mouseTracking);

  capture.set(CV_CAP_PROP_POS_FRAMES, 0);

  for (; !pleaseExit;) {
    if ( !stopvideo || dooneframe ) {
      framenum = gotoframe;

      dooneframe = false;

      frametime = (framenum-1.0) / Props.fps;

      if ( framenum <= ilatetagintime ) {
	origframe = frameintime[ rotatingTime(tagintimeidx - (ilatetagintime-framenum) - 2) ];
	masking_getCombinedMask(origframe, combinedmask);

	tags.clear();
	for (int itag = 0; itag < tagintime[ rotatingTime( tagintimeidx - 2 - (ilatetagintime-framenum) ) ]->size(); itag++)
	  tags.push_back( (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-framenum) - 2) ])[itag] );
	
	tracking_locateTags (tags, combinedmask);
      }
      else {
	capture >> frame;

	if (frame.empty())
	  break;

	frame.convertTo(origframe, CV_32FC3);
	resize(origframe, origframe, Size(0,0), comArgs.rescalingFactor, comArgs.rescalingFactor);

	masking_getCombinedMask(origframe, combinedmask);
      
	tracking_locateTags (tags, combinedmask);

	tagintime[tagintimeidx]->clear();
	for (int itag = 0; itag < tags.size(); itag++)
	  tagintime[tagintimeidx]->push_back(tags[itag]);

	frameintime[tagintimeidx] = origframe;
	
	ilatetagintime = framenum;
	tagintimeidx = (tagintimeidx+1)%(comArgs.backBufferSize*Props.fps);

	if ( ilatetagintime >= comArgs.backBufferSize*Props.fps-1 )
	  writeTanjaLog( framenum - (comArgs.backBufferSize*Props.fps-1), tagintime[ tagintimeidx ] );
      }

      gotoframe = framenum + 1;
 
      if ( ! isWindowClosed("stickletrack_original") )
	imshow("stickletrack_original", origframe/255.0);
    }
    
    drawTimes(tracking_getFrame());

    if ( tagsselected == 1 ) {
      circle( tracking_getFrame(), Point2f(tags[nearestTags[0]].x, tags[nearestTags[0]].y), Props.diagonal / 100.0, Scalar(0,0,255), -1, 8 );
    }
   
    if ( tracking_showFrame() )
      pleaseExit = true;

    char key;

    if (!stopvideo)
      key = (char)waitKey(5 + (double)moresleep/Props.fps);
    else
      key = (char)waitKey(5);

    switch (key) {
    case 'q':
      pleaseExit = true;
      break;
    default:
      break;
    }

  }
  
  if ( ilatetagintime >= comArgs.backBufferSize*Props.fps-1 ) {
    int timeidx = (tagintimeidx+1)%(comArgs.backBufferSize*Props.fps);
    for ( int itime = 0; itime < comArgs.backBufferSize*Props.fps-1; itime++ ) {
      writeTanjaLog( framenum+1+itime - (comArgs.backBufferSize*Props.fps-1), tagintime[ timeidx ] );
      timeidx = (timeidx+1)%(comArgs.backBufferSize*Props.fps);
    }
  }
  else {
    for ( int itime = 0; itime < ilatetagintime+1; itime++ )
      writeTanjaLog( itime, tagintime[ itime ] );
  }

  tanjaLogFile.close();

  beforeExit();

  cout << "Bye bye." << endl;

  return 0;
}

void showUsage() {
  cout << "Usage: stickletrack [OPTIONS]... <path to videofile>" << endl;
  cout << endl << "Available options:" << endl;
  cout << "-x <rescalingfactor>" << "\t" << "rescale video by a factor <rescalingfactor>" << endl;
  cout << "-o <maxfps>" << "\t" << "write to log maximally <maxfps> times per second" << endl;
  cout << "-b <bsize>" << "\t" << "buffer <bsize> seconds for going back in time" << endl;
}

bool parseComArgs (int ac, char* av[]) {
  if ( ac < 2 || ac%2 ) {
    showUsage();
    exit(0);
  }

  comArgs.videoFileName = NULL;
  comArgs.rescalingFactor = 1.0;
  comArgs.maxOutputFPS = 25;
  comArgs.backBufferSize = 10;

  for (int i = 1; i < ac-2; i = i+2)
    if ( strcmp(av[i], "-x") == 0 ) {
      comArgs.rescalingFactor = atof(av[i + 1]);
    } else if ( strcmp(av[i], "-o") == 0 ) {
      comArgs.maxOutputFPS = atoi(av[i + 1]);
    } else if ( strcmp(av[i], "-b") == 0 ) {
      comArgs.backBufferSize = atoi(av[i + 1]);
    } else {
      cout << "Unknown command line parameters!\n";
      showUsage();
      exit(0);
    }

  comArgs.videoFileName = av[ac-1];
}

int main(int ac, char* av[]) {
  parseComArgs(ac, av);

  capture.open( comArgs.videoFileName );

  genBaseDir();
  
  if (!capture.isOpened()) {
    cerr << "Failed to open a video device or video file!\n" << endl;
    return 1;
  }
  
  return process(capture);
}