/*	kemy.cu

	KentuckY Error Modeler

	This program memory maps a P6 PPM image file and then
	improves it by using texture synthesis to literally
	create a new image with more consistent statistics based
	on a probability density function for pixel value
	errors. The PDF is computed from the image being
	processed by examining neighboring pixel values; there
	is no training set.

	2022 by H. Dietz
*/

#include <ctype.h>
#include <errno.h>
#include <fcntl.h>
#include <float.h>
#include <limits.h>
#include <math.h>
#include <setjmp.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <time.h>
#include <sys/types.h>
#include <stdint.h>
#include <sys/mman.h>

// C++ and OpenCV stuff
#include <opencv2/opencv.hpp>
#include <iostream>
#include <cstdio>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
 
using namespace cv;
using namespace std;

// OpenMP
#include <omp.h>

#define	EMDIM	256		// how many gray levels for error model
#define	COLORS	3		// how many colors
#define	PDFTERMS (6*3)		// number of PDF terms in similarity
#define	MINW	(1.0/128)	// minimum distance weight

typedef	uint8_t	pixel_t;

char	*myname;	// program name for error messages 
char	*infile;	// input file name for messages
int	xdim, ydim;	// x,y dimension values
int	maxval;		// maximum value, should be <=255
int	dim;		// xdim * ydim
pixel_t	*org;		// original image

// structure for walk positions
#define	WALKRAD	24	// walk radius
#define	WALKDIM	(WALKRAD+WALKRAD+1)
typedef	struct {
	int	xoff, yoff;
	float	weight;
} walk_t;
walk_t	walk[WALKDIM*WALKDIM];	// generous space for it
int	walksp = 0;		// number actually allocated

// raw histogram
#define	HIST(C,Y,X)	hist[C][Y][X]
int	hist[COLORS][EMDIM][EMDIM];

// probability density function in 8 bits
#define	PDF8(C,Y,X)	pdf8[Y][X][C]
uint8_t	pdf8[EMDIM][EMDIM][COLORS];

// actual probability density function
#define	PDF(C,Y,X)	pdf[C][Y][X]
float	pdf[COLORS][EMDIM][EMDIM];

// patch offset table
#define	POFFS	8	// neighbors in a patch
int	poff[POFFS];

void
mkhist(void)
{
	// initialize hist; mostly 0, 1 on diagonal
	for (int c=0; c<COLORS; ++c) {
		for (int j=0; j<EMDIM; ++j) {
			for (int i=0; i<EMDIM; ++i) {
				HIST(c, j, i) = (i == j);
			}
		}
	}

	// everyone evals each pixel's neighbors,
	// but only marks itself (if appropriate)
	pixel_t *p = org + COLORS + COLORS*xdim;
	for (int row=1; row<ydim-1; ++row) {
		for (int col=1; col<xdim-1; ++col) {
			for (int c=0; c<COLORS; ++c) {
				int v = *p;

				// find most similar neighbor of same color
				int near = p[poff[0]];
				int dif = abs(near - v);

				for (int k=1; k<POFFS; ++k) {
					int n = p[poff[k]];
					int d = abs(n - v);
					if (d < dif) { dif = d; near = n; }
				}
				++p;

				// bump everything between near and v
				if (near > v) { int t = near; near = v; v = t; }
				for (int j=near; j<=v; ++j) {
					for (int i=near; i<=v; ++i) {
						++(HIST(c, j, i));
					}
				}
			}
		}
		p += (COLORS + COLORS);
	}
}

inline static void
mkpdf(int myx, int myy)
{
	// make the floating-point PDF for point myx, myy
	// also produces remdering of PDF as an image

	// for each color
	for (int c=0; c<COLORS; ++c) {
		// make myv monotonic and compute max in my row
		int myv = HIST(c, myy, myx);
		int max = myv;
		for (int i=0; i<EMDIM; ++i) {
			int v = HIST(c, myy, i);
			if (v > myv) {
				// new max?
				if (v > max) max = v;

				// new monotonic correction?
				if (((myx > myy) && (i > myx)) ||
				    ((myx < myy) && (i < myx))) {
					myv = v;
				}
			}
		}

		// normalize probability density function (PDF)
		float f = myv / ((float) max);
		PDF(c, myy, myx) = powf(f, (1.0 / PDFTERMS));
		PDF8(c, myy, myx) = 0.5 + (255.0 * f);
	}
}

void
mkerrmodel(void)
{
	// compute pixel error model from image

	// patch offsets
	poff[0] = -COLORS+-COLORS*xdim;
	poff[1] =       0+-COLORS*xdim;
	poff[2] =  COLORS+-COLORS*xdim;
	poff[3] = -COLORS+           0;
	poff[4] =  COLORS+           0;
	poff[5] = -COLORS+ COLORS*xdim;
	poff[6] =       0+ COLORS*xdim;
	poff[7] =  COLORS+ COLORS*xdim;

	// histogram map errors
	mkhist();

	// create probability density function from histogram
#pragma omp parallel for schedule(guided)
	for (int y=0; y<EMDIM; ++y) {
		for (int x=0; x<EMDIM; ++x) {
			mkpdf(x, y);
		}
	}

	// write a copy of PDF as an image
	Mat pdfMat = Mat(EMDIM, EMDIM, CV_8UC3);
	memcpy(((void *) pdfMat.data), ((void *) &PDF8(0,0,0)), sizeof(pdf8));
	imwrite("errmod.png", pdfMat);
}

void
mkwalk(void)
{
	// make walk offset table, weights decay at distance

	// a circular disc, walked in memory order...	
	for (int y=(-WALKRAD); y<=WALKRAD; ++y) {
		for (int x=(-WALKRAD); x<=WALKRAD; ++x) if (y || x) {
			float d = sqrtf((float) ((y * y) + (x * x)));
			if (d <= WALKRAD) {
			        walk[walksp].yoff = y;
				walk[walksp].xoff = x;
				d = (((WALKRAD * WALKRAD) - (d * d)) / (WALKRAD * WALKRAD));
				walk[walksp].weight = (d) * (1.0 - MINW) + MINW;
				++walksp;
			}
		}
	}

	fprintf(stderr, "%d points to sample for texture\n", walksp);
}

inline static float
similarity(pixel_t *a, pixel_t *b)
{
	// how similar are these multi-color pixel values?
	return(PDF(0, a[0], b[0]) * PDF(1, a[1], b[1]) * PDF(2, a[2], b[2]));
}

inline static float
simpatch(pixel_t *a, pixel_t *b)
{
	// compute simimarity of patches around a and b
	float sim[POFFS+1];

	for (int i=0; i<POFFS; ++i) {
		sim[i] = similarity(a + poff[i], b + poff[i]);
	}
	sim[POFFS] = sim[0];

	// find highest combined similarity of two adjacent neighbors
	float best1 = sim[0];
	float best = sim[0] * sim[1];
	for (int i=0; i<POFFS; ++i) {
		float s = sim[i] * sim[i+1];
		if (s > best) best = s;
		if (sim[i] > best1) best1 = sim[i];
	}

	// return best weighted by similarity of a and b
	// this determines number of PDF terms multiplied:
	// best1 and s are 1 similarity each, best is 2
	// and each similarity is 3
	float s = similarity(a, b);
	return(best1 * best1 * best * s * s);
}

inline static void
mkspot(pixel_t *res, int row, int col)
{
	// compute texture for one spot (pixel)
	int base = COLORS * ((row * xdim) + col);
	pixel_t *op = org + base;
	pixel_t *rp = res + base;

	// start assuming value is correct
	float qsum = MINW;
	float vsum0 = qsum * op[0];
	float vsum1 = qsum * op[1];
	float vsum2 = qsum * op[2];

	// need to check bounds of each walk element
	if ((row >= 1) &&
	    (col >= 1) &&
	    (row < ydim-1) &&
	    (col < xdim-1)) {
	        // scan in walk order
		int i = 0;

		do {
			int irow = row + walk[i].yoff;
			int icol = col + walk[i].xoff;

			// need to check bounds of each walk element
			if ((irow >= 1) &&
			    (icol >= 1) &&
			    (irow < ydim-1) &&
			    (icol < xdim-1)) {
				pixel_t *pp = org + COLORS * ((irow * xdim) + icol);
				float s = simpatch(op, pp) * walk[i].weight;
				vsum0 += (s * pp[0]);
				vsum1 += (s * pp[1]);
				vsum2 += (s * pp[2]);
				qsum += s;
			}
		} while (++i < walksp);
	}

	rp[0] = vsum0 / qsum;
	rp[1] = vsum1 / qsum;
	rp[2] = vsum2 / qsum;
}

Mat
mktexture(Mat resMat)
{
	// create walk weighting table
	mkwalk();

	// make texture at each spot
#pragma omp parallel for schedule(guided)
	for (int y=0; y<ydim; ++y) {
		for (int x=0; x<xdim; ++x) {
			mkspot(((pixel_t *) resMat.data), y, x);
		}
	}
	return(resMat);
}

int
main(int argc, char **argv)
{
	myname = argv[0];
	if (argc != 2) {
		fprintf(stderr,
			"Usage: %s input_image\n"
			"Outputs error model errmod.png, improved image better.png\n",
			myname);
		exit(1);
	}
	infile = argv[1];

	// read the color image in R, G, B byte sequence
	Mat orgMat = imread(infile);
	org = ((pixel_t *) orgMat.data);
	xdim = orgMat.cols;
	ydim = orgMat.rows;
	dim = xdim * ydim;

	fprintf(stderr, "Computing error model\n");
	mkerrmodel();

	fprintf(stderr, "Computing texture\n");
	// allocate result copy
	Mat resMat = orgMat.clone();
	resMat = mktexture(resMat);

	// write denoised P6 image
	imwrite("better.png", resMat);
	imshow("Original", orgMat);
	imshow("Better", resMat);
	waitKey(0);

	return(0);
}