#include <stdio.h>
#include <iostream>

using namespace std;

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

#define SEPCOST_AREA      0
#define SEPCOST_CIRCUM    1
#define SEPCOST_SHORTAXIS 2
#define SEPCOST_LONGAXIS  3
#define SEPCOST_ANGLE     4
#define SEPCOST_DIST      5

struct fishblobb {
  double x;
  double y;
  bool identified;
  int contourid;
  vector<Point> *points;
  double longaxis;
  double shortaxis;
  double angle;
  int area;
  int circum;
  Point2f head, tail;
};

static int enablecontours=0;
static Mat mContours;

struct tagcost {
  int tagid;
  double tagcost;
};

struct blobbcost {
  int blobbid;
  vector<tagcost> tagcosts;
};

bool sort_tags (tagcost i, tagcost j) {
  return i.tagcost < j.tagcost;
}

bool sort_blobbs (blobbcost i, blobbcost j) {
  return i.tagcosts[0].tagcost < j.tagcosts[0].tagcost;
}

Mat& tracking_getFrame() {
  return mContours;
}

int tracking_showFrame() {
  if ( ! isWindowClosed("stickletrack_tracking_screen") ) {
    imshow("stickletrack_tracking_screen", mContours);
    return 0;
  }
  else
    return 1;
}

bool tracking_isEnabled () {
  return enablecontours;
}

void tracking_init( void (*mouseTracking)(int, int, int, int, void*) ) {
  namedWindow("stickletrack_tracking_prefs", CV_WINDOW_KEEPRATIO);
  namedWindow("stickletrack_tracking_screen", CV_WINDOW_KEEPRATIO);

  createTrackbar("Enable", "stickletrack_tracking_prefs", &enablecontours, 1, NULL, 0);
  createTrackbar("Manyfish", "stickletrack_tracking_prefs", &Prefs.manyfish, 10, NULL, 0);
  createTrackbar("min area", "stickletrack_tracking_prefs", &Prefs.contours_minarea, 100, &trackbarCallbackUpdateNormPrefs, 0);
  createTrackbar("max area", "stickletrack_tracking_prefs", &Prefs.contours_maxarea, 100, &trackbarCallbackUpdateNormPrefs, 0);
  createTrackbar("min circum", "stickletrack_tracking_prefs", &Prefs.contours_mincircum, 100, &trackbarCallbackUpdateNormPrefs, 0);
  createTrackbar("max circum", "stickletrack_tracking_prefs", &Prefs.contours_maxcircum, 100, &trackbarCallbackUpdateNormPrefs, 0);
  createTrackbar("min shortaxis", "stickletrack_tracking_prefs", &Prefs.contours_minshortaxis, 100, &trackbarCallbackUpdateNormPrefs, 0);
  createTrackbar("max shortaxis", "stickletrack_tracking_prefs", &Prefs.contours_maxshortaxis, 100, &trackbarCallbackUpdateNormPrefs, 0);
  createTrackbar("min longaxis", "stickletrack_tracking_prefs", &Prefs.contours_minlongaxis, 100, &trackbarCallbackUpdateNormPrefs, 0);
  createTrackbar("max longaxis", "stickletrack_tracking_prefs", &Prefs.contours_maxlongaxis, 100, &trackbarCallbackUpdateNormPrefs, 0);
  createTrackbar("max speed", "stickletrack_tracking_prefs", &Prefs.contours_maxspeed, 100, &trackbarCallbackUpdateNormPrefs, 0);
  createTrackbar("max rotation speed", "stickletrack_tracking_prefs", &Prefs.contours_maxrot, 100, &trackbarCallbackUpdateNormPrefs, 0);

  cvSetMouseCallback("stickletrack_tracking_screen", mouseTracking, 0);

  mContours = Mat::zeros(Props.height, Props.width, CV_8UC3);
}

double kitvalue (double min, double max) {
  return min + pow( cos(frametime*10), 2)*(max-min);
}

void findFishBlobbs(vector<fishblobb>& fishblobbs, vector< vector<Point> >& contours, vector<Vec4i>& hierarchy) {
  for (int icontour=0; icontour >= 0; icontour = hierarchy[icontour][0]) {
    fishblobb newblobb;

    if ( contours[icontour].size() < 5 )
      continue;

    newblobb.circum = contours[icontour].size();
    if ( newblobb.circum < normalPrefs.contours_mincircum || newblobb.circum > normalPrefs.contours_maxcircum )
      continue;

    newblobb.area = contourArea(contours[icontour]);
    if ( newblobb.area < normalPrefs.contours_minarea || newblobb.area > normalPrefs.contours_maxarea )
      continue;

    RotatedRect minEllipse = fitEllipse(contours[icontour]);
    newblobb.longaxis = minEllipse.size.height;
    newblobb.shortaxis = minEllipse.size.width;
    if ( newblobb.shortaxis < normalPrefs.contours_minshortaxis || newblobb.shortaxis > normalPrefs.contours_maxshortaxis
	 || newblobb.longaxis < normalPrefs.contours_minlongaxis || newblobb.longaxis > normalPrefs.contours_maxlongaxis )
      continue;
     
    newblobb.x = minEllipse.center.x;
    newblobb.y = minEllipse.center.y;
     
    RotatedRect minRect;
    Mat mFish = Mat::zeros(Props.height, Props.width, CV_8U);
    Moments fishMom;
    Point2f massCenter;
    Point2f recHead, recTail;
	
    newblobb.identified = false;
    newblobb.points = &contours[icontour];

    newblobb.contourid = icontour;

    /* find head/tail sensitive angle
       maybe the method could be improved */
    drawContours( mFish, contours, icontour, Scalar(255, 255, 255), CV_FILLED );
    fishMom = moments(mFish);
      
    massCenter = Point2f( fishMom.m10/fishMom.m00 , fishMom.m01/fishMom.m00 );
    minRect = minAreaRect(contours[ icontour ]);
      
    recHead = Point2f( minRect.center.x + newblobb.longaxis*cos(minEllipse.angle * M_PI / 180 + M_PI/2), 
		       minRect.center.y + newblobb.longaxis*sin(minEllipse.angle * M_PI / 180 + M_PI/2) );
    recTail = Point2f( minRect.center.x + newblobb.longaxis*cos(minEllipse.angle * M_PI / 180 - M_PI/2), 
		       minRect.center.y + newblobb.longaxis*sin(minEllipse.angle * M_PI / 180 - M_PI/2) );
      
    if ( pow( recHead.x - massCenter.x, 2 ) + pow( recHead.y - massCenter.y, 2 ) > pow( recTail.x - massCenter.x, 2 ) + pow( recTail.y - massCenter.y, 2 ) ) {
      newblobb.angle = minEllipse.angle + 180;
	
      if (newblobb.angle > 360)
	newblobb.angle = newblobb.angle - 360;
    }
    else {
      newblobb.angle = minEllipse.angle;
    }

    newblobb.tail = Point2f( minEllipse.center.x + newblobb.longaxis*cos(newblobb.angle * M_PI / 180 - M_PI/2),
			     minEllipse.center.y + newblobb.longaxis*sin(newblobb.angle * M_PI / 180 - M_PI/2) );
    newblobb.head = Point2f( minEllipse.center.x + newblobb.longaxis*cos(newblobb.angle * M_PI / 180 + M_PI/2), 
			     minEllipse.center.y + newblobb.longaxis*sin(newblobb.angle * M_PI / 180 + M_PI/2) );

    fishblobbs.push_back(newblobb);
  }
}

void tagCare(vector<tag>& tags) {
  while ( tags.size() > Prefs.manyfish )
    tags.pop_back();

  while ( tags.size() < Prefs.manyfish ) {
    tag tmp;
    tmp.x = 0;
    tmp.y = 0;
    tmp.propZ = 0;
    tmp.lastseen = frametime;

    switch ( tags.size() ) {
    case 0:
      tmp.color = Scalar( 0, 255, 255 );
      break;
    case 1:
      tmp.color = Scalar( 255, 0, 255 );
      break;
    case 2:
      tmp.color = Scalar( 255, 255, 0 );
      break;
    case 3:
      tmp.color = Scalar( 255, 0, 0 );
      break;
    case 4:
      tmp.color = Scalar( 0, 255, 0 );
      break;
    case 5:
      tmp.color = Scalar( 0, 0, 255 );
      break;
    default:
      tmp.color = Scalar( rand()%255, rand()%255, rand()%255 );
    }

    tmp.angle = 0;
    tmp.area = 0;
    tmp.circum = 0;
    tmp.shortaxis = 0;
    tmp.longaxis = 0;
    tmp.virgin = true;
    tags.push_back(tmp);
  }
}

void matchBlobbsNTags(vector<tag>& tags, vector<fishblobb>& fishblobbs) {
  vector<blobbcost> blobbcosts;

  for (int iblobb = 0; iblobb < fishblobbs.size(); iblobb++) {
    double sepcost[tags.size()][6];
    double maxcost[6];
    blobbcost blobbCost;

    blobbCost.blobbid = iblobb;

    for (int imax = 0; imax < 6; imax++)
      maxcost[imax] = 0;

    for (int itag = 0; itag < tags.size(); itag++) {
      double anglediff = fabs (fishblobbs[iblobb].angle - tags[itag].angle);
	    
      if ( anglediff > 180 )
	anglediff = 360 - anglediff;

      sepcost[itag][SEPCOST_ANGLE] = anglediff;

      sepcost[itag][SEPCOST_AREA] = fabs ( fishblobbs[iblobb].area - tags[itag].area );
      sepcost[itag][SEPCOST_CIRCUM] = fabs ( fishblobbs[iblobb].circum - tags[itag].circum );
      sepcost[itag][SEPCOST_SHORTAXIS] = fabs ( fishblobbs[iblobb].shortaxis - tags[itag].shortaxis );
      sepcost[itag][SEPCOST_LONGAXIS] = fabs ( fishblobbs[iblobb].longaxis - tags[itag].longaxis );
      sepcost[itag][SEPCOST_DIST] = sqrt( pow(tags[itag].x - fishblobbs[iblobb].x, 2) + pow(tags[itag].y - fishblobbs[iblobb].y, 2) );

      for (int imax = 0; imax < 6; imax++)
	if ( maxcost[imax] < sepcost[itag][imax] )
	  maxcost[imax] = sepcost[itag][imax];
    }

    for (int itag = 0; itag < tags.size(); itag++) {
      tagcost newTagCost;
	    
      newTagCost.tagid = itag;
      newTagCost.tagcost = 0;

      if ( sepcost[itag][SEPCOST_DIST] > ( frametime - tags[itag].lastseen )*normalPrefs.contours_maxspeed 
	   || sepcost[itag][SEPCOST_ANGLE] > ( frametime - tags[itag].lastseen )*normalPrefs.contours_maxrot ) 
	continue;

      for (int isep = 0; isep < 6; isep++)
	newTagCost.tagcost += sepcost[itag][isep] / maxcost[isep];

      blobbCost.tagcosts.push_back( newTagCost );						
    }

    if ( blobbCost.tagcosts.size() > 0 ) { 
      sort( blobbCost.tagcosts.begin(), blobbCost.tagcosts.end(), sort_tags );
      blobbcosts.push_back( blobbCost );
    }
  }
	
  for (int itag = 0; itag < tags.size() && blobbcosts.size() > 0; itag++) {
    sort( blobbcosts.begin(), blobbcosts.end(), sort_blobbs );

    int tagnow = blobbcosts[0].tagcosts[0].tagid;

    tags[tagnow].virgin = false;
    tags[tagnow].x = fishblobbs[ blobbcosts[0].blobbid ].x;
    tags[tagnow].y = fishblobbs[ blobbcosts[0].blobbid ].y;
    tags[tagnow].shortaxis = fishblobbs[ blobbcosts[0].blobbid ].shortaxis;
    tags[tagnow].longaxis =  fishblobbs[ blobbcosts[0].blobbid ].longaxis;
    tags[tagnow].angle = fishblobbs[ blobbcosts[0].blobbid ].angle;
    tags[tagnow].lastseen = frametime;
    tags[tagnow].area = fishblobbs[ blobbcosts[0].blobbid ].area;
    tags[tagnow].circum = fishblobbs[ blobbcosts[0].blobbid ].circum;
    tags[tagnow].head = fishblobbs[ blobbcosts[0].blobbid ].head;
    tags[tagnow].tail = fishblobbs[ blobbcosts[0].blobbid ].tail;

    blobbcosts.erase( blobbcosts.begin() );

    for (int iblobb = 0; iblobb < blobbcosts.size(); iblobb++) {
	    
      for (int itag = 0; itag < blobbcosts[iblobb].tagcosts.size(); itag++)
	if ( blobbcosts[iblobb].tagcosts[itag].tagid == tagnow ) {
	  blobbcosts[iblobb].tagcosts.erase( blobbcosts[iblobb].tagcosts.begin() + itag );
	  break;
	}
	   
      if ( blobbcosts[iblobb].tagcosts.size() == 0 ) {
	blobbcosts.erase( blobbcosts.begin() + iblobb );
	iblobb--;
      }
    }
  }
}

void drawTags (const vector<tag>& tags) {
  RotatedRect rRect;
  Point2f vertices[4];

  for (int itag = 0; itag < tags.size(); itag++) {
    double recside = sqrt( kitvalue( normalPrefs.contours_minarea, normalPrefs.contours_maxarea) );
    rRect = RotatedRect( Point(tags[itag].x, tags[itag].y), Size2f(recside, recside), tags[itag].angle );
    rRect.points(vertices);
    for (int i = 0; i < 4; i++)
      line(mContours, vertices[i], vertices[(i+1)%4], Scalar(255,255,255), Props.diagonal/1000.0);

    ellipse(mContours, Point(tags[itag].x, tags[itag].y), 
	    Size( kitvalue( normalPrefs.contours_minshortaxis, normalPrefs.contours_maxshortaxis ), 
		  kitvalue( normalPrefs.contours_minlongaxis, normalPrefs.contours_maxlongaxis ) ), 
	    tags[itag].angle, 0, 360, Scalar(255,255,255), Props.diagonal/1000.0, 8);

    circle(mContours, Point(tags[itag].x, tags[itag].y), kitvalue( normalPrefs.contours_mincircum, normalPrefs.contours_maxcircum )/(2*M_PI), Scalar(255,255,255), 
	   Props.diagonal/1000.0, 8);
      
    if ( ! tags[itag].virgin ) {
      float searchradius = ( frametime - tags[itag].lastseen )*normalPrefs.contours_maxspeed;
      float searchangle = ( frametime - tags[itag].lastseen )*normalPrefs.contours_maxrot;
      ellipse(mContours, Point(tags[itag].x, tags[itag].y), Size(searchradius, searchradius), tags[itag].angle, 90-searchangle, 90+searchangle, 
	      tags[itag].color, Props.diagonal/1000.0, 8);
      
      rRect = RotatedRect( Point(tags[itag].x, tags[itag].y), Size2f(sqrt(tags[itag].area), sqrt(tags[itag].area)), tags[itag].angle );
      rRect.points(vertices);
      for (int i = 0; i < 4; i++)
	line(mContours, vertices[i], vertices[(i+1)%4], tags[itag].color, 2);
      
      ellipse(mContours, Point(tags[itag].x, tags[itag].y), 
	      Size( tags[itag].shortaxis, tags[itag].longaxis ), tags[itag].angle, 0, 360, tags[itag].color, 2*Props.diagonal/1000.0, 8);

      circle(mContours, Point(tags[itag].x, tags[itag].y), tags[itag].circum / (2*M_PI), tags[itag].color, 2*Props.diagonal/1000.0, 8);
      
      circle(mContours, tags[itag].head, 5*Props.diagonal/1000.0, tags[itag].color, -1, 8);
    }
  }
}

void tracking_locateTags (vector<tag>& tags, Mat combinedmask_contour) {
  mContours = Mat::zeros(Props.height, Props.width, CV_8UC3);

  if (enablecontours) {
    vector< vector<Point> > contours;
    vector<Vec4i> hierarchy;
    vector <fishblobb> fishblobbs;

    findContours( combinedmask_contour, contours, hierarchy, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_SIMPLE );

    if ( contours.size() > 0 ) {

      for(int idx = 0 ; idx >= 0; idx = hierarchy[idx][0] )
	drawContours(mContours, contours, idx, Scalar(255, 255, 255), CV_FILLED, 8, hierarchy);

      findFishBlobbs(fishblobbs, contours, hierarchy);

      for ( int iblobb = 0; iblobb < fishblobbs.size(); iblobb++) {
	Scalar color( rand()%255, rand()%255, rand()%255 );
	drawContours( mContours, contours, fishblobbs[iblobb].contourid, color, CV_FILLED, 8, hierarchy );
      }

      tagCare(tags);
      matchBlobbsNTags(tags, fishblobbs);

    }

    drawTags(tags);
  }
}