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

#include <iostream>

static Mat background, foregroundmask, colormask, mScissors, scissorsmask;
static double Z = 1;
static bool recomputeScissors = true;

void mouseScissors(int evt, int x, int y, int flags, void* param){
  if (evt==CV_EVENT_LBUTTONDOWN) {
    Prefs.scissorpoints->push_back( Point(x, y) );
    recomputeScissors = true;
  }

  else if (evt==CV_EVENT_RBUTTONDOWN) {
    int nearestid = -1;
    double nearestdist;

    for ( int ipoint = 0; ipoint < Prefs.scissorpoints->size(); ipoint++ ) {
      double dist = pow( x - (*Prefs.scissorpoints)[ipoint].x, 2 ) + pow( y - (*Prefs.scissorpoints)[ipoint].y, 2 );

      if ( dist < nearestdist || nearestid == -1 ) {
	nearestdist = dist;
	nearestid = ipoint;
      }
    }
    if ( nearestid != -1 ) {
      Prefs.scissorpoints->erase( Prefs.scissorpoints->begin() + nearestid );
      recomputeScissors = true;
    }
  }
}

void masking_init() {
  background = Mat::zeros(Props.height, Props.width, CV_32FC3);
  
  namedWindow("stickletrack_masking_background", CV_WINDOW_KEEPRATIO);
  namedWindow("stickletrack_masking_backgroundmask", CV_WINDOW_KEEPRATIO);
  namedWindow("stickletrack_masking_colormask", CV_WINDOW_KEEPRATIO);
  namedWindow("stickletrack_masking_scissorsmask", CV_WINDOW_KEEPRATIO);

  createTrackbar("Decay", "stickletrack_masking_background", &Prefs.halfdecay, 100, &trackbarCallbackUpdateNormPrefs, 0);
  createTrackbar("Threshold", "stickletrack_masking_backgroundmask", &Prefs.forethreshold, 255, &trackbarCallbackUpdateNormPrefs, 0);

  createTrackbar("min H", "stickletrack_masking_colormask", &Prefs.mincolor[0], 255, NULL, 0);
  createTrackbar("min S", "stickletrack_masking_colormask", &Prefs.mincolor[1], 255, NULL, 0);
  createTrackbar("min V", "stickletrack_masking_colormask", &Prefs.mincolor[2], 255, NULL, 0);
  createTrackbar("max H", "stickletrack_masking_colormask", &Prefs.maxcolor[0], 255, NULL, 0);
  createTrackbar("max S", "stickletrack_masking_colormask", &Prefs.maxcolor[1], 255, NULL, 0);
  createTrackbar("max V", "stickletrack_masking_colormask", &Prefs.maxcolor[2], 255, NULL, 0);

  cvSetMouseCallback("stickletrack_masking_scissorsmask", mouseScissors, 0);
}

void computeBackground (const Mat& origframe, Mat& background, double& Z, const double& decayfac) {
  background = ( exp(-decayfac) * background * Z + origframe ) / ( exp(-decayfac) * Z + 1 );
  Z = exp(-decayfac) * Z + 1;

  if ( ! isWindowClosed("stickletrack_masking_background") )
    imshow("stickletrack_masking_background", background/255.0);
}

void computeForegroundMask (const Mat& origframe, const Mat& background, Mat& foregroundMask) {
  absdiff(origframe, background, foregroundMask);
  cvtColor(foregroundMask, foregroundMask, CV_RGB2GRAY);
  threshold(foregroundMask, foregroundMask, Prefs.forethreshold, 255, CV_THRESH_BINARY);
  foregroundMask.convertTo(foregroundMask, CV_8UC3);

  if ( ! isWindowClosed("stickletrack_masking_backgroundmask") )
    imshow("stickletrack_masking_backgroundmask", foregroundmask);
}

void computeColorMask (const Mat& origframe, Mat& colormask) {
  origframe.convertTo(colormask, CV_8UC3);
  cvtColor(colormask, colormask, CV_RGB2HSV);
  inRange(colormask, Scalar(Prefs.mincolor[0],Prefs.mincolor[1],Prefs.mincolor[2]), Scalar(Prefs.maxcolor[0], Prefs.maxcolor[1], Prefs.maxcolor[2]), colormask);
  colormask.convertTo(colormask, CV_8UC3);

  if ( ! isWindowClosed("stickletrack_masking_colormask") )
    imshow("stickletrack_masking_colormask", colormask);
}

void computeScissors (const Mat& origframe, Mat& mScissors) {
  origframe.convertTo(mScissors, CV_8UC3);
  for (int ipoint = 0; ipoint < Prefs.scissorpoints->size(); ipoint++) {
    circle(mScissors, (*Prefs.scissorpoints)[ipoint], 5, Scalar(0,0,255), -1, 8);
    line(mScissors, (*Prefs.scissorpoints)[ipoint], (*Prefs.scissorpoints)[ (ipoint+1)%Prefs.scissorpoints->size() ], Scalar(0,0,255), 1, 8);
  }
}

void computeScissorsMask (Mat& scissorsmask) {
  if ( Prefs.scissorpoints->size() >= 3) {
    const Point* elementPoints[1] = { &((*Prefs.scissorpoints)[0]) };
    int numberOfPoints = (int)Prefs.scissorpoints->size();
    scissorsmask = Mat::zeros(Props.height, Props.width, CV_8U);
    fillPoly (scissorsmask, elementPoints, &numberOfPoints, 1, Scalar (255, 255, 255), 8);
  }
  else
    scissorsmask = Mat::ones(Props.height, Props.width, CV_8U);
}

void masking_getCombinedMask(const Mat& origframe, Mat& combinedmask) {
  double decayfac = (log(2) / normalPrefs.halfdecay) / Props.fps;

  computeBackground(origframe, background, Z, decayfac);
  computeForegroundMask(origframe, background, foregroundmask );

  computeColorMask(origframe, colormask);

  computeScissors(origframe, mScissors);

  if ( recomputeScissors ) {
    computeScissorsMask(scissorsmask);
    recomputeScissors = false;
  }

  if ( ! isWindowClosed("stickletrack_masking_scissorsmask") )
    imshow("stickletrack_masking_scissorsmask", mScissors);
  
  bitwise_and(foregroundmask, colormask, combinedmask);
  bitwise_and(combinedmask, scissorsmask, combinedmask);
}