//////////////////////////////////////////////////////////////////
//         Dense stereo using Semi-Global Matching
//
//  Implementation of Heiko Hirschmuller's algorithm described in
//  "Accurate and Efficient Stereo Processing by Semi-Global Matching
//  and Mutual Information"
//
//  Author: Ilya Rosenberg, NYU Media Research Lab 12/15/2005
//
//
#ifndef __SGM_ROSENBERG_H
#define __SGM_ROSENBERG_H

#pragma once

#include "SmartArray.h"
#include "SGM_config.h"

#define DEPTH (MAX_DISP - MIN_DISP + 1)

////////////////////////////////////////////////////////////////////
// The following define errt and ERRT_MAX
// Errt is the type used for the error values in the big error array
// Using int rather than float increases speed without noticeably
// worsening the quality. Using short will reduce the array size by
// two but may hurt the quality of the result.

//#define USE_SHORTS
//#define USE_INTS

#define FLT_MAX 100000000

#define USE_FLOATS

#if defined(USE_SHORTS)
#define errt short
#define ERRT_MAX SHRT_MAX
#elif defined(USE_INTS)
#define errt int
#define ERRT_MAX INT_MAX
#elif defined(USE_FLOATS)
#define errt float
#define ERRT_MAX FLT_MAX
#endif

/////////////////////////////////////////////////////////////////////
// These constants affect how the code is compiled
// Can be left as they are unless tuning for performance

const bool subtract_common = true; // Subtracts minimum error from all previous errors (to prevent overflows for shorts)
								   // This limits the accumulated error at any pixel,depth to 16(MaxError + penalty2 + penalty3)
								   // If you enable this, you might as well enable allow jumps cuz it does half the work.									
const bool allow_jumps = true; // Increases quality (by allowing over one pixel jumps) but reduces speed
const bool use_bitmin = false; // Compares errors using bitwise ops rather than actual comparisons. May be faster



enum Side
{
	Left, Center, Right
};

enum CostStyle
{
	Regular, Clamp, Tomasi
};

enum Dir
{
	N, NNE, NE, NEE, E, SEE, SE, SSE, S, SSW, SW, SWW, W, NWW, NW, NNW
};

class Todo
{
public:
	int x, y, stepx[2], stepy[2], dir[2];

	Todo(int x, int y, Dir dir) : x(x), y(y)
	{
		switch(dir)
		{
		case N:	  setSteps(N, N); break;
		case NNE: setSteps(N, NE); break;
		case NE:  setSteps(NE, NE); break;
		case NEE: setSteps(E, NE); break;
		case E:   setSteps(E, E); break;
		case SEE: setSteps(E, SE); break;
		case SE:  setSteps(SE, SE); break;
		case SSE: setSteps(S, SE); break;
		case S:   setSteps(S, S); break;
		case SSW: setSteps(S, SW); break;
		case SW:  setSteps(SW, SW); break;
		case SWW: setSteps(W, SW); break;
		case W:   setSteps(W, W); break;
		case NWW: setSteps(W, NW); break;
		case NW:  setSteps(NW, NW); break;
		case NNW: setSteps(N, NW); break;
		default: printf("Invalid dir\n");
		}
	}

	void setSteps(Dir d1, Dir d2)
	{
		setStep(d1, stepx[0], stepy[0], dir[0]);
		setStep(d2, stepx[1], stepy[1], dir[1]);
	}

	void setStep(Dir d1, int& stepx, int& stepy, int& dir)
	{
		dir = d1;

		switch(d1)
		{
		case N:  stepx = 0; stepy = -1; break;
		case NE: stepx = 1; stepy = -1; break;
		case E:  stepx = 1; stepy = 0; break;
		case SE: stepx = 1; stepy = 1; break;
		case S:  stepx = 0; stepy = 1; break;
		case SW: stepx = -1; stepy = 1; break;
		case W:  stepx = -1; stepy = 0; break;
		case NW: stepx = -1; stepy = -1; break;
		default: printf("Invalid dir\n");
		}
	}
};

struct ResultPixel
{
	int depth;
	errt error;
	float subpixel_correction;
	float subpixel_depth;
	bool consistent;
};

//#define SGM_USE_COLOR
#ifdef SGM_USE_COLOR
struct RGB
{
    unsigned char r, g, b;
};
#else
struct RGB
{
    unsigned short int r;
};
#endif

class SGM
{
private:
	int* m_sqrtTable_p;

	bool m_imgsLocallyAllocated_b;
	RGB (*m_leftImg_p)[WIDTH];
	RGB (*m_rightImg_p)[WIDTH];

	errt (*m_diffErrVol_p)[WIDTH][2][DEPTH];

	errt (*color_penalty)[HEIGHT][WIDTH];

	ResultPixel (*m_leftResult_p)   [WIDTH];
	ResultPixel (*m_centerResult_p) [WIDTH];
	ResultPixel (*m_rightResult_p)  [WIDTH];
	ResultPixel (*m_tempResult_p)   [WIDTH];

	SmartArray<Todo> m_todos;

	errt colorDif(RGB& a, RGB& b);

public:

	SGM();
	~SGM();

	RGB (*get_img(Side side))[WIDTH];
	void set_img(Side side, RGB (*img)[WIDTH]);
	ResultPixel (*get_result_arr(Side side))[WIDTH];
	void initDataStructs(bool allocate_imgs);
	void initTodos(bool sixteen_dirs, bool optimize);
	void colorCorrect(Side side);
	void calcDifferences(Side side, errt default_err, errt cutoff, CostStyle cost_style);
	void calcColorPenalty(Side side, errt penalty2, errt penalty3);
	void semiGlobalSmoothing(errt penalty1, errt penalty2, bool use_color_penalty);
	void extractResults(Side source_side, Side result_side);
	void consistencyCheck(int thresh = 2);
	void medianFilter(Side side, int steps = 1, int thresh = 2);
	void medianFill(Side side, int steps = 500);
};


void* sgm_setup(int height, int width, int min_disp, int max_disp, bool sixteen_dirs = true);
void sgm_stereo(void* context, char* left_image, char* right_image, ResultPixel** left_results, ResultPixel** right_results, bool do_color_correct = true, bool do_median_filter = true, bool do_median_fill = true, errt penalty1 = 30, errt penalty2 = 500, errt penalty3 = 0, errt default_err = 50, CostStyle cost_style = Regular, errt cutoff = ERRT_MAX);
void sgm_free(void* context);


#endif // __SGM_ROSENBERG_H
