// BarnesHut.cpp : Defines the entry point for the DLL application.
//

#include "stdafx.h"
#include "assert.h"
#include "D3dx8math.h"
#include "BarnesHut.h"
#include <cmath>


// Ratio of diameter-of-box / distance-to-center-of-box, where if over this number, 
// we continue down the tree.
#define CUTOFF_RATIO	(0.9)


BOOL APIENTRY DllMain( HANDLE hModule, 
                       DWORD  ul_reason_for_call, 
                       LPVOID lpReserved
					 )
{
    switch (ul_reason_for_call)
	{
		case DLL_PROCESS_ATTACH:
		case DLL_THREAD_ATTACH:
		case DLL_THREAD_DETACH:
		case DLL_PROCESS_DETACH:
			break;
    }
    return TRUE;
}


// This is an example of an exported variable
BARNESHUT_API int nBarnesHut=0;

// This is an example of an exported function.
BARNESHUT_API int fnBarnesHut(void)
{
	return 42;
}


// ****************************************
// ** POINT3
// ****************************************

POINT3 POINT3::operator= (POINT3 &pt)
{
	this->x = pt.x;
	this->y = pt.y;
	this->z = pt.z;

	return *this;
}

void::POINT3::Init (void)
{
	x = 0.0;
	y = 0.0;
	z = 0.0;
}


// ****************************************
// ** BODY
// ****************************************

BODY::BODY()
{
	m_dMass = 0.0;
	m_ptCenter.x = 0.0;
	m_ptCenter.y = 0.0;
	m_ptCenter.z = 0.0;
	m_ptCenterNext.x = 0.0;
	m_ptCenterNext.y = 0.0;
	m_ptCenterNext.z = 0.0;
	m_vVelocity.x = 0.0;
	m_vVelocity.y = 0.0;
	m_vVelocity.z = 0.0;
	//memset (&m_dMass, 0, sizeof(GFLOAT)+
}

void BODY::Init (void)
{
	m_dMass = 0.0;
	m_ptCenter.x = 0.0;
	m_ptCenter.y = 0.0;
	m_ptCenter.z = 0.0;
	m_ptCenterNext.x = 0.0;
	m_ptCenterNext.y = 0.0;
	m_ptCenterNext.z = 0.0;
	m_vVelocity.x = 0.0;
	m_vVelocity.y = 0.0;
	m_vVelocity.z = 0.0;
}

BODY BODY::operator= (BODY &body)
{
	this->m_dMass = body.m_dMass;
	this->m_ptCenter.x = body.m_ptCenter.x;
	this->m_ptCenter.y = body.m_ptCenter.y;
	this->m_ptCenter.z = body.m_ptCenter.z;
	this->m_ptCenterNext.x = body.m_ptCenterNext.x;
	this->m_ptCenterNext.y = body.m_ptCenterNext.y;
	this->m_ptCenterNext.z = body.m_ptCenterNext.z;
	this->m_vVelocity.x = body.m_vVelocity.x;
	this->m_vVelocity.y = body.m_vVelocity.y;
	this->m_vVelocity.z = body.m_vVelocity.z;

	return *this;
}


// ****************************************
// ** NODE
// ****************************************

// NODE constructor
NODE::NODE()
{
	int i;

	for (i=0; i<8; i++)
		m_apnodeChild[i] = NULL;

	m_fChildren = FALSE;
	m_iBodyCurrent = 0;
	m_vForce.x = 0.0;
	m_vForce.y = 0.0;
	m_vForce.z = 0.0;
	
}
NODE::NODE(PNODE pnodeParent)
{
	NODE();
	m_pnodeParent = pnodeParent;
}

NODE::~NODE()
{

}


// Set the extents of a node.  Really only useful for the top node, as use otherwise
// can screw up the tree.
void NODE::SetExtents (POINT3 ptOrigin, GFLOAT dWidth)
{
	this->m_ptOrigin = ptOrigin;
	this->m_dWidth = dWidth;
}


// ****************************************
// ** CBarnesHut
// ****************************************

// Default constructor.
CBarnesHut::CBarnesHut()
{ 
	BODY	bodyNULL;
	size_t	i;

	
	i=sizeof(NODE);
	i=sizeof(BODY);
	i=sizeof(CBarnesHut);

	m_iCalculations = 0;

	// Unit cube, w/ a null body.
	CreateNode (&m_pnodeUniverse, NULL, 0, bodyNULL);
	CalculateTimeStep();

}

// Constructor when we've got a list of bodies.
CBarnesHut::CBarnesHut (int *pcAdded, BODY *aBody, int cBodies)
{
	HRESULT	hr;
	BODY	bodyNULL;
	int		i;
	POINT3	ptMax, 
			ptMin;
	GFLOAT	dWidth, 
			dWidthX,
			dWidthY,
			dWidthZ;

	
	m_iCalculations = 0;

	// Create the universe.
	CreateNode (&m_pnodeUniverse, NULL, 0, bodyNULL);

	// Calculate the size of the universe.
	ptMax = aBody[0].m_ptCenter;
	ptMin = aBody[0].m_ptCenter;
	for (i=1; i<cBodies; i++)
	{
		ptMax.x = max(ptMax.x, aBody[i].m_ptCenter.x);
		ptMax.y = max(ptMax.y, aBody[i].m_ptCenter.y);
		ptMax.z = max(ptMax.z, aBody[i].m_ptCenter.z);

		ptMin.x = min(ptMin.x, aBody[i].m_ptCenter.x);
		ptMin.y = min(ptMin.y, aBody[i].m_ptCenter.y);
		ptMin.z = min(ptMin.z, aBody[i].m_ptCenter.z);
	}
	dWidthX = ptMax.x-ptMin.x;
	dWidthY = ptMax.y-ptMin.y;
	dWidthZ = ptMax.z-ptMin.z;
	dWidth = max(dWidthX, max(dWidthY,dWidthZ));	

	// Add 10% buffer zone on all sides of cube.
	ptMin.x -= dWidth * 0.10;
	ptMin.y -= dWidth * 0.10;
	ptMin.z -= dWidth * 0.10;
	dWidth *= 1.20;

	// Set the size of the universe.
	m_pnodeUniverse->SetExtents (ptMin, dWidth);

	// Populate the universe w/ bodies.
	hr = AddBodyList (pcAdded, aBody, cBodies);
	CalculateTimeStep();
}

// Destructor.
CBarnesHut::~CBarnesHut()
{
	DestroyNode (m_pnodeUniverse);
	delete m_pnodeUniverse;
	m_pnodeUniverse = NULL;
}


// Recursive.
void CBarnesHut::DestroyNode (PNODE pnode)
{
	int	i;

	// Kill the children!
	if (pnode->m_fChildren)
	{
		for (i=0; i<8; i++)
		{
			if (pnode->m_apnodeChild[i])
			{
				DestroyNode (pnode->m_apnodeChild[i]);
				delete pnode->m_apnodeChild[i];
				pnode->m_apnodeChild[i] = NULL;
			}
		}
	}


} // DestroyNode()


// Get which child the coordinates reside in.
HRESULT	CBarnesHut::ChildIndexFromPoint(INDEX *piChild, PNODE pnodeThis, POINT3 pt3)
{
	*piChild = 0;

	if (pt3.x >= pnodeThis->m_ptOrigin.x+pnodeThis->m_dWidth/2.0)
		*piChild |= 0x4;

	if (pt3.y >= pnodeThis->m_ptOrigin.y+pnodeThis->m_dWidth/2.0)
		*piChild |= 0x2;

	if (pt3.z >= pnodeThis->m_ptOrigin.z+pnodeThis->m_dWidth/2.0)
		*piChild |= 0x1;

	return S_OK;
} // ChildIndexFromPoint()


// Propagate the mass & location up the tree.
void NODE::Propagate (void)
{
	PNODE	pnodeTmp;
	GFLOAT	dMass;
	POINT3	ptCenter;
	int		cBodies=0;
	int		i;


	pnodeTmp = this;
	
	do
	{
		// Sum the mass and get center of mass
		dMass = 0.0;
		ptCenter.x = 0.0;
		ptCenter.y = 0.0;
		ptCenter.z = 0.0;
		if (pnodeTmp->m_fChildren)
		{
			for (i=0; i<8; i++)
				if (pnodeTmp->m_apnodeChild[i])
				{
					dMass += pnodeTmp->m_apnodeChild[i]->m_body.m_dMass;
					ptCenter.x += pnodeTmp->m_apnodeChild[i]->m_body.m_dMass * pnodeTmp->m_apnodeChild[i]->m_body.m_ptCenter.x;
					ptCenter.y += pnodeTmp->m_apnodeChild[i]->m_body.m_dMass * pnodeTmp->m_apnodeChild[i]->m_body.m_ptCenter.y;
					ptCenter.z += pnodeTmp->m_apnodeChild[i]->m_body.m_dMass * pnodeTmp->m_apnodeChild[i]->m_body.m_ptCenter.z;
				}
			ptCenter.x /= dMass;
			ptCenter.y /= dMass;
			ptCenter.z /= dMass;
	
			// Put child sum into this node
			pnodeTmp->m_body.m_dMass = dMass;
			pnodeTmp->m_body.m_ptCenter = ptCenter;
		}
		// else the sum is already in this node.

		// Go up one level in the tree.
		pnodeTmp = pnodeTmp->m_pnodeParent;

	} while (pnodeTmp);
} // Propagate()


// Create a new child node underneath the specified parent.
HRESULT CBarnesHut::CreateNode (PNODE *ppnodeNew, PNODE pnodeParent, INDEX iChild, BODY bodyNew)
{
	int	i;

	// Allocate memory
	(*ppnodeNew) = new NODE;
	if (!(*ppnodeNew))
		return E_OUTOFMEMORY;

	// Set parent & children nodes
	(*ppnodeNew)->m_pnodeParent = pnodeParent;
	for (i=0; i<8; i++) (*ppnodeNew)->m_apnodeChild[i] = NULL;
	(*ppnodeNew)->m_fChildren = FALSE;

	// Contains the specified body.
	(*ppnodeNew)->m_body = bodyNew;

	// If it's a subnode...
	if (pnodeParent)
	{
		// Give this node the proper subcoordinates.
		(*ppnodeNew)->m_ptOrigin.x = pnodeParent->m_ptOrigin.x + (pnodeParent->m_dWidth/2.0)*((iChild & 0x04) >> 2);
		(*ppnodeNew)->m_ptOrigin.y = pnodeParent->m_ptOrigin.y + (pnodeParent->m_dWidth/2.0)*((iChild & 0x02) >> 1);
		(*ppnodeNew)->m_ptOrigin.z = pnodeParent->m_ptOrigin.z + (pnodeParent->m_dWidth/2.0)*((iChild & 0x01));
		(*ppnodeNew)->m_dWidth = pnodeParent->m_dWidth / 2.0;

		// Put this new node into its parent.
		pnodeParent->m_apnodeChild[iChild] = *ppnodeNew;
		pnodeParent->m_fChildren = TRUE;
	}
	else // else it must be the top node (i.e. whole universe)
	{
		(*ppnodeNew)->m_ptOrigin.x=0.0;
		(*ppnodeNew)->m_ptOrigin.y=0.0;
		(*ppnodeNew)->m_ptOrigin.z=0.0;
		(*ppnodeNew)->m_dWidth = 1.0;
	}


	// Propagate body to top of tree
	(*ppnodeNew)->Propagate();


	return S_OK;
} // CreateNode()


// Add a body to a pre-existing node.  Returns the node (possibly new) of the body.
// pnodeThis==NULL means the whole universe.  This function is recursive, and may
// require extra stack space, although the tree depth should be moderate unless the
// dynamic range of the universe is excessive, (e.g. depth==9 for 100,000 randomly
// placed bodies).
// Recursive.
HRESULT	CBarnesHut::AddBody(PNODE *ppnodeNew, PNODE pnodeThis, BODY bodyNew)
{
	HRESULT	 hr=S_OK;
	INDEX	 ichildExisting,
			 ichildNew;


	if (!pnodeThis)
		pnodeThis = m_pnodeUniverse;

	if (!pnodeThis->m_fChildren)
	{
		// Case 1: no children, node is empty.  Just add the body to it.
		if (pnodeThis->m_body.m_dMass==0.0)
			pnodeThis->m_body = bodyNew;
		

		// Case 2: no children, node contains a body.  
		// If node contains exactly one body, then create childA, move the existing 
		// body to new childA, get childB index of new body, and either create a new childB 
		// node for it, or (if occupied) recurse down to now-existing new childA
		// and start over.
		else if (pnodeThis->m_body.m_dMass > 0.0)
		{
			ChildIndexFromPoint (&ichildExisting, pnodeThis, pnodeThis->m_body.m_ptCenter);
			hr = CreateNode (ppnodeNew, pnodeThis, ichildExisting, pnodeThis->m_body);
			if (FAILED(hr))
				return hr;

			ChildIndexFromPoint (&ichildNew, pnodeThis, bodyNew.m_ptCenter);

			if (ichildExisting != ichildNew)
				// Case 2a: new body goes on null child
				hr = CreateNode (ppnodeNew, pnodeThis, ichildNew, bodyNew);
			else
				// Case 2b: new body goes on occupied child
				hr = AddBody (ppnodeNew, *ppnodeNew, bodyNew);

		}
	}
	else	// Case 3: there are children
	{
		ChildIndexFromPoint (&ichildNew, pnodeThis, bodyNew.m_ptCenter);

		if (!pnodeThis->m_apnodeChild[ichildNew])
			// Case 3a: new body goes on null child
			hr = CreateNode (ppnodeNew, pnodeThis, ichildNew, bodyNew);
		else
			// Case 3b: new body goes on occupied child
			hr = AddBody (ppnodeNew, pnodeThis->m_apnodeChild[ichildNew], bodyNew);
	}

	return hr;
} // AddBody()


// Add an array of bodies to the universe.
HRESULT CBarnesHut::AddBodyList (int *cAdded, BODY *aBody, int cBodies)
{
	HRESULT	hr;
	PNODE	pnode;

	for (*cAdded=0; *cAdded<cBodies; (*cAdded)++)
	{
		hr = AddBody (&pnode, NULL, aBody[*cAdded]);
		if (FAILED(hr))
			return hr;
	}

	return S_OK;
} // AddBodyList()


// Get an array of bodies from the universe.  Caller must allocate space for aBody. 
// Number of bodies retrieved is placed in *pcBodies.
void CBarnesHut::GetBodyList (BODY *aBody, int *pcBodies)
{
	*pcBodies = 0;
	GetBodyListA (m_pnodeUniverse, aBody, pcBodies);

} // GetBodyList()

// Recursive part of GetBodyList().
void CBarnesHut::GetBodyListA (PNODE pnodeThis, BODY *aBody, int *pcBodies)
{
	// Recurse down to the leaves.
	if (pnodeThis->m_fChildren)
	{
		for (int i=0; i<8; i++)
		{
			if (pnodeThis->m_apnodeChild[i])
			{
				GetBodyListA (pnodeThis->m_apnodeChild[i], aBody, pcBodies);
			}
		}
		return;
	}

	// When we've reached a leaf, add it to aBody.
	aBody[*pcBodies].m_dMass = pnodeThis->m_body.m_dMass;
	aBody[*pcBodies].m_ptCenter = pnodeThis->m_body.m_ptCenterNext;
	aBody[*pcBodies].m_vVelocity = pnodeThis->m_body.m_vVelocity;
	(*pcBodies)++;

} // GetBodyListA()


// Return the number of levels the tree has.
int CBarnesHut::GetTreeDepth (void)
{
	m_iTreeDepth = 0;
	GetTreeDepthA (m_pnodeUniverse, m_iTreeDepth);

	return m_iTreeDepth;

} // GetTreeDepth()

// Recursive part of GetTreeDepth().
void CBarnesHut::GetTreeDepthA (PNODE pnodeThis, int iTreeDepth)
{
	// Recurse down to the leaves.
	if (pnodeThis->m_fChildren)
	{
		for (int i=0; i<8; i++)
		{
			if (pnodeThis->m_apnodeChild[i])
			{
				GetTreeDepthA (pnodeThis->m_apnodeChild[i], iTreeDepth+1);
			}
		}
		return;
	}

	// We've reached a leaf.
	m_iTreeDepth = max(m_iTreeDepth, iTreeDepth);	

} // GetTreeDepthA()


// How many calculations we did for one time step.
int CBarnesHut::GetNumberOfCalculations (void)
{
	return m_iCalculations;
}


// Move the whole universe ahead one step
void CBarnesHut::AdvanceUniverse (void)
{
	AdvanceNode (m_pnodeUniverse);
}
// The recursive part of AdvanceUniverse()
void CBarnesHut::AdvanceNode (PNODE pnode)
{
	int		i;
	VECTOR3	a;	// acceleration

	// Recurse down to the leaves.
	if (pnode->m_fChildren)
	{
		for (i=0; i<8; i++)
			if (pnode->m_apnodeChild[i])
				AdvanceNode (pnode->m_apnodeChild[i]);

		return;
	}

	// Once at a leaf, get the forces from all the other bodies.
	pnode->m_vForce.x = 0.0;
	pnode->m_vForce.y = 0.0;
	pnode->m_vForce.z = 0.0;
	SumForces (pnode, m_pnodeUniverse);
	
	// F=ma, a=F/m
	if (pnode->m_body.m_dMass > 0.0)
	{
		a.x = pnode->m_vForce.x / pnode->m_body.m_dMass;
		a.y = pnode->m_vForce.y / pnode->m_body.m_dMass;
		a.z = pnode->m_vForce.z / pnode->m_body.m_dMass;
	}
	else
	{
		a.Init();
	}

	// vx = vx0 + ax * t
	pnode->m_body.m_vVelocity.x = pnode->m_body.m_vVelocity.x + a.x * m_tStep;
	pnode->m_body.m_vVelocity.y = pnode->m_body.m_vVelocity.y + a.y * m_tStep;
	pnode->m_body.m_vVelocity.z = pnode->m_body.m_vVelocity.z + a.z * m_tStep;

	// x = x0 + vx * t
	pnode->m_body.m_ptCenterNext.x = pnode->m_body.m_ptCenter.x + pnode->m_body.m_vVelocity.x * m_tStep;
	pnode->m_body.m_ptCenterNext.y = pnode->m_body.m_ptCenter.y + pnode->m_body.m_vVelocity.y * m_tStep;
	pnode->m_body.m_ptCenterNext.z = pnode->m_body.m_ptCenter.z + pnode->m_body.m_vVelocity.z * m_tStep;

	// ...and we're there!
}


// Sum up the forces acting on pLeaf.
// (Recursive).
void CBarnesHut::SumForces (PNODE pnodeLeaf, PNODE pnodeThis)
{
	int				i;
	VECTOR3			vDistance;
	VECTOR3			vDistanceSq;
	//GFLOAT		dRadiusLeaf, dRadiusThis;
	GFLOAT			dDistanceCubed;
	GFLOAT			dDistance;
	GFLOAT			k;
	const GFLOAT	G = 6.67E-11F;


	// Don't calculate force from ourself.
	if (pnodeLeaf==pnodeThis)
		return;

	// Distance from center of cube to center of cube.
	//dRadiusLeaf = pnodeLeaf->m_dWidth/2.0;
	//dRadiusThis = pnodeThis->m_dWidth/2.0;
	//vDistance.x = ((pnodeLeaf->m_ptOrigin.x + dRadiusLeaf) - (pnodeThis->m_ptOrigin.x + dRadiusThis));
	//vDistance.y = ((pnodeLeaf->m_ptOrigin.y + dRadiusLeaf) - (pnodeThis->m_ptOrigin.y + dRadiusThis));
	//vDistance.z = ((pnodeLeaf->m_ptOrigin.z + dRadiusLeaf) - (pnodeThis->m_ptOrigin.z + dRadiusThis));
	//vDistanceSq.x = vDistance.x * vDistance.x;
	//vDistanceSq.y = vDistance.y * vDistance.y;
	//vDistanceSq.z = vDistance.z * vDistance.z;

	// Get distance between bodies.
	vDistance.x = (pnodeLeaf->m_body.m_ptCenter.x - pnodeThis->m_body.m_ptCenter.x);
	vDistance.y = (pnodeLeaf->m_body.m_ptCenter.y - pnodeThis->m_body.m_ptCenter.y);
	vDistance.z = (pnodeLeaf->m_body.m_ptCenter.z - pnodeThis->m_body.m_ptCenter.z);
	vDistanceSq.x = vDistance.x * vDistance.x;
	vDistanceSq.y = vDistance.y * vDistance.y;
	vDistanceSq.z = vDistance.z * vDistance.z;
	dDistance = sqrt(vDistanceSq.x + vDistanceSq.y + vDistanceSq.z);

	
	// Recurse down until hit a node where D/r < cutoff (but if there's no children, then fall through).
	if ((pnodeThis->m_dWidth / dDistance) >= CUTOFF_RATIO)
	{
		if (pnodeThis->m_fChildren)
		{
			for (i=0; i<8; i++)
				if (pnodeThis->m_apnodeChild[i])
					SumForces (pnodeLeaf, pnodeThis->m_apnodeChild[i]);
			return;
		}
	}



	// Add force between pnodeLeaf and pnodeThis
	dDistanceCubed = pow(vDistanceSq.x + vDistanceSq.y + vDistanceSq.z, 3.0/2.0);
	k = (G 
			* pnodeLeaf->m_body.m_dMass 
			* pnodeThis->m_body.m_dMass) 
		  / 
			dDistanceCubed;

	pnodeLeaf->m_vForce.x += k * vDistance.x;
	pnodeLeaf->m_vForce.y += k * vDistance.y;
	pnodeLeaf->m_vForce.z += k * vDistance.z;

	m_iCalculations++;

} // AdvanceNode()


// Best (?) time step depends on the closest two bodies, and their relative velocity.
void CBarnesHut::CalculateTimeStep(void)
{
	m_tStep = 1.0;
}