// Listing 15.3. Obliczanie nienależących do przekątnej elementów modelu kowariancji
#include <opencv2/opencv.hpp>
#include <vector>
#include <iostream>
#include <cstdlib>
#include <fstream>

using namespace std;

vector<cv::Mat> planes(3);
vector<cv::Mat> sums(3);
vector<cv::Mat> xysums(6);
cv::Mat sum, sqsum;
int image_count = 0;

// funkcja akumulująca
// informacje potrzebne do obliczenia wariancji
//
void accumulateVariance(
	cv::Mat& I) {
	if( sum.empty() ) {
		sum = cv::Mat::zeros( I.size(), CV_32FC(I.channels()) );
		sqsum = cv::Mat::zeros( I.size(), CV_32FC(I.channels()) );
		image_count = 0;
	}
	cv::accumulate( I, sum );
	cv::accumulateSquare( I, sqsum );
	image_count++;
}

// funkcja obliczania wariancji:
// (variance to sigma^2)
//
void computeVariance(
	cv::Mat& variance) {
	double one_by_N = 1.0 / image_count;
	variance = (one_by_N * sqsum) - ((one_by_N * one_by_N) * sum.mul(sum));
}

// to samo, co powyżej, ale oblicza odchylenie standardowe
void computeStdev(
	cv::Mat& std__) {
	double one_by_N = 1.0 / image_count;
	cv::sqrt(((one_by_N * sqsum) -((one_by_N * one_by_N) * sum.mul(sum))), std__);
}

// And avg images
void computeAvg(
	cv::Mat& av) {
	double one_by_N = 1.0 / image_count;
	av = one_by_N * sum;
}
	
// ===================================================================//


void accumulateCovariance(
cv::Mat& I
) {
	int i, j, n;

	if( sum.empty() ) {
		image_count = 0;
		for( i=0; i<3; i++ ) {
			// sumy r, g i b
			sums[i]
			= cv::Mat::zeros( I.size(), CV_32FC1 );
		}
		for( n=0; n<6; n++ ) {
			// elementy rr, rg, rb, gg, gb i bb
			xysums[n] = cv::Mat::zeros( I.size(), CV_32FC1 );
		}
	}
	cv::split( I, planes );
	for( i=0; i<3; i++ ) {
		cv::accumulate( planes[i], sums[i] );
	}
	n = 0;
	for( i=0; i<3; i++ ) {
		// wiersz sigmy
		for( j=i; j<3; j++ ) {
			// kolumna sigmy
			n++;
			cv::accumulateProduct( planes[i], planes[j], xysums[n] );
		}
	}
	image_count++;
}

// Odpowiednia funkcja obliczeń również będzie nieznacznym rozszerzeniem opisanej wcześniej analogicznej funkcji do obliczania wariancji
//
void computeCoariance(
	cv::Mat& covariance
	// 6-kanałowa tablica, kanały to elementy rr, rg, rb, gg, gb oraz bb tablicy Sigma_xy
	// elementy rr, rg, rb, gg, gb i bb z Sigma_xy
) {
	double one_by_N = 1.0 / image_count;
	
	// użycie tablic xysum do przechowywania indywidualnych elementów
	//
	int n = 0;
	for( int i=0; i<3; i++ ) {
		// wiersz sigmy
		for( int j=i; j<3; j++ ) {
			// kolumna sigmy
			n++;
			xysums[n] = (one_by_N * xysums[n])
			- ((one_by_N * one_by_N) * sums[i].mul(sums[j]));
		}
	}
	
	// złożenie 6 indywidualnych elementów w 6-kanałową tablicę
	//
	cv::merge( xysums, covariance );
}

////////////////////////////////////////////////////////////////////////
/////////////tilities to run///////////////////////////////////////////

void help(char** argv ) {
	cout << "\n"
	<< "Oblicza mean i std na <#frames to train on> klatkach filmu wejściowego, następnie uruchamia model\n"
	<< argv[0] <<" <#frames to train on> <avi_path/filename>\n"
	<< "Na przykład:\n"
	<< argv[0] << " 50 ../tree.avi\n"
	<< "a do regulacji progów, esc, q i Q do kończenia"
	<< endl;
}
	
//////////////  Kod pożyczony z listingu 15.02  //////////////////////

// magazyn globalny
//
// zmiennoprzecinkowe, trzykanałowe obrazy
//
cv::Mat image; // klatka filmu
cv::Mat IavgF, IdiffF, IhiF, IlowF; // próg
cv::Mat tmp, mask;

// zmiennoprzecinkowe jednokanałowe obrazy
//
vector<cv::Mat> Igray(3); 
vector<cv::Mat> Ilow(3);
vector<cv::Mat> Ihi(3);

// jednokanałowy obraz bajtowy
//
cv::Mat Imaskt; // tymczasowa maska

// progi
//
float high_thresh = 21.0;  // skalowanie progów w backgroundDiff()
float low_thresh = 2.0;    //

// I jest tylko przykładowym obrazem na potrzeby alokacji
// (stanowi wzór rozmiaru)

//
void AllocateImages( const cv::Mat& I ) {
	cv::Size sz = I.size();
	IavgF = cv::Mat::zeros(sz, CV_32FC3 ); 
	IdiffF = cv::Mat::zeros(sz, CV_32FC3 ); 
	IhiF = cv::Mat::zeros(sz, CV_32FC3 ); 
	IlowF = cv::Mat::zeros(sz, CV_32FC3 );
	tmp = cv::Mat::zeros( sz, CV_32FC3 ); 
	Imaskt = cv::Mat( sz, CV_32FC1 ); 
}


void setHighThreshold( float scale ) {
	IhiF = IavgF + (IdiffF * scale);
	cv::split( IhiF, Ihi );
}

void setLowThreshold( float scale ) {
	IlowF = IavgF - (IdiffF * scale);
	cv::split( IlowF, Ilow );
}

void createModelsfromStats() {
	// wartość IavgF jest już ustawiona
	// wartość IdiffF obrazem odchylenia standardowego
	
	// diff zawsze musi mieć jakąś wartość
	//
	IdiffF += cv::Scalar( 0.1, 0.1, 0.1 );
	setHighThreshold( high_thresh);
	setLowThreshold( low_thresh);
}


// tworzy binarną maskę 0,255, w której 255 oznacza piksel pierwszego planu
// I jest obrazem wejściowym, 3 kanały, 8u
// Imask to obraz maski, który ma zostać utworzony, 1 kanał, 8u
//
void backgroundDiff(
		cv::Mat& I,
		cv::Mat& Imask) {
	
	I.convertTo( tmp, CV_32F ); // na typ zmiennoprzecinkowy
	cv::split( tmp, Igray );
	
	// kanał 1
	//
	cv::inRange( Igray[0], Ilow[0], Ihi[0], Imask );

	// kanał 2
	//
	cv::inRange( Igray[1], Ilow[1], Ihi[1], Imaskt );
	Imask = cv::min( Imask, Imaskt );

	// kanał 3
	//
	cv::inRange( Igray[2], Ilow[2], Ihi[2], Imaskt );
	Imask = cv::min( Imask, Imaskt );

	// na koniec odwracamy wyniki
	//
	Imask = 255 - Imask;
}


void showForgroundInRed( char** argv, const cv::Mat &img) {
		cv::Mat rawImage;
		cv::split( img, Igray );
		Igray[2] = cv::max( mask, Igray[2] );
		cv::merge( Igray, rawImage );
		cv::imshow( argv[0], rawImage );
		cv::imshow("Segmentation", mask);
}

void adjustThresholds(char** argv, cv::Mat &img) {
	int key = 1;
	while((key = cv::waitKey()) != 27 && key != 'Q' && key != 'q')  // Esc or Q or q to exit
	{
		if(key == 'L') { low_thresh += 0.2;}
		if(key == 'l') { low_thresh -= 0.2;}	
		if(key == 'H') { high_thresh += 0.2;}
		if(key == 'h') { high_thresh -= 0.2;}
		cout << "H or h, L or l, esq or q to quit;  high_thresh = " << high_thresh << ", " << "low_thresh = " << low_thresh << endl;
		setHighThreshold(high_thresh);
		setLowThreshold(low_thresh);
		backgroundDiff(img, mask);
		showForgroundInRed(argv, img);
	}
}

int main( int argc, char** argv) {
	cv::namedWindow( argv[0], cv::WINDOW_AUTOSIZE );
	cv::VideoCapture cap;
	if((argc < 3)|| !cap.open(argv[2])) {
		cerr << "Nie udało się uruchomić programu" << endl;
		help(argv);
		cap.open(0);
		return -1;
	}
	int number_to_train_on = atoi( argv[1] );

	// PIERWSZA PĘTLA PRZETWARZANIA (SZKOLENIE):
	//
	int image_count = 0;
	int key;
	bool first_frame = true;
	cout << "Total frames to train on = " << number_to_train_on << endl; //db
	while(1) {
		cout << "frame#: " << image_count << endl;
		cap >> image;
		if( !image.data ) exit(1); // coś się nie udało
		if(image_count == 0) AllocateImages( image ); 
		accumulateVariance(image);
		cv::imshow( argv[0], image );
		image_count++;
		if( (key = cv::waitKey(7)) == 27 || key == 'q' || key == 'Q' || image_count >= number_to_train_on) break; //Allow early exit on space, esc, q
	}

	// zakończono szkolenie, czas na utworzenie modeli
	//
	cout << "Tworzenie modelu tła" << endl;
	computeAvg(IavgF);
	computeStdev(IdiffF);
	createModelsfromStats();
	cout << "Skończone!  Naciśnij dowolny klawisz, aby kontynuować. Naciśnij klawisz a lub A, aby ustawić progi lub esq, q albo Q, aby zakończyć.\n" << endl;

	// DRUGA PĘTLA PRZETWARZANIA (TESTOWANIE):
	//
	cv::namedWindow("Segmentation", cv::WINDOW_AUTOSIZE ); // dla obrazu maski
	while((key = cv::waitKey()) != 27 || key == 'q' || key == 'Q'  ) { // esc, q lub Q, aby zakończyć
		cap >> image;
		if( !image.data ) exit(0);
		cout <<  image_count++ << endl;
		backgroundDiff( image, mask );
		cv::imshow("Segmentation", mask);
		
		// prosta wizualizacja, polegająca na zapisie w kanale czerwonym
		//
		showForgroundInRed( argv, image);
		if(key == 'a') {
			cout << "H lub h == wysoki próg w górę lub w dół; L lub l == niski próg w górę lub w dół." << endl;
			cout << " esc, q lub Q kończy program" << endl;
			adjustThresholds(argv, image);
			cout << "Zakończono adjustThreshold, powrót do przeglądania klatek, esc, q lub Q, aby zakończyć." << endl;
		}
	}
	exit(0);
}

	
