/*	kvirp.cpp

	Run this by something like ./kvirp </dev/ttyACM0
	It will create two windows as well as spewing text.

	January 2020, by H. Dietz
*/


//Uncomment the following line if you are compiling this code in Visual Studio
//#include "stdafx.h"

#include <opencv2/opencv.hpp>
#include <iostream>
#include <math.h>
#include <cstdio>

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/ipc.h>
#include <sys/shm.h>

using namespace cv;
using namespace std;

const	char	*name360 = "Insta360 + AMG8833";
const	char	*nameraw = "Insta360 raw video";
const	char	*nameTemp = "AMG8833";

#define	MINCONF	(1.0/255)	// minimum non-0 confidence
#define PI	3.1415926536
#define	VIEWANG	(180)
#define	HALFANG	(VIEWANG/2)
#define	TEMPANG	(8 * 7.5)	// approx view of thermal imager
#define	WALK	(10.0)
#define	WALKANG	(WALK * 7.5)	// walk covering view of thermal imager
#define	LOOKRAD	(15)		// look radius (square, actually) for alignment
#define	SHIFTPOINTS	1024	// number of point to examine to determine motion
#define	BADSHIFTQ	(SHIFTPOINTS * 5 * 256)	// bad (high) value for shiftq
#define	TICKAGE		(254.0 / 255.0)		// confidence drop per update

int	cameraNum = 1;		// which /dev/video device?
int	noiseFilter = 0;	// filter-out painted values if there's noise
int	rawvout = 0;		// raw video out?
int	stitchCert = 0;		// how certain is stitch alignment?
int	verbose = 0;		// enable verbose output

//	Stuff for parsing thermal sensor data

int	nextc = ' ';

int
advance(void)
{
	nextc = getchar();
	while (nextc == EOF) {
		nextc = getchar();
	}
	return(nextc);
}

int
see(int c)
{
	if (c == nextc) { advance(); return(1); }
	return(0);
}

int
readint(void)
{
	int s = 0;
	int v = 0;

	while ((nextc != '-') && !((nextc >= '0') && (nextc <= '9'))) advance();
	if (see('-')) s = -1;
	while ((nextc >= '0') && (nextc <= '9')) {
		v = (nextc - '0') + (v * 10);
		advance();
	}
	return(s ? -v : v);
}

#define	SHMRECS	1024

// nodes of shared memory circular buffer
// sending data from Walabot to 360 view
typedef	struct {
	int64_t		serial;
	uchar 		amg[8][8];
} shmrec_t;

volatile shmrec_t	*shm;
int64_t			shmserial = 0;

void
shmstart(void)
{
	// Allocate System V shared memory circular buffer
	int shmid = shmget(IPC_PRIVATE,
			   (SHMRECS * sizeof(shmrec_t)),
			   (IPC_CREAT | 0600));

	// Attach to it
	shm = ((volatile shmrec_t *) shmat(shmid, 0, 0));

	// Set to self-destruct when all done
	shmctl(shmid, IPC_RMID, 0);

	// Clear all the serials
	for (int i=0; i<SHMRECS; ++i) {
		shm[i].serial = -1;
	}
}

void
HeatMap(void)
{
	while (1) {
		// read 8x8 data from AMG8833
		int s = (shmserial % SHMRECS);
		int v;

		do {
			v = readint();
		} while (v >= 0);

		// paint into image
		for (int j=0; j<8; ++j)
		for (int i=0; i<8; ++i) {
			// AMG8833 not in the obvious orientation...
			*(&(shm[s].amg[0][0]) + ((7 - j) + (8 * i))) = readint();
		}

		shm[s].serial = shmserial;
		++shmserial;
	}
}


Vec3b
frac2bgr(register float f)
{
	// Computes BGR from a fraction 0..1
	// 0 becomes blue, 1 is red
	// formula from NodeScape
	Vec3b c;
	uchar r, g, b;

	/* Only use colors to blue, red is hot */
	f = 0.8 - (f * 0.8);

	switch ((int) (f * 5)) {
	case 0:
		r = 255;
		g = 5 * 255 * f;
		b = 0;
		break;
	case 1:
		r = 5 * 255 * (0.4 - f);
		g = 255;
		b = 0;
		break;
	case 2:
		r = 0;
		g = 255;
		b = 5 * 255 * (f - 0.4);
		break;
	case 3:
		r = 0;
		g = 5 * 255 * (0.8 - f);
		b = 255;
		break;
	default:
		r = 5 * 255 * (f - 0.8);
		g = 0;
		b = 255;
	}

	c.val[0] = b;
	c.val[1] = g;
	c.val[2] = r;
	return(c);
}


Point2f
tofishdex(int x, int y, int tall)
{
	Point2f pfish;
	float theta, phi, r, r2;
	Point3f psph;
	float FOV =(float)PI/180 * 180;
	float FOV2 = (float)PI/180 * 180;
	float ftall = tall;
	// Active tall crop to get 180 degrees
	float atall = ftall * 0.92;

	// Polar angles
	theta = PI * (x / ftall - 0.5); // -pi/2 to pi/2
	phi = PI * (y / ftall - 0.5);  // -pi/2 to pi/2

	// Vector in 3D space
	psph.x = cos(phi) * sin(theta);
	psph.y = cos(phi) * cos(theta);
	psph.z = sin(phi);		// cylindrical about 92%

	// Calculate fisheye angle and radius
	theta = atan2(psph.z,psph.x);
	phi = atan2(sqrt(psph.x*psph.x + psph.z*psph.z), psph.y);

	r = atall * phi / FOV;
	r2 = atall * phi / FOV2;

	// Pixel in fisheye space
	pfish.x = 0.5 * atall + r * cos(theta) + ((ftall - atall) / 2);
	pfish.y = 0.5 * atall + r2 * sin(theta) + ((ftall - atall) / 2);
	return(pfish);
}


static struct {
	int x, y;
	float frac;
} mapcache[1504][3008];
static int hascache = 0;

Mat
fastmap(Mat frame)
{
	// cache-based fisheye remapping
	Mat map(frame.rows, frame.cols, CV_8UC3, Scalar(0));
	int off = (frame.rows / 2);
	int off2 = frame.rows + off;
	int mod = frame.rows * 2;

	if (!hascache) {
		for (int j=0; j<frame.rows; ++j) {
			for (int i=0; i<frame.rows; ++i) {
				Point2f fpos = tofishdex(i, j, frame.rows);
				float fraci = (fpos.x - frame.rows/2.0) / (frame.rows/2.0) / 0.92;
				float fracj = (fpos.y - frame.rows/2.0) / (frame.rows/2.0) / 0.92;
				fraci *= fraci;
				fracj *= fracj;
				float frac = sqrt(fraci + fracj);
				frac *= frac;
				frac *= frac;
				frac = 1.0 - frac;
				frac = (stitchCert / 100.0) + (frac * (1.0 - (stitchCert / 100.0)));
				frac = ((frac > 1.0) ? 1.0 : ((frac < 0) ? 0.0 : frac));

				mapcache[j][i + off].x = fpos.x;
				mapcache[j][i + off].y = fpos.y;
				mapcache[j][i + off].frac = frac;

				mapcache[j][(i + off2) % mod].x = fpos.x + frame.rows;
				mapcache[j][(i + off2) % mod].y = fpos.y;
				mapcache[j][(i + off2) % mod].frac = frac;
			}
		}
		hascache = 1;
	}

	// Perform the actual mapping
	for (int j=0; j<frame.rows; ++j) {
		for (int i=0; i<mod; ++i) {
			Vec3b color = frame.at<Vec3b>( Point(mapcache[j][i].x, mapcache[j][i].y) );
			float frac = mapcache[j][i].frac;
			float ifrac = 127.5 * (1.0 - frac);
			color.val[0] = (color.val[0] * frac) + ifrac;
			color.val[1] = (color.val[1] * frac) + ifrac;
			color.val[2] = (color.val[2] * frac) + ifrac;
			map.at<Vec3b>(Point(i, j)) = color;
		}
	}

	return(map);
}

float
shift(int &rdx, int &rdy, Mat one, Mat two)
{
	// determine pixel shift by sampling SHIFTPOINTS random positions
	// returns a quality metric, smaller value is better
	int mody = (one.rows / 2) - (4 * LOOKRAD);
	int modx = (one.rows / 2) - (4 * LOOKRAD); // could be cols, but limit to center
	int cy = one.rows / 2;
	int sub = LOOKRAD;
	float dif[2*LOOKRAD+1][2*LOOKRAD+1];

	for (int dy=0; dy<((2*LOOKRAD)+1); ++dy)
	for (int dx=0; dx<((2*LOOKRAD)+1); ++dx) {
		dif[dy][dx] = 0;
	}

	for (int k=0; k<SHIFTPOINTS; ++k) {
		int j = LOOKRAD + (rand() % mody);
		int i = LOOKRAD + (rand() % modx);
		Vec3b rgb1 = one.at<Vec3b>(LOOKRAD + j, LOOKRAD + i);
		float r1 = rgb1[0]; 
		float g1 = rgb1[1]; 
		float b1 = rgb1[2]; 

		for (int dy=0; dy<((2*LOOKRAD)+1); ++dy)
		for (int dx=0; dx<((2*LOOKRAD)+1); ++dx) {
			Vec3b rgb2 = two.at<Vec3b>(j + dy, i + dx);
			float r2 = r1 - rgb2[0];
			float g2 = g1 - rgb2[1];
			float b2 = b1 - rgb2[2]; 
			dif[dy][dx] += (r2 * r2) + (g2 * g2) + (b2 * b2);
		}
	}

	float lo = dif[LOOKRAD][LOOKRAD];
	int bestdy = LOOKRAD;
	int bestdx = LOOKRAD;
	for (int dy=0; dy<((2*LOOKRAD)+1); ++dy)
	for (int dx=0; dx<((2*LOOKRAD)+1); ++dx) {
		if (dif[dy][dx] < lo) {
			lo = dif[dy][dx];
			bestdy = dy;
			bestdx = dx;
		}
	}

	rdy = bestdy - LOOKRAD;
	rdx = bestdx - LOOKRAD;
}

int
main(int argc, char** argv)
{
	Mat frame;
	Mat map;
	Mat oldmap;
	Mat fmap;
	int firsttime = 1;

	// command line args
	for (int i=1; i<argc; ++i) {
		char *p = argv[i];
 
		if (*p != '-') {
usage:
			fprintf(stderr,
"%s: Usage: %s -n -v\n"
"-c#\tCamera device, /dev/video# (default %d)\n"
"-n\tNoise filter image motion\n"
"-r\tEnable raw video out\n"
"-s#\tStitch certainty #%% (default %d)\n"
"-v\tVerbose output\n",
				argv[0],
				argv[0],
				cameraNum,
				stitchCert);
			exit(1);
		}

		while (*++p) {
			switch (*p) {
			case 'c':	cameraNum = atoi(p+1);
					while (p[1]) ++p;
					break;
			case 'n':	noiseFilter = 1;
					break;
			case 'r':	rawvout = 1;
					break;
			case 's':	stitchCert = atoi(p+1);
					stitchCert = ((stitchCert < 1) ? 0 : ((stitchCert > 100) ? 100 : stitchCert));
					while (p[1]) ++p;
					break;
			case 'v':	verbose = 1;
					break;
			default:
					goto usage;
			}
		}
	}


	// Create shared memory segment
	shmstart();

	// Run thermal sensor in a separate process
	if (!fork()) {
		HeatMap();
		exit(1);
	}

	//Open the default video camera
	VideoCapture cap(cameraNum);

	// if not success, exit program
	if (cap.isOpened() == false) {
		cout << "Cannot open the video camera" << endl;
		cin.get(); //wait for any key press
		return(-1);
	} 

	namedWindow(name360, WINDOW_NORMAL); //create a window for 360 painted view

	// read a video frame
	bool bSuccess = cap.read(frame); // read a new frame from video 

	// exit if the frame cannot be captured
	if (bSuccess == false) {
		cout << "Video camera is disconnected" << endl;
		cin.get(); //Wait for any key press
		exit(1);
	}

	if (rawvout) {
		namedWindow(nameraw, WINDOW_NORMAL); //create a window for raw view
		imshow(nameraw, frame);
	}

	map = fastmap(frame);
	imshow(name360, map);

	Mat conf(frame.rows, frame.cols, CV_32F, Scalar(0)); //confidence
	Mat temp(frame.rows, frame.cols, CV_32F, Scalar(1)); //temperature
	Mat ttemp(frame.rows, frame.cols, CV_32F, Scalar(2)); //temperature
	Mat tconf(frame.rows, frame.cols, CV_32F, Scalar(3)); //confidence

	Mat image(8, 8, CV_8UC3, Scalar(0));

	namedWindow(nameTemp, WINDOW_NORMAL); // Create a window for temp sensor only


	while (1) {
		float mintemp = 255;
		float maxtemp = 0;
		int s = (shmserial % SHMRECS);
		int newamg = 0;

		// get latest data
		if (shm[s].serial >= shmserial) {
			newamg = 1;

			// skip any old records
			// video is faster, so should only skip at startup
			do {
				s = (shmserial % SHMRECS);
				shmserial = shm[s].serial + 1;
			} while (shm[shmserial % SHMRECS].serial >= shmserial);

			// display temp image
			// scale temps
			for (int j=0; j<8; ++j)
			for (int i=0; i<8; ++i) {
				uchar v = shm[s].amg[j][i];
				if (v < mintemp) mintemp = v;
				if (v > maxtemp) maxtemp = v;
			}

			// make the image
			uchar *p = image.data;
			for (int j=0; j<8; ++j)
			for (int i=0; i<8; ++i) {
				uchar v = shm[s].amg[j][i];
				Vec3b tbgr = frac2bgr((v - mintemp) / (maxtemp - mintemp));
				p[0] = tbgr.val[0];
				p[1] = tbgr.val[1];
				p[2] = tbgr.val[2];
				p += 3;
			}

			imshow(nameTemp, image);                // Show our image inside it.
		}


		// read a video frame
		bool bSuccess = cap.read(frame); // read a new frame from video 

		//Breaking the while loop if the frames cannot be captured
		if (bSuccess == false) {
			cout << "Video camera is disconnected" << endl;
			cin.get(); //Wait for any key press
			break;
		}

		if (rawvout) {
			imshow(nameraw, frame);
		}

		// make color version of remapped
		map = fastmap(frame);
		if (firsttime) { map.copyTo(oldmap); firsttime=0; }

		mintemp = 255;
		maxtemp = 0;

		// determine and apply shift
		int bestdx, bestdy;
		float shiftq = shift(bestdx, bestdy, oldmap, map);
//		cout << "shift is " << bestdx << "," << bestdy << " quality " << shiftq << endl;

		// shift the previous temp & adjust confidence
		// confidence reduces with age & low confidence shifts
		float confmul = ((shiftq > BADSHIFTQ) ? 1.0 : (shiftq / BADSHIFTQ));
		confmul *= confmul;
		confmul *= confmul;
		confmul = TICKAGE * (1.0 - confmul);
		temp.copyTo(ttemp);
		conf.copyTo(tconf);
		for (int j=0; j<temp.rows; ++j) {
			int tj = (j + bestdy + temp.rows) % temp.rows;
			for (int i=0; i<temp.cols; ++i) {
				int ti = (i + bestdx + temp.cols) % temp.cols;
				float t = ttemp.at<float>(j, i);
				float c = tconf.at<float>(j, i);

				// also age by local motion
				if (noiseFilter && (c > MINCONF)) {
					Vec3b rgb0 = oldmap.at<Vec3b>(tj, ti);
					Vec3b rgb1 = map.at<Vec3b>(tj, ti);
					float b = (rgb0.val[0] - rgb1.val[0]) / 255.0;
					float g = (rgb0.val[1] - rgb1.val[1]) / 255.0;
					float r = (rgb0.val[2] - rgb1.val[2]) / 255.0;
					float d = (b * b) + (g * g) + (r * r);
					d = 1.0 - d;
					if (d > 0) c *= d;
				}

				if (c > MINCONF) {
					if (t < mintemp) mintemp = t;
					if (t > maxtemp) maxtemp = t;
				} else {
					t = 0;
					c = 0;
				}
				temp.at<float>(tj, ti) = t;
				conf.at<float>(tj, ti) = confmul * c;
			}
		}

		// done using old 360 map
		map.copyTo(oldmap);

		// add-in the latest temperature readings
		if (newamg) {
			int tspan = temp.rows * (WALKANG / VIEWANG);
			int minj = (temp.rows / 2) - (tspan / 2);
			int mini = (temp.cols / 2) - (tspan / 2);
			for (int j=0; j<=tspan; ++j) {
				float yco = ((j * 10.0) / tspan);
				int ybase = yco;
				float yw = yco - ybase;
				--ybase;
				for (int i=0; i<=tspan; ++i) {
					float xco = ((i * 10.0) / tspan);
					int xbase = xco;
					float xw = xco - xbase;
					float v = 0;
					float m = 0;
					--xbase;

#define	VALIDTX(V)	(!((V)&~7))	// is a valid index for 8x8 AMG8833 thermal data?

					if (VALIDTX(ybase) && VALIDTX(xbase)) {
						float w = (1.0 - yw) * (1.0 - xw);
						v += (shm[s].amg[ybase][xbase] * w);
						m += w;
					}
					if (VALIDTX(ybase) && VALIDTX(xbase+1)) {
						float w = (1.0 - yw) * (xw);
						v += (shm[s].amg[ybase][xbase+1] * w);
						m += w;
					}
					if (VALIDTX(ybase+1) && VALIDTX(xbase)) {
						float w = (yw) * (1.0 - xw);
						v += (shm[s].amg[ybase+1][xbase] * w);
						m += w;
					}
					if (VALIDTX(ybase+1) && VALIDTX(xbase+1)) {
						float w = (yw) * (xw);
						v += (shm[s].amg[ybase+1][xbase+1] * w);
						m += w;
					}

					// v is the temperature value
					v = ((m > 0) ? (v / m) : 0);
					// m is the confidence, with max value 2

					// merge with previous values
					float ov = temp.at<float>(minj+j, mini+i);
					float om = conf.at<float>(minj+j, mini+i);
					if (om > 1.0) om = 1.0;
					ov = ((ov * om) + (v * m)) / (om + m);
					om += m;
					if (om > 1.0) om = 1.0;
					if (ov < mintemp) mintemp = ov;
					if (ov > maxtemp) maxtemp = ov;
					temp.at<float>(minj+j, mini+i) = ov;
					conf.at<float>(minj+j, mini+i) = om;
				}
			}
		}

		// now paint the temperatures
		map.copyTo(fmap);
		if (mintemp < maxtemp) {
			float tempscale = 1.0 / (maxtemp - mintemp);
			for (int j=0; j<fmap.rows; ++j)
			for (int i=0; i<fmap.cols; ++i) {
				float c = conf.at<float>(j, i);

				// only bother painting if some confidence
				if (c > (1.0 / 255.0)) {
					float t = tempscale * (temp.at<float>(j, i) - mintemp);
					Vec3b tbgr = frac2bgr(t);
					Vec3b bgr = fmap.at<Vec3b>(j, i);
//					if (c > 1) c = 1;	// should be true already
					c = (c * c * c) / 2.0;	// cube fraction to make shading more subtle
					bgr.val[0] = (bgr.val[0] * (1.0 - c)) + (tbgr.val[0] * c);
					bgr.val[1] = (bgr.val[1] * (1.0 - c)) + (tbgr.val[1] * c);
					bgr.val[2] = (bgr.val[2] * (1.0 - c)) + (tbgr.val[2] * c);
					fmap.at<Vec3b>(j, i) = bgr;
				}
			}
		}

		imshow(name360, fmap);

		// cause display to update
		waitKey(1);
	}
}