/*
  This program is used to process GeoTiff format DEM, using RegionGrowing method
to extract the water region in DEM, and set them value to 0.
*/

#include <iostream>  
#include <stack>
#include <vector>
#include "gdal_priv.h"
using namespace std;  

struct Point{
	int x;
	int y;
	Point(int x0, int y0):x(x0), y(y0){}
};

// 8 neighbors
const Point connects[8] = { Point(-1, -1), Point(0, -1), Point(1, -1), Point(1, 0), Point(1, 1), Point(0, 1), Point(-1, 1), Point(-1, 0)};  
  
int regiongrowing(const char* img, const char* img_out, vector<Point> &vseed, double reg_maxdist){ 
	//Check Parameter
	if(reg_maxdist < 0){
		cout << "Please input the correct value of reg_maxdist!\n";
		exit(1);
	}

	//Open Image
	GDALDataset *poDataset;
	GDALAllRegister();
	poDataset = (GDALDataset*)GDALOpen(img, GA_ReadOnly);
	if(poDataset == NULL){
		cout << "Please input the right image path!\n";
		exit(1);
	}

	//Get Image Size
	int size_x = poDataset->GetRasterXSize();
	int size_y = poDataset->GetRasterYSize();

	//Read Raster Data
	GDALRasterBand *poBand = poDataset->GetRasterBand(1);
	unsigned short *I = new unsigned short[size_x*size_y];
	poBand->RasterIO(GF_Read, 0, 0, size_x, size_y, I, size_x, size_y, GDT_UInt16, 0, 0);

    //Output Image And Flag Matrix
	unsigned char *res  = new unsigned char[size_x*size_y];
	unsigned char *flag = new unsigned char[size_x*size_y];
    memset(res , 0, size_x*size_y*sizeof(unsigned char));
	memset(flag, 0, size_x*size_y*sizeof(unsigned char));
    
    // Init Seed Point and Region
    stack<Point> seeds; 
	double reg_mean = 0;
	int reg_size = vseed.size();
	for(int i = 0; i < vseed.size(); ++i){
		if( vseed[i].x > size_x || vseed[i].y > size_y || vseed[i].x < 1 || vseed[i].y < 1){
			cout << "Please input the correct seed point!\n";
			exit(1);
		}
		seeds.push(vseed[i]);  
		res[(vseed[i].y-1)*size_x + vseed[i].x-1] = 255;
		reg_mean += static_cast<double>(I[(vseed[i].y-1)*size_x + vseed[i].x-1])/vseed.size();
	}
      
	//Start Growing
    while (!seeds.empty())  {  
        Point seed = seeds.top();  
        seeds.pop();  
  
        //flag the point 
		flag[seed.y*size_x+seed.x] = 1; 
  
        // through 8 neighbors 
        for (size_t i = 0; i < 8; i++)  {  
            int tmpx = seed.x + connects[i].x;  
            int tmpy = seed.y + connects[i].y;  
  
            if (tmpx < 0 || tmpy < 0 || tmpx >= size_x || tmpy >= size_y)  
                continue;  
              
			double reg_dist = fabs(static_cast<double>(I[tmpy*size_x+tmpx])-reg_mean);
            if (reg_dist < reg_maxdist && flag[tmpy*size_x+tmpx] == 0){ 
				//Re-calculate the mean value
				reg_mean = ( reg_mean*reg_size + static_cast<double>(I[tmpy*size_x+tmpx]) )/(reg_size+1);
				++reg_size;

				res[tmpy*size_x+tmpx] = 255;
				I[tmpy*size_x+tmpx] = 0;
                flag[tmpy*size_x+tmpx] = 1;   

                seeds.push(Point(tmpx, tmpy));  
            }  
        }  
    }  

	//Write result to output image
	const char *pszFormat = "GTiff";
	GDALDriver *poDriver = GetGDALDriverManager()->GetDriverByName(pszFormat);
	if(!poDriver) exit(1);

	GDALDataset *poDstDS = poDriver->CreateCopy( img_out, poDataset, TRUE, NULL, NULL, NULL );
	poBand = poDstDS->GetRasterBand(1);
	poBand->RasterIO(GF_Write, 0, 0, size_x, size_y, I, size_x, size_y, GDT_UInt16, 0, 0);

	GDALDataset *poResDS = poDriver->CreateCopy( "Mask.tif", poDataset, TRUE, NULL, NULL, NULL );
	poBand = poResDS->GetRasterBand(1);
	poBand->RasterIO(GF_Write, 0, 0, size_x, size_y, res, size_x, size_y, GDT_Byte, 0, 0);

	//Close Image and release memory
    GDALClose((GDALDatasetH)poDataset);
	GDALClose((GDALDatasetH)poDstDS);
	GDALClose((GDALDatasetH)poResDS);
	delete[] I; I = NULL;
	delete[] res; res = NULL;
	delete[] flag; flag = NULL;

	return 1;  
}  

int main(){
	const char* img = "first.tif";
	const char* img_out = "second.tif";
	vector<Point> vseed;
	/*vseed.push_back(Point(172, 2513));
	vseed.push_back(Point(572, 226));
	vseed.push_back(Point(2931, 2056));
	vseed.push_back(Point(2674, 1504));
	vseed.push_back(Point(2932, 2793));
	vseed.push_back(Point(3474, 3151));
	vseed.push_back(Point(3180, 3346));
	vseed.push_back(Point(2504, 2592));*/

	vseed.push_back(Point(2484, 2469));
	vseed.push_back(Point(2079, 1855));
	vseed.push_back(Point(2049, 1746));

	double reg_maxdist = 0.5;

	regiongrowing(img, img_out, vseed, reg_maxdist);
}