
#include <math.h>
#include <i3d/FastLevelSet.h>
#include <i3d/image3d.h>
#include <i3d/descriptors/shape.h>
#include <i3d/morphology.h>

#include "../settings.h"
#include "shape_rendering.h"
#include "rnd_generators.h"
#include "deformations.h"

//----------------------------------------------------------------
bool VALUE1_OR_VALUE2(bool value1, bool value2) 
{ 
	 return (value1 || value2); 
}

bool VALUE1_AND_NEG_VALUE2(bool value1, bool value2) 
{ 
	 return (value1 && (!value2)); 
}

//----------------------------------------------------------------
void CombineImagesRespectingItsOffset(i3d::Image3d<bool> &imgBase,
												  const i3d::Image3d<bool> &imgProjected,
												  bool (*binOperator)(bool, bool))
{
	 if (imgBase.GetResolution().GetRes() != imgProjected.GetResolution().GetRes())
	 {
		  throw ERROR_REPORT("Merging the images with different resolution!?");
	 }

	 i3d::Vector3d<float> res = imgBase.GetResolution().GetRes();

	 i3d::Vector3d<float> shiftInMicrons = 
				imgProjected.GetOffset() - imgBase.GetOffset();

	 i3d::Vector3d<int> shiftInPixels;

	 shiftInPixels.x = static_cast<int>(roundf(shiftInMicrons.x * res.x));
	 shiftInPixels.y = static_cast<int>(roundf(shiftInMicrons.y * res.y));
	 shiftInPixels.z = static_cast<int>(roundf(shiftInMicrons.z * res.z));

	 for (size_t i=0; i<imgProjected.GetImageSize(); i++)
	 {
		  bool projectedValue = imgProjected.GetVoxel(i);

		  if (projectedValue == true)
		  {
				i3d::Vector3d<size_t> pos = imgProjected.GetPos(i);

				if (imgBase.Include(shiftInPixels.x + (int)pos.x,
										  shiftInPixels.y + (int)pos.y,
										  shiftInPixels.z + (int)pos.z))
				{
					 i3d::Vector3d<size_t> posII((size_t)(shiftInPixels.x + pos.x),
														  (size_t)(shiftInPixels.y + pos.y),
														  (size_t)(shiftInPixels.z + pos.z));

					 bool baseValue = imgBase.GetVoxel(posII);

					 bool newBaseValue = binOperator(baseValue, projectedValue);

					 imgBase.SetVoxel(posII, newBaseValue);
				}
		  }
	 }
}

//----------------------------------------------------------------
void AddImageRespectingItsOffset(i3d::Image3d<bool> &imgBase,
											const i3d::Image3d<bool> &imgAdded)
{
	 CombineImagesRespectingItsOffset(imgBase, imgAdded, VALUE1_OR_VALUE2);
}

//----------------------------------------------------------------
void SubtractImageRespectingItsOffset(i3d::Image3d<bool> &imgBase, 
													const i3d::Image3d<bool> &imgSubtracted)
{
	 CombineImagesRespectingItsOffset(imgBase, imgSubtracted, VALUE1_AND_NEG_VALUE2);
}


//----------------------------------------------------------------
void AlignMasks(const i3d::Image3d<bool> &firstMask,
					 i3d::Image3d<bool> &secondMask)
{
	 // find principal axes of the first and the second mask
	 i3d::Shape<bool> shape;

	 // Transform matrix A1 consists of PCA basis vectors 'pas1'
	 // stored in rows.
	 shape.SetImage(&firstMask);
	 shape.SetLabel(true);
	 i3d::PrincipalAxes pas1 = shape.GetPrincipalAxes();

	 // Transform matrix A2 consists of PCA basis vectors 'pas2'
	 // stored in rows.
	 shape.SetImage(&secondMask);
	 shape.SetLabel(true);
	 i3d::PrincipalAxes pas2 = shape.GetPrincipalAxes();

	 // Find out the centre of mask2
	 i3d::Vector3d<float> maskCentre = 
				i3d::PixelsToMicrons(shape.Centroid(),
											secondMask.GetResolution());

	 // Matrix 'B' defining the transform of pixel positions
	 // from image 'firstMask' into image 'secondMask' is
	 // given by: A2^(-1) * A1, where '*' is a matrix multiplication.
	 //
	 // A2 is an ortohogal matrix, i.e. the inversion of A2 is 
	 // equal to the transposition of A2.
	 i3d::Vector3d<float> B[3];

	 B[0].x = 
				pas2.axe1.x * pas1.axe1.x + 
				pas2.axe2.x * pas1.axe2.x +
				pas2.axe3.x * pas1.axe3.x;

	 B[0].y = 
				pas2.axe1.x * pas1.axe1.y + 
				pas2.axe2.x * pas1.axe2.y +
				pas2.axe3.x * pas1.axe3.y;

	 B[0].z =
				pas2.axe1.x * pas1.axe1.z + 
				pas2.axe2.x * pas1.axe2.z +
				pas2.axe3.x * pas1.axe3.z;

	 B[1].x = 
				pas2.axe1.y * pas1.axe1.x + 
				pas2.axe2.y * pas1.axe2.x +
				pas2.axe3.y * pas1.axe3.x;

	 B[1].y = 
				pas2.axe1.y * pas1.axe1.y + 
				pas2.axe2.y * pas1.axe2.y +
				pas2.axe3.y * pas1.axe3.y;

	 B[1].z = 
				pas2.axe1.y * pas1.axe1.z + 
				pas2.axe2.y * pas1.axe2.z +
				pas2.axe3.y * pas1.axe3.z;

	 B[2].x = 
				pas2.axe1.z * pas1.axe1.x + 
				pas2.axe2.z * pas1.axe2.x +
				pas2.axe3.z * pas1.axe3.x;

	 B[2].y = 
				pas2.axe1.z * pas1.axe1.y + 
				pas2.axe2.z * pas1.axe2.y +
				pas2.axe3.z * pas1.axe3.y;

	 B[2].z = 
				pas2.axe1.z * pas1.axe1.z + 
				pas2.axe2.z * pas1.axe2.z +
				pas2.axe3.z * pas1.axe3.z;


	 // backup secondMask
	 i3d::Image3d<bool> backupMask(secondMask), alignedMask;
	 i3d::Vector3d<size_t> ipos, ipos_next;
	 i3d::Vector3d<float> fpos, fpos_next;
	 i3d::Vector3d<float> res = secondMask.GetResolution().GetRes();

	 alignedMask.CopyMetaData(secondMask);
	
	 // Perform transform of each position using the transform matrix B
	 for (size_t i=0; i<alignedMask.GetImageSize(); i++)
	 {
		  ipos = secondMask.GetPos(i);
		  fpos = i3d::PixelsToMicrons(ipos, res) - maskCentre;
		  fpos_next.x = i3d::Dot(fpos, B[0]);
		  fpos_next.y = i3d::Dot(fpos, B[1]);
		  fpos_next.z = i3d::Dot(fpos, B[2]);
		  fpos_next += maskCentre;
		  ipos_next.x = (size_t) roundf(fpos_next.x * res.x + maskCentre.x);
		  ipos_next.y = (size_t) roundf(fpos_next.y * res.y + maskCentre.y);
		  ipos_next.z = (size_t) roundf(fpos_next.z * res.z + maskCentre.z);

		  if ((backupMask.Include(ipos_next)) &&
				(backupMask.GetVoxel(ipos_next) == true))
		  {
				alignedMask.SetVoxel(i, true);
		  }
		  else
		  {
				alignedMask.SetVoxel(i, false);
		  }
	 }

	 i3d::Closing(alignedMask, secondMask, i3d::nb3D_26);
}

//----------------------------------------------------------------
void RenderBlankCellMask(const i3d::Vector3d<float> &res,
								 i3d::Image3d<bool> &mask,
								 size_t cellSize,
								 float shape_variance)
{
	 // prepare the initial estimate of object (hl60) dimensions
	 i3d::Vector3d<float> r;
	 r.x = GetRandomGauss (cellSize/2.0f, (shape_variance/100.0f) * cellSize/2.0f);
	 r.y = cellSize/2.0f + (cellSize/2.0f - r.x);
	 r.z = GetRandomGauss (cellSize/3.0f, (shape_variance/100.0f) * cellSize/3.0f);

	 float xy_rotation = GetRandomGauss(0.0f, 45.0f);
	 float xz_rotation = GetRandomGauss(0.0f, 10.0f);

	 i3d::Vector3d<size_t> sz;
	 float max_r = std::max(r.x, std::max(r.y, r.z));
	 i3d::Vector3d<float> centre(max_r + 1.0f), dims = centre * 2.0f; 
	 
	 sz = i3d::MicronsToPixels(dims, res);
	 mask.MakeRoom(sz); // plus some boundary
	 mask.SetResolution(res);
	 mask.Clear();

	 // create the cellMask object
	 RenderEllipsoid (mask, centre, r, xy_rotation, xz_rotation, true);

	 // slightly deform the initial object
	 i3d::Image3d<float> speedImg;
	 speedImg.CopyMetaData(mask);
	
	 GenerateSpeedImage(speedImg, (r.x + r.y + r.z)/4.5f);
	 float deformationLevel = 0.7f;
	 Deformation(mask, NULL, speedImg, deformationLevel); 
}

//----------------------------------------------------------------

template <class T> void GetLabelMask(const i3d::Image3d<T> &img,
												 T label,
												 i3d::Image3d<bool> &mask)
{
	 size_t x,y,z;
	 size_t min_x, min_y, min_z, max_x, max_y, max_z;

	 // initialize the variables
	 min_x = img.GetSizeX();
	 min_y = img.GetSizeY();
	 min_z = img.GetSizeZ();

	 max_x = max_y = max_z = 0;

	 // find the voxels with given label
	 for (size_t i=0; i<img.GetImageSize(); i++)
	 {
		  if (img.GetVoxel(i) == label)
		  {
				x = img.GetX(i);
				y = img.GetY(i);
				z = img.GetZ(i); 
				
				min_x = (min_x > x) ? x : min_x;
				min_y = (min_y > y) ? y : min_y;
				min_z = (min_z > z) ? z : min_z; 
				
				max_x = (max_x < x) ? x : max_x; 
				max_y = (max_y < y) ? y : max_y;
				max_z = (max_z < z) ? z : max_z;
		  }
	 }

	 //enlarge the mask so that the object, HOPEFULLY, does not touch the border
	 //this is required for proper work of ResetBoundaryPoints() function
	 if (min_x > 2) min_x-=2; else min_x=0;
	 if (min_y > 2) min_y-=2; else min_y=0;
	 if (min_z > 2) min_z-=2; else min_z=0;
	 if (max_x < img.GetSizeX()-2) max_x+=2; else max_x=img.GetSizeX()-1;
	 if (max_y < img.GetSizeY()-2) max_y+=2; else max_y=img.GetSizeY()-1;
	 if (max_z < img.GetSizeZ()-2) max_z+=2; else max_z=img.GetSizeZ()-1;

	 // form the mask boundaries
	 i3d::Vector3d<size_t> label_offset(min_x, min_y, min_z);
	 i3d::Vector3d<size_t> label_size(max_x - min_x + 1,
										  max_y - min_y + 1,
										  max_z - min_z + 1);

	 mask.SetResolution(img.GetResolution());
	 mask.SetOffset(img.GetOffset() + 
						 PixelsToMicrons(label_offset, img.GetResolution()));
	 mask.MakeRoom(label_size);
	 mask.Clear();

	 for (x=0; x<label_size.x; x++)
		  for (y=0; y<label_size.y; y++)
				for (z=0; z<label_size.z; z++)
				{
					 if (img.GetVoxel(label_offset.x + x,
											label_offset.y + y,
											label_offset.z + z) == label)
					 {
						  mask.SetVoxel(x,y,z,true);
					 }
				}
}

//----------------------------------------------------------------


inline bool IsInEllipsoid(float x, float y, float z,
		float sx, float sy, float sz,
		float rx, float ry, float rz,
		float rotate_xy, float rotate_xz)
{
	 // first proceed fast rough check
	if (SQR(x-sx) + SQR(y-sy) + SQR(z-sz) > 
		 SQR(std::max(std::max(rx, ry), rz)))
		  return false;

	float x1, x2, x3, y1, y2, y3, z1, z2, z3;

	// shift to center
	x1 = x - sx;
	y1 = y - sy;
	z1 = z - sz;

	// xy-plane rotation
	x2 = float(x1 * cos(DEG_TO_RAD(rotate_xy)) 
				  - y1 * sin(DEG_TO_RAD(rotate_xy)));

	y2 = float(x1 * sin(DEG_TO_RAD(rotate_xy)) 
				  + y1 * cos(DEG_TO_RAD(rotate_xy)));

	z2 = z1;

	// xz-plane rotation
	x3 = float(x2 * cos(DEG_TO_RAD(rotate_xz)) 
				  - z2 * sin(DEG_TO_RAD(rotate_xz)));

	y3 = y2;

	z3 = float(x2 * sin(DEG_TO_RAD(rotate_xz)) 
				  + z2 * cos(DEG_TO_RAD(rotate_xz)));

	return ((SQR(x3/rx) + SQR(y3/ry) + SQR(z3/rz)) <= 1.0f);
}


//----------------------------------------------------------------

template <class VOXEL> void RenderEllipsoid(
		i3d::Image3d<VOXEL> &img,
		const i3d::Vector3d<float> &center,
		const i3d::Vector3d<float> &radius,
		const float angle_xy,
		const float angle_xz,
		const VOXEL foregroundIntensity)
{
	 i3d::Vector3d<float> res = img.GetResolution().GetRes();
	
	for (size_t i=0; i<img.GetImageSize(); i++)
	{
		if (IsInEllipsoid(
					img.GetX(i)/res.x, img.GetY(i)/res.y, img.GetZ(i)/res.z,
					center.x, center.y, center.z,
					radius.x, radius.y, radius.z,
					angle_xy, angle_xz))
		{
			img.SetVoxel(i, foregroundIntensity);
		}
	}
}

/*************************************************************************/
/* explicit instantiations */
template 
void RenderEllipsoid(i3d::Image3d<bool> &,
							const i3d::Vector3d<float> &,
							const i3d::Vector3d<float> &,
							const float angle_xy,
							const float angle_xz,
							const bool foreground_intensity);

template 
void RenderEllipsoid(i3d::Image3d<i3d::GRAY8> &,
							const i3d::Vector3d<float> &,
							const i3d::Vector3d<float> &,
							const float angle_xy,
							const float angle_xz,
							const i3d::GRAY8 foreground_intensity);

template 
void RenderEllipsoid(i3d::Image3d<i3d::GRAY16> &,
							const i3d::Vector3d<float> &,
							const i3d::Vector3d<float> &,
							const float angle_xy,
							const float angle_xz,
							const i3d::GRAY16 foreground_intensity);

template 
void RenderEllipsoid(i3d::Image3d<size_t> &,
							const i3d::Vector3d<float> &,
							const i3d::Vector3d<float> &,
							const float angle_xy,
							const float angle_xz,
							const size_t foreground_intensity);


template 
void GetLabelMask(const i3d::Image3d<i3d::GRAY8> &img,
						i3d::GRAY8 label,
						i3d::Image3d<bool> &mask);
template 
void GetLabelMask(const i3d::Image3d<i3d::GRAY16> &img,
						i3d::GRAY16 label,
						i3d::Image3d<bool> &mask);

