// Listing 15.4. Implementacja algorytmu ksiąg kodów
#include <opencv2/opencv.hpp>
#include <vector>
#include <iostream>
#include <cstdlib>
#include <fstream>

using namespace std;

#define CHANNELS 3			// zawsze trzy, ponieważ yuv
int cbBounds[CHANNELS];		// jeśli piksel mieści się w tych granicach poza księgą kodów, użyj go, w przeciwnym razie utwórz nowy kod
int minMod[CHANNELS];		// jeśli piksel jest niższy niż księga kodów o tę ilość, jest dopasowany
int maxMod[CHANNELS];		// jeśli piksel jest wyższy niż księga kodów o tę ilość, jest dopasowany


// Zmienna t rejestruje liczbę punktów zgromadzonych od początku lub od ostatniej operacji czyszczenia
//
class CodeElement {
	public:
		uchar learnHigh[CHANNELS]; // górna granica uczenia
		uchar learnLow[CHANNELS];  // dolna granica uczenia
		uchar max[CHANNELS];       // górny bok pola
		uchar min[CHANNELS];       // dolny bok pola
		int t_last_update;         // umożliwi nam kasowanie przestarzałych elementów
		int stale;                 // maksymalny czas nieaktywności

		CodeElement() {
			for(int i = 0; i < CHANNELS; i++)
				learnHigh[i] = learnLow[i] = max[i] = min[i] = 0;
			t_last_update = stale = 0;
		}
		
		CodeElement& operator=( const CodeElement& ce ) {
			for(int i=0; i<CHANNELS; i++ ) {
				learnHigh[i] = ce.learnHigh[i];
				learnLow[i] = ce.learnLow[i];
				min[i] = ce.min[i];
				max[i] = ce.max[i];
			}
			t_last_update = ce.t_last_update;
			stale = ce.stale;
			return *this;
		}
		
		CodeElement( const CodeElement& ce ) { *this = ce; }
};


class CodeBook : public vector<CodeElement> {
	public:
	int t;     // Count of every image learned on
	
	// liczy wszystkie operacje dostępu
	CodeBook() { t=0; }
	
	// domyślna jest pusta księga
	CodeBook( int n ) : vector<CodeElement>(n) { t=0; } //  tworzy księgę o rozmiarze n
};


// aktualizuje element księgi kodów o nowy punkt danych
// uwaga: cbBounds musi mieć długość równą numChannels
//
//
int updateCodebook(   // zwraca indeks w CodeBook
	const cv::Vec3b& p, // wejściowy piksel YUV
	CodeBook& c,        // CodeBook dla piksela
	unsigned* cbBounds, // granice księgi kodów (zazwyczaj: {10,10,10})
	int numChannels     // liczba interesujących nas kanałów kolorów
	) {
	if(c.size() == 0) 
		c.t = 0;
	c.t += 1;		// rejestruje zdarzenie uczenia
	// ustawienia górnych i dolnych granic
	unsigned int high[3], low[3], n;
	for( n=0; n<numChannels; n++ ) {
		high[n] = p[n] + *(cbBounds+n); 
		if( high[n] > 255 ) high[n] = 255;
		low[n] = p[n] - *(cbBounds+n);
		if( low[n] < 0) low[n] = 0;
	}
	
	// sprawdzenie, czy pasuje do istniejącego kodu
	//
	int i;
	int matchChannel;
	for( i=0; i<c.size(); i++ ) {
		matchChannel = 0;
		for( n=0; n<numChannels; n++ ) {
			if( // znaleziono element dla tego kanału
				( c[i].learnLow[n] <= p[n] ) && ( p[n] <= c[i].learnHigh[n]))
				matchChannel++;
			}
		
		if( matchChannel == numChannels ) {// jeśli znaleziono element
			c[i].t_last_update = c.t; 
			
			// regulacja tego kodu dla pierwszego kanału
			//
			for( n=0; n<numChannels; n++ ) {
				if( c[i].max[n] < p[n] )
					c[i].max[n] = p[n];
				else if( c[i].min[n] > p[n] ) 
					c[i].min[n] = p[n];
			}
			break;
		}
	}
	
	// NARZUT NA SZUKANIE POTENCJALNIE PRZESTARZAŁYCH ELEMENTÓW
	//
	for( int s=0; s<c.size(); s++ ) {
	
		// szuka starzejących się elementów księgi kodów:
		//
		int negRun = c.t - c[s].t_last_update;
		if( c[s].stale < negRun ) c[s].stale = negRun;
	}
	
	// W RAZIE POTRZEBY DODAJE NOWY KOD
	//
	if( i == c.size() ) {
		// jeśli nie znajdzie żadnego kodu, tworzy jeden
		CodeElement ce;
		for( n=0; n<numChannels; n++ ) {
			ce.learnHigh[n] = high[n];
			ce.learnLow[n] = low[n];
			ce.max[n] = p[n];
			ce.min[n] = p[n];
		}
			ce.t_last_update = c.t;
			ce.stale = 0;
			c.push_back( ce );
	}
	
	// POWOLNE KORYGOWANIE GRANIC UCZENIA
	//
	for( n=0; n<numChannels; n++ ) {
		if( c[i].learnHigh[n] < high[n]) c[i].learnHigh[n] += 1;
		if( c[i].learnLow[n] > low[n] ) c[i].learnLow[n] -= 1;
	}
	return c.size();
}

// Po upływie pewnego czasu nauki okresowo wywołujemy tę funkcję, aby skasować przestarzałe elementy
int foo = 0;
int clearStaleEntries(
	// zwraca liczbę usuniętych elementów
	CodeBook &c
	// księga do wyczyszczenia
){
	int staleThresh = c.t>>1;
	int *keep = new int[c.size()];
	int keepCnt = 0;
	
	// SPRAWDZENIE, KTÓRE ELEMENTY SĄ PRZESTARZAŁE
	//
	int foogo2 = 0;
	for( int i=0; i<c.size(); i++ ){
		if(c[i].stale > staleThresh) 
			keep[i] = 0; // oznaczenie do skasowania
		else
		{
			keep[i] = 1; // oznaczenie do zachowania
			keepCnt += 1;
		}
	}
	
	// przenosimy elementy, które chcemy zachować, na początek wektora, a następnie
	// po zapisaniu wszystkiego, co nas interesuje, obcinamy go do odpowiedniej długości
	//
	int k = 0;
	int numCleared = 0;
	for( int ii=0; ii<c.size(); ii++ ) {
		if( keep[ii] ) {
			c[k] = c[ii];
			// musimy odświeżyć te elementy dla następnej sesji
			c[k].t_last_update = 0;
			k++;
		} else {
			numCleared++;
		}
	}
	c.resize( keepCnt );
	delete[] keep;
	return numCleared;
}
	
// dla danego piksela i księgi kodów określa, czy ta księga obejmuje ten piksel
//
// UWAGI:
// minMod i maxMod muszą mieć długość numChannels, np. 3 kanały => minMod[3], maxMod[3] 
// dla każdego kanału jest po jednym progu minimalnym i maksymalnym
//
uchar backgroundDiff( // zwraca 0 => tło, 255 => pierwszy plan
	const cv::Vec3b& p, // piksel (YUV)
	CodeBook& c,        // księga kodów
	int numChannels,    // liczba testowanych kanałów
	int* minMod,        // dodaje tę (potencjalnie ujemną) liczbę do poziomu maksymalnego
						// podczas sprawdzania, czy nowy piksel należy do pierwszego planu
	int* maxMod         // odejmuje tę (potencjalnie ujemną) liczbę od poziomu minimalnego
						// podczas sprawdzania, czy nowy piksel należy do pierwszego planu
) {
	int matchChannel;
	
	// SPRAWDZA, CZY PASUJE DO ISTNIEJĄCEGO KODU
	//
	int i;
	for( i=0; i<c.size(); i++ ) {
		matchChannel = 0;
		for( int n=0; n<numChannels; n++ ) {
			if((c[i].min[n] - minMod_[n] <= p[n] ) && (p[n] <= c[i].max[n] + maxMod_[n]))
			{
				matchChannel++; // znaleziono element dla tego kanału
			} else {
				break;
			}
		}
		if(matchChannel == numChannels) {
			break; // znaleziono element pasujący do wszystkich kanałów
		}
	}
	if( i >= c.size() )	// brak dopasowania z codebook => foreground
		return 255;
	return 0; 			// w przeciwnym razie tło
}

///////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////// Ta część dodaje funkcję główną ////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////////

// klasa pomocnicza (zakłada, że obraz w filmie nie zmieni rozmiaru);
class CbBackgroudDiff {
	public:
	cv::Mat Iyuv;					// przechowuje przekonwertowany obraz
	cv::Mat mask;					// przechowuje maskę różnicy tła
	vector<CodeBook> codebooks;   	// przechowuje księgę kodó dla każdego piksela
	int row, col, image_length;		// ile jest pikseli w obrazie
	
	// konstruktor
	void init(cv::Mat &I_from_video) {
		vector<int> v(3,10);
		set_global_vecs(cbBounds, v);
		v[0] = 6; v[1] = 20; v[2] = 8; // domyślne wartości na dole
		set_global_vecs(minMod, v);
		v[0] = 70; v[1] = 20; v[2] = 6; // domyślne wartości na górze
		set_global_vecs(maxMod, v);
		Iyuv.create(I_from_video.size(), I_from_video.type());
		mask.create(I_from_video.size(), CV_8UC1);
		row = I_from_video.rows;
		col = I_from_video.cols;
		image_length = row*col;
		cout << "(row,col,len) = (" << row << ", " << col << ", " << image_length << ")" << endl;
		codebooks.resize(image_length);
	}
	
	CbBackgroudDiff(cv::Mat &I_from_video) {
		init(I_from_video);
	}
	
	CbBackgroudDiff(){};
	
	// konwersja na YUV
	void convert_to_yuv(cv::Mat &Irgb)
	{
		cvtColor(Irgb, Iyuv, cv::COLOR_BGR2YUV);
	}
	
	int size_check(cv::Mat &I) { // sprawdzenie czy obraz nie zmienił rozmiaru, zwrot -1, jeśli rozmiar nie pasuje, w przeciwnym razie 0
		int ret = 0;
		if((row != I.rows) || (col != I.cols)) {
			cerr << "BŁĄD: rozmiar się zmienił! Wcześniej [" << row << ", " << col << "], teraz [" << I.rows << ", " << I.cols << "]!" << endl;
			ret = -1;
		}
		return ret;
	}
	
	// narzędzia do ustawień globalnych
	void set_global_vecs(int *globalvec, vector<int> &vec) {
		if(vec.size() != CHANNELS) {
			cerr << "Wejściowy vec[" << vec.size() << "] powinien być równy CHANNELS [" << CHANNELS << "]" << endl;
			vec.resize(CHANNELS, 10);
		}
		int i = 0;
		for (vector<int>::iterator it = vec.begin(); it != vec.end(); ++it, ++i) {
			 globalvec[i] = *it;
		 }
	 }

	// operacje na tle
	int updateCodebookBackground(cv::Mat &Irgb) { // Learn codebook, -1 if error, else total # of codes
		convert_to_yuv(Irgb);
		if(size_check(Irgb))
			return -1;
		int total_codebooks = 0;
		cv::Mat_<cv::Vec3b>::iterator Iit = Iyuv.begin<cv::Vec3b>(), IitEnd = Iyuv.end<cv::Vec3b>();
		vector<CodeBook>::iterator Cit = codebooks.begin(), CitEnd = codebooks.end();
		for(; Iit != IitEnd; ++Iit,++Cit) {
			total_codebooks += updateCodebook(*Iit,*Cit,cbBounds,CHANNELS);
		}
		if(Cit != CitEnd)
			cerr << "ERROR: Cit != CitEnd in updateCodeBackground(...) " << endl;
		return(total_codebooks);
	}

	int clearStaleEntriesBackground() { //Clean out stuff that hasn't been updated for a long time
		int total_cleared = 0;
		vector<CodeBook>::iterator Cit = codebooks.begin(), CitEnd = codebooks.end();
		for(; Cit != CitEnd; ++Cit) {
			total_cleared += clearStaleEntries(*Cit);
		}
		return(total_cleared);
	}

	int backgroundDiffBackground(cv::Mat &Irgb) {  //Take the background difference of the image
		convert_to_yuv(Irgb);
		if(size_check(Irgb))
			return -1;
		cv::Mat_<cv::Vec3b>::iterator Iit = Iyuv.begin<cv::Vec3b>(), IitEnd = Iyuv.end<cv::Vec3b>();
		vector<CodeBook>::iterator Cit = codebooks.begin(), CitEnd = codebooks.end();
		cv::Mat_<uchar>::iterator Mit = mask.begin<uchar>(), MitEnd = mask.end<uchar>();
		for(; Iit != IitEnd; ++Iit,++Cit,++Mit) {
			*Mit = backgroundDiff(*Iit,*Cit,CHANNELS,minMod,maxMod);
		}
		if((Cit != CitEnd)||(Mit != MitEnd)){
			cerr << "ERROR: Cit != CitEnd and, or Mit != MitEnd in updateCodeBackground(...) " << endl;
			return -1;
		}
		return 0;
	}
}; // end CbBackgroudDiff


void help(char** argv ) {
	cout << "\n"
	<< "Train a codebook background model on the first <#frames to train on> frames of an incoming video, then run the model\n"
	<< argv[0] <<" <#frames to train on> <avi_path/filename>\n"
	<< "For example:\n"
	<< argv[0] << " 50 ../tree.avi\n"
	<< "'A' or 'a' to adjust thresholds, esc, 'q' or 'Q' to quit"
	<< endl;
}

//Adjusting the distance you have to be on the low side (minMod) or high side (maxMod) of a codebook
//to be considered as recognized/background
//
void adjustThresholds(char* charstr, cv::Mat &Irgb, CbBackgroudDiff &bgd) {
	int key = 1;
	int y = 1,u = 0,v = 0, index = 0, thresh = 0;
	if(thresh)
		cout << "yuv[" << y << "][" << u << "][" << v << "] maxMod active" << endl;
	else
		cout << "yuv[" << y << "][" << u << "][" << v << "] minMod active" << endl;
	cout << "minMod[" << minMod[0] << "][" << minMod[1] << "][" << minMod[2] << "]" << endl;
	cout << "maxMod[" << maxMod[0] << "][" << maxMod[1] << "][" << maxMod[2] << "]" << endl;
	while((key = cv::waitKey()) != 27 && key != 'Q' && key != 'q')  // Esc or Q or q to exit
	{
		if(thresh)
			cout << "yuv[" << y << "][" << u << "][" << v << "] maxMod active" << endl;
		else
			cout << "yuv[" << y << "][" << u << "][" << v << "] minMod active" << endl;
		cout << "minMod[" << minMod[0] << "][" << minMod[1] << "][" << minMod[2] << "]" << endl;
		cout << "maxMod[" << maxMod[0] << "][" << maxMod[1] << "][" << maxMod[2] << "]" << endl;
		
		if(key == 'y') { y = 1; u = 0; v = 0; index = 0;}
		if(key == 'u') { y = 0; u = 1; v = 0; index = 1;}
		if(key == 'v') { y = 0; u = 0; v = 1; index = 2;}
		if(key == 'l') { thresh = 0;} //minMod
		if(key == 'h') { thresh = 1;} //maxMod
		if(key == '.') { //Up
			if(thresh == 0) { minMod[index] += 4;}
			if(thresh == 1) { maxMod[index] += 4;}
		}
		if(key == ',') { //Down
			if(thresh == 0) { minMod[index] -= 4;}
			if(thresh == 1) { maxMod[index] -= 4;}
		}
		cout << "y,u,v for channel; l for minMod, h for maxMod threshold; , for down, . for up; esq or q to quit;" << endl;
		bgd.backgroundDiffBackground(Irgb);
		cv::imshow(charstr, bgd.mask);
	}
}


////////////////////////////////////////////////////////////////
int main( int argc, char** argv) {
	cv::namedWindow( argv[0], cv::WINDOW_AUTOSIZE );
	cv::VideoCapture cap;
	if((argc < 3)|| !cap.open(argv[2])) {
		cerr << "Couldn't run the program" << endl;
		help(argv);
		cap.open(0);
		return -1;
	}
	int number_to_train_on = atoi( argv[1] );
	cv::Mat image; 
	CbBackgroudDiff bgd;

	// FIRST PROCESSING LOOP (TRAINING):
	//
	int frame_count = 0;
	int key;
	bool first_frame = true;
	cout << "Total frames to train on = " << number_to_train_on << endl; //db
	char seg[] = "Segmentation";
	while(1) {
		cout << "frame#: " << frame_count;
		cap >> image;
		if( !image.data ) exit(1); // Something went wrong, abort
		if(frame_count == 0) { bgd.init(image);}
		
		cout << ", Codebooks: " << bgd.updateCodebookBackground(image) << endl;
		
		cv::imshow( argv[0], image );
		frame_count++;
		if( (key = cv::waitKey(7)) == 27 || key == 'q' || key == 'Q' || frame_count >= number_to_train_on) break; //Allow early exit on space, esc, q
	}

	// We have accumulated our training, now create the models
	//
	cout << "Created the background model" << endl;
	cout << "Total entries cleared = " << bgd.clearStaleEntriesBackground() << endl;
	cout << "Press a key to start background differencing, 'a' to set thresholds, esc or q or Q to quit" << endl;
	
	// SECOND PROCESSING LOOP (TESTING):
	//
	cv::namedWindow( seg, cv::WINDOW_AUTOSIZE );
	while((key = cv::waitKey()) != 27 || key == 'q' || key == 'Q'  ) { // esc, 'q' or 'Q' to exit
		cap >> image;
		if( !image.data ) exit(0);
		cout <<  frame_count++ << " 'a' to adjust threholds" << endl;
		if(key == 'a') {
			cout << "Adjusting thresholds" << endl;
		cout << "y,u,v for channel; l for minMod, h for maxMod threshold; , for down, . for up; esq or q to quit;" << endl;
			adjustThresholds(seg,image,bgd);
		}
		else {
			if(bgd.backgroundDiffBackground(image)) {
				cerr << "ERROR, bdg.backgroundDiffBackground(...) failed" << endl;
				exit(-1);
			}
		}
		cv::imshow("Segmentation",bgd.mask);
	}
	exit(0);
}




