#include <ImageStack.h>

using namespace ImageStack;

int main(int argc, char **argv) {
    Image dx = Load::apply(argv[1]);
    Image dy = dx.copy();

    Gradient::apply(dx, 'x');
    Gradient::apply(dy, 'y');

    Image gx = dx.copy();
    Image gy = dy.copy();

    float *dxPtr;
    float *gxPtr;
    for (int t = 0; t < dx.frames; t++) {
	for (int y = 0; y < dx.height; y++) {
	    dxPtr = dx(t, 1, y);
	    gxPtr = gx(t, 1, y);
	    for (int x = 1; x < dx.width-1; x++) {
		// look for the largest neighbour
		float mag[3] = {0, 0, 0};
		for (int c = 0; c < dx.channels; c++) {
		    mag[0] += dxPtr[-dx.channels+c]*dxPtr[-dx.channels+c];
		    mag[1] += dxPtr[c]*dxPtr[c];
		    mag[2] += dxPtr[dx.channels+c]*dxPtr[dx.channels+c];
		}
		
		if (mag[0] >= mag[1] && mag[0] >= mag[2]) {
		    // it's the guy on the left
		    for (int c = 0; c < dx.channels; c++) {
			gxPtr[-dx.channels+c] += dxPtr[c];
			gxPtr[c] -= dxPtr[c];
		    }
		} else if (mag[2] > mag[1]) {
		    // it's the guy on the right
		    for (int c = 0; c < dx.channels; c++) {
			gxPtr[dx.channels+c] += dxPtr[c];
			gxPtr[c] -= dxPtr[c];
		    }
		}

		dxPtr += dx.channels;
		gxPtr += gx.channels;
	    }
	}
    }

    float *dyPtr;
    float *gyPtr;
    for (int t = 0; t < dy.frames; t++) {
	for (int y = 1; y < dy.height-1; y++) {
	    dyPtr = dy(t, 0, y);
	    gyPtr = gy(t, 0, y);
	    for (int x = 0; x < dy.width; x++) {
		// look for the largest neighbour
		float mag[3] = {0, 0, 0};
		for (int c = 0; c < dy.channels; c++) {
		    mag[0] += dyPtr[-dy.ystride+c]*dyPtr[-dy.ystride+c];
		    mag[1] += dyPtr[c]*dyPtr[c];
		    mag[2] += dyPtr[dy.ystride+c]*dyPtr[dy.ystride+c];
		}
		
		if (mag[0] >= mag[1] && mag[0] >= mag[2]) {
		    // it's the guy above
		    for (int c = 0; c < dy.channels; c++) {
			gyPtr[-dy.ystride+c] += dyPtr[c];
			gyPtr[c] -= dyPtr[c];
		    }
		} else if (mag[2] > mag[1]) {
		    // it's the guy below
		    for (int c = 0; c < dy.channels; c++) {
			gyPtr[dy.ystride+c] += dyPtr[c];
			gyPtr[c] -= dyPtr[c];
		    }
		}

		dyPtr += dy.channels;
		gyPtr += gy.channels;
	    }
	}
    }

    Save::apply(gx, "gx.tmp");
    Save::apply(gy, "gy.tmp");
    
    Image out = Poisson::apply(gx, gy, 10);
    Save::apply(out, "out.tmp");

    return 0;
}