#ifndef USING_R
#include "nifti1.h"
#endif
#include "nifti1_io_core.h"
#include "nii_ortho.h"
#include <ctype.h>
#include <float.h>
#include <math.h>
#include <stdbool.h> //requires VS 2015 or later
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
// #include <unistd.h>
#include <stdio.h>
#ifndef _MSC_VER

#include <unistd.h>

#endif
// #define MY_DEBUG //verbose text reporting

#include "print.h"

typedef struct {
	int v[3];
} vec3i;

mat33 matDotMul33(mat33 a, mat33 b)
// in Matlab: ret = a'.*b
{
	mat33 ret;
	for (int i = 0; i < 3; i++) {
		for (int j = 0; j < 3; j++) {
			ret.m[i][j] = a.m[i][j] * b.m[j][i];
		}
	}
	return ret;
}

mat33 matMul33(mat33 a, mat33 b)
// mult = a * b
{
	mat33 mult;
	for (int i = 0; i < 3; i++) {
		for (int j = 0; j < 3; j++) {
			mult.m[j][i] = 0;
			for (int k = 0; k < 3; k++)
				mult.m[j][i] += a.m[j][k] * b.m[k][i];
		}
	}
	return mult;
}

float getOrthoResidual(mat33 orig, mat33 transform) {
	mat33 mat = matDotMul33(orig, transform);
	float ret = 0;
	for (int i = 0; i < 3; i++) {
		for (int j = 0; j < 3; j++) {
			ret = ret + (mat.m[i][j]);
		}
	}
	return ret;
}

mat33 getBestOrient(mat44 R, vec3i flipVec)
// flipVec reports flip: [1 1 1]=no flips, [-1 1 1] flip X dimension
{
	mat33 ret, newmat, orig;
	LOAD_MAT33(orig, R.m[0][0], R.m[0][1], R.m[0][2],
			   R.m[1][0], R.m[1][1], R.m[1][2],
			   R.m[2][0], R.m[2][1], R.m[2][2]);
	float best = 0; // FLT_MAX;
	float newval;
	for (int rot = 0; rot < 6; rot++) { // 6 rotations
		switch (rot) {
		case 0:
			LOAD_MAT33(newmat, flipVec.v[0], 0, 0, 0, flipVec.v[1], 0, 0, 0, flipVec.v[2]);
			break;
		case 1:
			LOAD_MAT33(newmat, flipVec.v[0], 0, 0, 0, 0, flipVec.v[1], 0, flipVec.v[2], 0);
			break;
		case 2:
			LOAD_MAT33(newmat, 0, flipVec.v[0], 0, flipVec.v[1], 0, 0, 0, 0, flipVec.v[2]);
			break;
		case 3:
			LOAD_MAT33(newmat, 0, flipVec.v[0], 0, 0, 0, flipVec.v[1], flipVec.v[2], 0, 0);
			break;
		case 4:
			LOAD_MAT33(newmat, 0, 0, flipVec.v[0], flipVec.v[1], 0, 0, 0, flipVec.v[2], 0);
			break;
		case 5:
			LOAD_MAT33(newmat, 0, 0, flipVec.v[0], 0, flipVec.v[1], 0, flipVec.v[2], 0, 0);
			break;
		}
		newval = getOrthoResidual(orig, newmat);
		if (newval > best) {
			best = newval;
			ret = newmat;
		}
	}
	return ret;
}

bool isMat44Canonical(mat44 R)
// returns true if diagonals >0 and all others =0
//   no rotation is necessary - already in perfect orthogonal alignment
{
	for (int i = 0; i < 3; i++) {
		for (int j = 0; j < 3; j++) {
			if ((i == j) && (R.m[i][j] <= 0))
				return false;
			if ((i != j) && (R.m[i][j] != 0))
				return false;
		} // j
	} // i
	return true;
}

vec3i setOrientVec(mat33 m)
// Assumes isOrthoMat NOT computed on INVERSE, hence return INVERSE of solution...
// e.g. [-1,2,3] means reflect x axis, [2,1,3] means swap x and y dimensions
{
	vec3i ret = {{0, 0, 0}};
	// mat33 m = {-1,0,0, 0,1,0, 0,0,1};
	for (int i = 0; i < 3; i++) {
		for (int j = 0; j < 3; j++) {
			if (m.m[i][j] > 0)
				ret.v[j] = i + 1;
			if (m.m[i][j] < 0)
				ret.v[j] = -(i + 1);
		} // j
	} // i
	return ret;
}

mat44 setMat44Vec(mat33 m33, vec3 Translations)
// convert a 3x3 rotation matrix to a 4x4 matrix where the last column stores translations and the last row is 0 0 0 1
{
	mat44 m44;
	for (int i = 0; i < 3; i++) {
		for (int j = 0; j < 3; j++) {
			m44.m[i][j] = m33.m[i][j];
		}
	}
	m44.m[0][3] = Translations.v[0];
	m44.m[1][3] = Translations.v[1];
	m44.m[2][3] = Translations.v[2];
	m44.m[3][0] = 0;
	m44.m[3][1] = 0;
	m44.m[3][2] = 0;
	m44.m[3][3] = 1;
	return m44;
}

mat44 sFormMat(struct nifti_1_header *h) {
	mat44 s;
	s.m[0][0] = h->srow_x[0];
	s.m[0][1] = h->srow_x[1];
	s.m[0][2] = h->srow_x[2];
	s.m[0][3] = h->srow_x[3];
	s.m[1][0] = h->srow_y[0];
	s.m[1][1] = h->srow_y[1];
	s.m[1][2] = h->srow_y[2];
	s.m[1][3] = h->srow_y[3];
	s.m[2][0] = h->srow_z[0];
	s.m[2][1] = h->srow_z[1];
	s.m[2][2] = h->srow_z[2];
	s.m[2][3] = h->srow_z[3];
	s.m[3][0] = 0;
	s.m[3][1] = 0;
	s.m[3][2] = 0;
	s.m[3][3] = 1;
	return s;
}

void mat2sForm(struct nifti_1_header *h, mat44 s) {
	h->srow_x[0] = s.m[0][0];
	h->srow_x[1] = s.m[0][1];
	h->srow_x[2] = s.m[0][2];
	h->srow_x[3] = s.m[0][3];
	h->srow_y[0] = s.m[1][0];
	h->srow_y[1] = s.m[1][1];
	h->srow_y[2] = s.m[1][2];
	h->srow_y[3] = s.m[1][3];
	h->srow_z[0] = s.m[2][0];
	h->srow_z[1] = s.m[2][1];
	h->srow_z[2] = s.m[2][2];
	h->srow_z[3] = s.m[2][3];
}

size_t *orthoOffsetArray(int dim, int stepBytesPerVox) {
	// return lookup table of length dim with values incremented by stepBytesPerVox
	//  e.g. if Dim=10 and stepBytes=2: 0,2,4..18, is stepBytes=-2 18,16,14...0
	size_t *lut = (size_t *)malloc(dim * sizeof(size_t));
	if (stepBytesPerVox > 0)
		lut[0] = 0;
	else
		lut[0] = -stepBytesPerVox * (dim - 1);
	if (dim > 1)
		for (int i = 1; i < dim; i++)
			lut[i] = lut[i - 1] + (size_t)stepBytesPerVox;
	return lut;
} // orthoOffsetArray()

// void  reOrientImg( unsigned char * restrict img, vec3i outDim, vec3i outInc, int bytePerVox, int nvol) {
void reOrientImg(unsigned char *img, vec3i outDim, vec3i outInc, int bytePerVox, int nvol) {
	// reslice data to new orientation
	// generate look up tables
	size_t *xLUT = orthoOffsetArray(outDim.v[0], bytePerVox * outInc.v[0]);
	size_t *yLUT = orthoOffsetArray(outDim.v[1], bytePerVox * outInc.v[1]);
	size_t *zLUT = orthoOffsetArray(outDim.v[2], bytePerVox * outInc.v[2]);
	// convert data
	size_t bytePerVol = bytePerVox * outDim.v[0] * outDim.v[1] * outDim.v[2]; // number of voxels in spatial dimensions [1,2,3]
	size_t o = 0;															  // output address
	uint8_t *inbuf = (uint8_t *)malloc(bytePerVol);							  // we convert 1 volume at a time
	uint8_t *outbuf = (uint8_t *)img;										  // source image
	for (int vol = 0; vol < nvol; vol++) {
		memcpy(&inbuf[0], &outbuf[vol * bytePerVol], bytePerVol); // copy source volume
		for (int z = 0; z < outDim.v[2]; z++)
			for (int y = 0; y < outDim.v[1]; y++)
				for (int x = 0; x < outDim.v[0]; x++) {
					memcpy(&outbuf[o], &inbuf[xLUT[x] + yLUT[y] + zLUT[z]], bytePerVox);
					o = o + bytePerVox;
				} // for each x
	} // for each volume
	// free arrays
	free(inbuf);
	free(xLUT);
	free(yLUT);
	free(zLUT);
} // reOrientImg

unsigned char *reOrient(unsigned char *img, struct nifti_1_header *h, vec3i orientVec, mat33 orient, vec3 minMM)
// e.g. [-1,2,3] means reflect x axis, [2,1,3] means swap x and y dimensions
{
	size_t nvox = h->dim[1] * h->dim[2] * h->dim[3];
	if (nvox < 1)
		return img;
	vec3i outDim = {{0, 0, 0}};
	vec3i outInc = {{0, 0, 0}};
	for (int i = 0; i < 3; i++) { // set dimension, pixdim and
		outDim.v[i] = h->dim[abs(orientVec.v[i])];
		if (abs(orientVec.v[i]) == 1)
			outInc.v[i] = 1;
		if (abs(orientVec.v[i]) == 2)
			outInc.v[i] = h->dim[1];
		if (abs(orientVec.v[i]) == 3)
			outInc.v[i] = h->dim[1] * h->dim[2];
		if (orientVec.v[i] < 0)
			outInc.v[i] = -outInc.v[i]; // flip
	} // for each dimension
	int nvol = 1; // convert all non-spatial volumes from source to destination
	for (int vol = 4; vol < 8; vol++) {
		if (h->dim[vol] > 1)
			nvol = nvol * h->dim[vol];
	}
	reOrientImg(img, outDim, outInc, h->bitpix / 8, nvol);
	// now change the header....
	vec3 outPix = {{h->pixdim[abs(orientVec.v[0])], h->pixdim[abs(orientVec.v[1])], h->pixdim[abs(orientVec.v[2])]}};
	for (int i = 0; i < 3; i++) {
		h->dim[i + 1] = outDim.v[i];
		h->pixdim[i + 1] = outPix.v[i];
	}
	mat44 s = sFormMat(h);
	mat33 mat; // computer transform
	LOAD_MAT33(mat, s.m[0][0], s.m[0][1], s.m[0][2],
			   s.m[1][0], s.m[1][1], s.m[1][2],
			   s.m[2][0], s.m[2][1], s.m[2][2]);
	mat = matMul33(mat, orient);
	s = setMat44Vec(mat, minMM); // add offset
	mat2sForm(h, s);
	h->qform_code = h->sform_code; // apply to the quaternion as well
	float dumdx, dumdy, dumdz;
	nifti_mat44_to_quatern(s, &h->quatern_b, &h->quatern_c, &h->quatern_d, &h->qoffset_x, &h->qoffset_y, &h->qoffset_z, &dumdx, &dumdy, &dumdz, &h->pixdim[0]);
	return img;
} // reOrient()

float getDistance(vec3 v, vec3 min)
// scalar distance between two 3D points - Pythagorean theorem
{
	return sqrt(pow((v.v[0] - min.v[0]), 2) + pow((v.v[1] - min.v[1]), 2) + pow((v.v[2] - min.v[2]), 2));
}

vec3 xyz2mm(mat44 R, vec3 v) {
	vec3 ret;
	for (int i = 0; i < 3; i++) {
		ret.v[i] = ((R.m[i][0] * v.v[0]) + (R.m[i][1] * v.v[1]) + (R.m[i][2] * v.v[2]) + R.m[i][3]);
	}
	return ret;
}

vec3 minCornerFlip(struct nifti_1_header *h, vec3i *flipVec)
// orthogonal rotations and reflections applied as 3x3 matrices will cause the origin to shift
//  a simple solution is to first compute the most left, posterior, inferior voxel in the source image
//  this voxel will be at location i,j,k = 0,0,0, so we can simply use this as the offset for the final 4x4 matrix...
{
	int i, j, minIndex;
	vec3i flipVecs[8];
	vec3 corner[8], min;
	mat44 s = sFormMat(h);
	for (int i = 0; i < 8; i++) {
		if (i & 1)
			flipVecs[i].v[0] = -1;
		else
			flipVecs[i].v[0] = 1;
		if (i & 2)
			flipVecs[i].v[1] = -1;
		else
			flipVecs[i].v[1] = 1;
		if (i & 4)
			flipVecs[i].v[2] = -1;
		else
			flipVecs[i].v[2] = 1;
		corner[i] = setVec3(0, 0, 0); // assume no reflections
		if ((flipVecs[i].v[0]) < 1)
			corner[i].v[0] = h->dim[1] - 1; // reflect X
		if ((flipVecs[i].v[1]) < 1)
			corner[i].v[1] = h->dim[2] - 1; // reflect Y
		if ((flipVecs[i].v[2]) < 1)
			corner[i].v[2] = h->dim[3] - 1; // reflect Z
		corner[i] = xyz2mm(s, corner[i]);
	}
	// find extreme edge from ALL corners....
	min = corner[0];
	for (i = 1; i < 8; i++) {
		for (j = 0; j < 3; j++) {
			if (corner[i].v[j] < min.v[j])
				min.v[j] = corner[i].v[j];
		}
	}
	float dx; // observed distance from corner
	float min_dx = getDistance(corner[0], min);
	minIndex = 0; // index of corner closest to min
	// see if any corner is closer to absmin than the first one...
	for (i = 1; i < 8; i++) {
		dx = getDistance(corner[i], min);
		if (dx < min_dx) {
			min_dx = dx;
			minIndex = i;
		}
	}
	min = corner[minIndex]; // this is the single corner closest to min from all
	*flipVec = flipVecs[minIndex];
	return min;
}

#ifdef MY_DEBUG
void reportMat44o(char *str, mat44 A) {
	printMessage("%s = [%g %g %g %g; %g %g %g %g; %g %g %g %g; 0 0 0 1]\n", str,
				 A.m[0][0], A.m[0][1], A.m[0][2], A.m[0][3],
				 A.m[1][0], A.m[1][1], A.m[1][2], A.m[1][3],
				 A.m[2][0], A.m[2][1], A.m[2][2], A.m[2][3]);
}
#endif

unsigned char *nii_setOrtho(unsigned char *img, struct nifti_1_header *h) {
	if ((h->dim[1] < 1) || (h->dim[2] < 1) || (h->dim[3] < 1))
		return img;
	if ((h->sform_code == NIFTI_XFORM_UNKNOWN) && (h->qform_code != NIFTI_XFORM_UNKNOWN)) { // only q-form provided
		mat44 q = nifti_quatern_to_mat44(h->quatern_b, h->quatern_c, h->quatern_d,
										 h->qoffset_x, h->qoffset_y, h->qoffset_z,
										 h->pixdim[1], h->pixdim[2], h->pixdim[3], h->pixdim[0]);
		mat2sForm(h, q); // convert q-form to s-form
		h->sform_code = h->qform_code;
	}
	if (h->sform_code == NIFTI_XFORM_UNKNOWN) {
#ifdef MY_DEBUG
		printMessage("No Q or S spatial transforms - assuming canonical orientation");
#endif
		return img;
	}
	mat44 s = sFormMat(h);
	if (isMat44Canonical(s)) {
#ifdef MY_DEBUG
		printMessage("Image in perfect alignment: no need to reorient");
#endif
		return img;
	}
	vec3i flipV;
	vec3 minMM = minCornerFlip(h, &flipV);
	mat33 orient = getBestOrient(s, flipV);
	vec3i orientVec = setOrientVec(orient);
	if ((orientVec.v[0] == 1) && (orientVec.v[1] == 2) && (orientVec.v[2] == 3)) {
#ifdef MY_DEBUG
		printMessage("Image already near best orthogonal alignment: no need to reorient\n");
#endif
		return img;
	}
	bool is24 = false;
	if (h->bitpix == 24) { // RGB stored as planar data. treat as 3 8-bit slices
		return img;
		/*is24 = true;
		h->bitpix = 8;
		h->dim[3] = h->dim[3] * 3;*/
	}
	img = reOrient(img, h, orientVec, orient, minMM);
	if (is24) {
		h->bitpix = 24;
		h->dim[3] = h->dim[3] / 3;
	}
#ifdef MY_DEBUG
	printMessage("NewRotation= %d %d %d\n", orientVec.v[0], orientVec.v[1], orientVec.v[2]);
	printMessage("MinCorner= %.2f %.2f %.2f\n", minMM.v[0], minMM.v[1], minMM.v[2]);
	reportMat44o((char *)"input", s);
	s = sFormMat(h);
	reportMat44o((char *)"output", s);
#endif
	return img;
}
