//========= Copyright Valve Corporation, All rights reserved. ============//
//
// Purpose: 
//
//=============================================================================//

#include "vrad.h"
#include "leaf_ambient_lighting.h"
#include "bsplib.h"
#include "vraddetailprops.h"
#include "mathlib/anorms.h"
#include "pacifier.h"
#include "coordsize.h"
#include "vstdlib/random.h"
#include "bsptreedata.h"
#include "messbuf.h"
#include "vmpi.h"
#include "vmpi_distribute_work.h"

static TableVector g_BoxDirections[6] = 
{
	{  1,  0,  0 }, 
	{ -1,  0,  0 },
	{  0,  1,  0 }, 
	{  0, -1,  0 }, 
	{  0,  0,  1 }, 
	{  0,  0, -1 }, 
};



static void ComputeAmbientFromSurface( dface_t *surfID, dworldlight_t* pSkylight, 
									   Vector& radcolor )
{
	if ( !surfID )
		return;

	texinfo_t *pTexInfo = &texinfo[surfID->texinfo];

	// If we hit the sky, use the sky ambient
	if ( pTexInfo->flags & SURF_SKY )
	{
		if ( pSkylight )
		{
			// add in sky ambient
			VectorCopy( pSkylight->intensity, radcolor );
		}
	}
	else
	{
		Vector reflectivity = dtexdata[pTexInfo->texdata].reflectivity;
		VectorMultiply( radcolor, reflectivity, radcolor );
	}
}


// TODO: it's CRAZY how much lighting code we share with the engine. It should all be shared code.
float Engine_WorldLightAngle( const dworldlight_t *wl, const Vector& lnormal, const Vector& snormal, const Vector& delta )
{
	float dot, dot2;

	Assert( wl->type == emit_surface );

	dot = DotProduct( snormal, delta );
	if (dot < 0)
		return 0;

	dot2 = -DotProduct (delta, lnormal);
	if (dot2 <= ON_EPSILON/10)
		return 0; // behind light surface

	return dot * dot2;
}


// TODO: it's CRAZY how much lighting code we share with the engine. It should all be shared code.
float Engine_WorldLightDistanceFalloff( const dworldlight_t *wl, const Vector& delta )
{
	Assert( wl->type == emit_surface );

	// Cull out stuff that's too far
	if (wl->radius != 0)
	{
		if ( DotProduct( delta, delta ) > (wl->radius * wl->radius))
			return 0.0f;
	}

	return InvRSquared(delta);
}


void AddEmitSurfaceLights( const Vector &vStart, Vector lightBoxColor[6] )
{
	fltx4 fractionVisible;

	FourVectors vStart4, wlOrigin4;
	vStart4.DuplicateVector ( vStart );

	for ( int iLight=0; iLight < *pNumworldlights; iLight++ )
	{
		dworldlight_t *wl = &dworldlights[iLight];

		// Should this light even go in the ambient cubes?
		if ( !( wl->flags & DWL_FLAGS_INAMBIENTCUBE ) )
			continue;

		Assert( wl->type == emit_surface );

		// Can this light see the point?
		wlOrigin4.DuplicateVector ( wl->origin );
		TestLine ( vStart4, wlOrigin4, &fractionVisible );
		if ( !TestSignSIMD ( CmpGtSIMD ( fractionVisible, Four_Zeros ) ) )
			continue;

		// Add this light's contribution.
		Vector vDelta = wl->origin - vStart;
		float flDistanceScale = Engine_WorldLightDistanceFalloff( wl, vDelta );

		Vector vDeltaNorm = vDelta;
		VectorNormalize( vDeltaNorm );
		float flAngleScale = Engine_WorldLightAngle( wl, wl->normal, vDeltaNorm, vDeltaNorm );

		float ratio = flDistanceScale * flAngleScale * SubFloat ( fractionVisible, 0 );
		if ( ratio == 0 )
			continue;

		for ( int i=0; i < 6; i++ )
		{
			float t = DotProduct( g_BoxDirections[i], vDeltaNorm );
			if ( t > 0 )
			{
				lightBoxColor[i] += wl->intensity * (t * ratio);
			}
		}
	}	
}


void ComputeAmbientFromSphericalSamples( int iThread, const Vector &vStart, Vector lightBoxColor[6] )
{
	// Figure out the color that rays hit when shot out from this position.
	Vector radcolor[NUMVERTEXNORMALS];
	float tanTheta = tan(VERTEXNORMAL_CONE_INNER_ANGLE);

	for ( int i = 0; i < NUMVERTEXNORMALS; i++ )
	{
		Vector vEnd = vStart + g_anorms[i] * (COORD_EXTENT * 1.74);

		// Now that we've got a ray, see what surface we've hit
		Vector lightStyleColors[MAX_LIGHTSTYLES];
		lightStyleColors[0].Init();	// We only care about light style 0 here.
		CalcRayAmbientLighting( iThread, vStart, vEnd, tanTheta, lightStyleColors );
	
		radcolor[i] = lightStyleColors[0];
	}

	// accumulate samples into radiant box
	for ( int j = 6; --j >= 0; )
	{
		float t = 0;

		lightBoxColor[j].Init();

		for (int i = 0; i < NUMVERTEXNORMALS; i++)
		{
			float c = DotProduct( g_anorms[i], g_BoxDirections[j] );
			if (c > 0)
			{
				t += c;
				lightBoxColor[j] += radcolor[i] * c;
			}
		}
		
		lightBoxColor[j] *= 1/t;
	}

	// Now add direct light from the emit_surface lights. These go in the ambient cube because
	// there are a ton of them and they are often so dim that they get filtered out by r_worldlightmin.
	AddEmitSurfaceLights( vStart, lightBoxColor );
}


bool IsLeafAmbientSurfaceLight( dworldlight_t *wl )
{
	static const float g_flWorldLightMinEmitSurface = 0.005f;
	static const float g_flWorldLightMinEmitSurfaceDistanceRatio = ( InvRSquared( Vector( 0, 0, 512 ) ) );

	if ( wl->type != emit_surface )
		return false;

	if ( wl->style != 0 )
		return false;

	float intensity = max( wl->intensity[0], wl->intensity[1] );
	intensity = max( intensity, wl->intensity[2] );
	
	return (intensity * g_flWorldLightMinEmitSurfaceDistanceRatio) < g_flWorldLightMinEmitSurface;
}


class CLeafSampler
{
public:
	CLeafSampler( int iThread ) : m_iThread(iThread) {}

	// Generate a random point in the leaf's bounding volume
	// reject any points that aren't actually in the leaf
	// do a couple of tracing heuristics to eliminate points that are inside detail brushes 
	// or underneath displacement surfaces in the leaf
	// return once we have a valid point, use the center if one can't be computed quickly
	void GenerateLeafSamplePosition( int leafIndex, const CUtlVector<dplane_t> &leafPlanes, Vector &samplePosition )
	{
		dleaf_t *pLeaf = dleafs + leafIndex;

		float dx = pLeaf->maxs[0] - pLeaf->mins[0];
		float dy = pLeaf->maxs[1] - pLeaf->mins[1];
		float dz = pLeaf->maxs[2] - pLeaf->mins[2];
		bool bValid = false;
		for ( int i = 0; i < 1000 && !bValid; i++ )
		{
			samplePosition.x = pLeaf->mins[0] + m_random.RandomFloat(0, dx);
			samplePosition.y = pLeaf->mins[1] + m_random.RandomFloat(0, dy);
			samplePosition.z = pLeaf->mins[2] + m_random.RandomFloat(0, dz);
			bValid = true;

			for ( int j = leafPlanes.Count(); --j >= 0 && bValid; )
			{
				float d = DotProduct(leafPlanes[j].normal, samplePosition) - leafPlanes[j].dist;
				if ( d < DIST_EPSILON )
				{
					// not inside the leaf, try again 
					bValid = false;
					break;
				}
			}
			if ( !bValid )
				continue;

			for ( int j = 0; j < 6; j++ )
			{
				Vector start = samplePosition;
				int axis = j%3;
				start[axis] = (j<3) ? pLeaf->mins[axis] : pLeaf->maxs[axis];
				float t;
				Vector normal;
				CastRayInLeaf( m_iThread, samplePosition, start, leafIndex, &t, &normal );
				if ( t == 0.0f )
				{
					// inside a func_detail, try again.
					bValid = false;
					break;
				}
				if ( t != 1.0f )
				{
					Vector delta = start - samplePosition;
					if ( DotProduct(delta, normal) > 0 )
					{
						// hit backside of displacement, try again.
						bValid = false;
						break;
					}
				}
			}
		}
		if ( !bValid )
		{
			// didn't generate a valid sample point, just use the center of the leaf bbox
			samplePosition = ( Vector( pLeaf->mins[0], pLeaf->mins[1], pLeaf->mins[2] ) + Vector( pLeaf->maxs[0], pLeaf->maxs[1], pLeaf->maxs[2] ) ) * 0.5f;
		}
	}

private:
	int		m_iThread;
	CUniformRandomStream m_random;
};

// gets a list of the planes pointing into a leaf
void GetLeafBoundaryPlanes( CUtlVector<dplane_t> &list, int leafIndex )
{
	list.RemoveAll();
	int nodeIndex = leafparents[leafIndex];
	int child = -(leafIndex + 1);
	while ( nodeIndex >= 0 )
	{
		dnode_t *pNode = dnodes + nodeIndex;
		dplane_t *pNodePlane = dplanes + pNode->planenum;
		if ( pNode->children[0] == child )
		{
			// front side
			list.AddToTail( *pNodePlane );
		}
		else
		{
			// back side
			int plane = list.AddToTail();
			list[plane].dist = -pNodePlane->dist;
			list[plane].normal = -pNodePlane->normal;
			list[plane].type = pNodePlane->type;
		}
		child = nodeIndex;
		nodeIndex = nodeparents[child];
	}
}

// this stores each sample of the ambient lighting
struct ambientsample_t
{
	Vector pos;
	Vector cube[6];
};

// add the sample to the list.  If we exceed the maximum number of samples, the worst sample will
// be discarded.  This has the effect of converging on the best samples when enough are added.
void AddSampleToList( CUtlVector<ambientsample_t> &list, const Vector &samplePosition, Vector *pCube )
{
	const int MAX_SAMPLES = 16;

	int index = list.AddToTail();
	list[index].pos = samplePosition;
	for ( int i = 0; i < 6; i++ )
	{
		list[index].cube[i] = pCube[i];
	}

	if ( list.Count() <= MAX_SAMPLES )
		return;

	int nearestNeighborIndex = 0;
	float nearestNeighborDist = FLT_MAX;
	float nearestNeighborTotal = 0;
	for ( int i = 0; i < list.Count(); i++ )
	{
		int closestIndex = 0;
		float closestDist = FLT_MAX;
		float totalDC = 0;
		for ( int j = 0; j < list.Count(); j++ )
		{
			if ( j == i )
				continue;
			float dist = (list[i].pos - list[j].pos).Length();
			float maxDC = 0;
			for ( int k = 0; k < 6; k++ )
			{
				// color delta is computed per-component, per cube side
				for (int s = 0; s < 3; s++ )
				{
					float dc = fabs(list[i].cube[k][s] - list[j].cube[k][s]);
					maxDC = max(maxDC,dc);
				}
				totalDC += maxDC;
			}
			// need a measurable difference in color or we'll just rely on position
			if ( maxDC < 1e-4f )
			{
				maxDC = 0;
			}
			else if ( maxDC > 1.0f )
			{
				maxDC = 1.0f;
			}
			// selection criteria is 10% distance, 90% color difference
			// choose samples that fill the space (large distance from each other)
			// and have largest color variation
			float distanceFactor = 0.1f + (maxDC * 0.9f);
			dist *= distanceFactor;

			// find the "closest" sample to this one
			if ( dist < closestDist )
			{
				closestDist = dist;
				closestIndex = j;
			}
		}
		// the sample with the "closest" neighbor is rejected
		if ( closestDist < nearestNeighborDist || (closestDist == nearestNeighborDist && totalDC < nearestNeighborTotal) )
		{
			nearestNeighborDist = closestDist;
			nearestNeighborIndex = i;
		}
	}
	list.FastRemove( nearestNeighborIndex );
}

// max number of units in gamma space of per-side delta
int CubeDeltaGammaSpace( Vector *pCube0, Vector *pCube1 )
{
	int maxDelta = 0;
	// do this comparison in gamma space to try and get a perceptual basis for the compare
	for ( int i = 0; i < 6; i++ )
	{
		for ( int j = 0; j < 3; j++ )
		{
			int val0 = LinearToScreenGamma( pCube0[i][j] );
			int val1 = LinearToScreenGamma( pCube1[i][j] );
			int delta = abs(val0-val1);
			if ( delta > maxDelta )
				maxDelta = delta;
		}
	}
	return maxDelta;
}
// reconstruct the ambient lighting for a leaf at the given position in worldspace
// optionally skip one of the entries in the list
void Mod_LeafAmbientColorAtPos( Vector *pOut, const Vector &pos, const CUtlVector<ambientsample_t> &list, int skipIndex )
{
	for ( int i = 0; i < 6; i++ )
	{
		pOut[i].Init();
	}
	float totalFactor = 0;
	for ( int i = 0; i < list.Count(); i++ )
	{
		if ( i == skipIndex )
			continue;
		// do an inverse squared distance weighted average of the samples to reconstruct 
		// the original function
		float dist = (list[i].pos - pos).LengthSqr();
		float factor = 1.0f / (dist + 1.0f);
		totalFactor += factor;
		for ( int j = 0; j < 6; j++ )
		{
			pOut[j] += list[i].cube[j] * factor;
		}
	}
	for ( int i = 0; i < 6; i++ )
	{
		pOut[i] *= (1.0f / totalFactor);
	}
}

// this samples the lighting at each sample and removes any unnecessary samples
void CompressAmbientSampleList( CUtlVector<ambientsample_t> &list )
{
	Vector testCube[6];
	for ( int i = 0; i < list.Count(); i++ )
	{
		if ( list.Count() > 1 )
		{
			Mod_LeafAmbientColorAtPos( testCube, list[i].pos, list, i );
			if ( CubeDeltaGammaSpace(testCube, list[i].cube) < 3 )
			{
				list.FastRemove(i);
				i--;
			}
		}
	}
}

// basically this is an intersection routine that returns a distance between the boxes
float AABBDistance( const Vector &mins0, const Vector &maxs0, const Vector &mins1, const Vector &maxs1 )
{
	Vector delta;
	for ( int i = 0; i < 3; i++ )
	{
		float greatestMin = max(mins0[i], mins1[i]);
		float leastMax = min(maxs0[i], maxs1[i]);
		delta[i] = (greatestMin < leastMax) ? 0 : (leastMax - greatestMin);
	}
	return delta.Length();
}

// build a list of leaves from a query
class CLeafList : public ISpatialLeafEnumerator
{
public:
	virtual bool EnumerateLeaf( int leaf, int context )
	{
		m_list.AddToTail(leaf);
		return true;
	}

	CUtlVector<int> m_list;
};

// conver short[3] to vector
static void LeafBounds( int leafIndex, Vector &mins, Vector &maxs )
{
	for ( int i = 0; i < 3; i++ )
	{
		mins[i] = dleafs[leafIndex].mins[i];
		maxs[i] = dleafs[leafIndex].maxs[i];
	}
}

// returns the index of the nearest leaf with ambient samples
int NearestNeighborWithLight(int leafID)
{
	Vector mins, maxs;
	LeafBounds( leafID, mins, maxs );
	Vector size = maxs - mins;
	CLeafList leafList;
	ToolBSPTree()->EnumerateLeavesInBox( mins-size, maxs+size, &leafList, 0 );
	float bestDist = FLT_MAX;
	int bestIndex = leafID;
	for ( int i = 0; i < leafList.m_list.Count(); i++ )
	{
		int testIndex = leafList.m_list[i];
		if ( !g_pLeafAmbientIndex->Element(testIndex).ambientSampleCount )
			continue;

		Vector testMins, testMaxs;
		LeafBounds( testIndex, testMins, testMaxs );
		float dist = AABBDistance( mins, maxs, testMins, testMaxs );
		if ( dist < bestDist )
		{
			bestDist = dist;
			bestIndex = testIndex;
		}
	}
	return bestIndex;
}

// maps a float to a byte fraction between min & max
static byte Fixed8Fraction( float t, float tMin, float tMax )
{
	if ( tMax <= tMin )
		return 0;

	float frac = RemapValClamped( t, tMin, tMax, 0.0f, 255.0f );
	return byte(frac+0.5f);
}

CUtlVector< CUtlVector<ambientsample_t> > g_LeafAmbientSamples;

void ComputeAmbientForLeaf( int iThread, int leafID, CUtlVector<ambientsample_t> &list )
{
	CUtlVector<dplane_t> leafPlanes;
	CLeafSampler sampler( iThread );

	GetLeafBoundaryPlanes( leafPlanes, leafID );
	list.RemoveAll();
	// this heuristic tries to generate at least one sample per volume (chosen to be similar to the size of a player) in the space
	int xSize = (dleafs[leafID].maxs[0] - dleafs[leafID].mins[0]) / 32;
	int ySize = (dleafs[leafID].maxs[1] - dleafs[leafID].mins[1]) / 32;
	int zSize = (dleafs[leafID].maxs[2] - dleafs[leafID].mins[2]) / 64;
	xSize = max(xSize,1);
	ySize = max(xSize,1);
	zSize = max(xSize,1);
	// generate update 128 candidate samples, always at least one sample
	int volumeCount = xSize * ySize * zSize;
	if ( g_bFastAmbient )
	{
		// save compute time, only do one sample
		volumeCount = 1;
	}
	int sampleCount = clamp( volumeCount, 1, 128 );
	if ( dleafs[leafID].contents & CONTENTS_SOLID )
	{
		// don't generate any samples in solid leaves
		// NOTE: We copy the nearest non-solid leaf sample pointers into this leaf at the end
		return;
	}
	Vector cube[6];
	for ( int i = 0; i < sampleCount; i++ )
	{
		// compute each candidate sample and add to the list
		Vector samplePosition;
		sampler.GenerateLeafSamplePosition( leafID, leafPlanes, samplePosition );
		ComputeAmbientFromSphericalSamples( iThread, samplePosition, cube );
		// note this will remove the least valuable sample once the limit is reached
		AddSampleToList( list, samplePosition, cube );
	}

	// remove any samples that can be reconstructed with the remaining data
	CompressAmbientSampleList( list );
}

static void ThreadComputeLeafAmbient( int iThread, void *pUserData )
{
	CUtlVector<ambientsample_t> list;
	while (1)
	{
		int leafID = GetThreadWork ();
		if (leafID == -1)
			break;
		list.RemoveAll();
		ComputeAmbientForLeaf(iThread, leafID, list);
		// copy to the output array
		g_LeafAmbientSamples[leafID].SetCount( list.Count() );
		for ( int i = 0; i < list.Count(); i++ )
		{
			g_LeafAmbientSamples[leafID].Element(i) = list.Element(i);
		}
	}
}

void VMPI_ProcessLeafAmbient( int iThread, uint64 iLeaf, MessageBuffer *pBuf )
{
	CUtlVector<ambientsample_t> list;
	ComputeAmbientForLeaf(iThread, (int)iLeaf, list);

	VMPI_SetCurrentStage( "EncodeLeafAmbientResults" );

	// Encode the results.
	int nSamples = list.Count();
	pBuf->write( &nSamples, sizeof( nSamples ) );
	if ( nSamples )
	{
		pBuf->write( list.Base(), list.Count() * sizeof( ambientsample_t ) );
	}
}

//-----------------------------------------------------------------------------
// Called on the master when a worker finishes processing a static prop.
//-----------------------------------------------------------------------------
void VMPI_ReceiveLeafAmbientResults( uint64 leafID, MessageBuffer *pBuf, int iWorker )
{
	// Decode the results.
	int nSamples;
	pBuf->read( &nSamples, sizeof( nSamples ) );

	g_LeafAmbientSamples[leafID].SetCount( nSamples );
	if ( nSamples )
	{
		pBuf->read(g_LeafAmbientSamples[leafID].Base(), nSamples * sizeof(ambientsample_t) );
	}
}


void ComputePerLeafAmbientLighting()
{
	// Figure out which lights should go in the per-leaf ambient cubes.
	int nInAmbientCube = 0;
	int nSurfaceLights = 0;
	for ( int i=0; i < *pNumworldlights; i++ )
	{
		dworldlight_t *wl = &dworldlights[i];
		
		if ( IsLeafAmbientSurfaceLight( wl ) )
			wl->flags |= DWL_FLAGS_INAMBIENTCUBE;
		else
			wl->flags &= ~DWL_FLAGS_INAMBIENTCUBE;
	
		if ( wl->type == emit_surface )
			++nSurfaceLights;

		if ( wl->flags & DWL_FLAGS_INAMBIENTCUBE )
			++nInAmbientCube;
	}

	Msg( "%d of %d (%d%% of) surface lights went in leaf ambient cubes.\n", nInAmbientCube, nSurfaceLights, nSurfaceLights ? ((nInAmbientCube*100) / nSurfaceLights) : 0 );

	g_LeafAmbientSamples.SetCount(numleafs);

	if ( g_bUseMPI )
	{
		// Distribute the work among the workers.
		VMPI_SetCurrentStage( "ComputeLeafAmbientLighting" );
		DistributeWork( numleafs, VMPI_DISTRIBUTEWORK_PACKETID, VMPI_ProcessLeafAmbient, VMPI_ReceiveLeafAmbientResults );
	}
	else
	{
		RunThreadsOn(numleafs, true, ThreadComputeLeafAmbient);
	}

	// now write out the data
	Msg("Writing leaf ambient...");
	g_pLeafAmbientIndex->RemoveAll();
	g_pLeafAmbientLighting->RemoveAll();
	g_pLeafAmbientIndex->SetCount( numleafs );
	g_pLeafAmbientLighting->EnsureCapacity( numleafs*4 );
	for ( int leafID = 0; leafID < numleafs; leafID++ )
	{
		const CUtlVector<ambientsample_t> &list = g_LeafAmbientSamples[leafID];
		g_pLeafAmbientIndex->Element(leafID).ambientSampleCount = list.Count();
		if ( !list.Count() )
		{
			g_pLeafAmbientIndex->Element(leafID).firstAmbientSample = 0;
		}
		else
		{
			g_pLeafAmbientIndex->Element(leafID).firstAmbientSample = g_pLeafAmbientLighting->Count();
			// compute the samples in disk format.  Encode the positions in 8-bits using leaf bounds fractions
			for ( int i = 0; i < list.Count(); i++ )
			{
				int outIndex = g_pLeafAmbientLighting->AddToTail();
				dleafambientlighting_t &light = g_pLeafAmbientLighting->Element(outIndex);

				light.x = Fixed8Fraction( list[i].pos.x, dleafs[leafID].mins[0], dleafs[leafID].maxs[0] );
				light.y = Fixed8Fraction( list[i].pos.y, dleafs[leafID].mins[1], dleafs[leafID].maxs[1] );
				light.z = Fixed8Fraction( list[i].pos.z, dleafs[leafID].mins[2], dleafs[leafID].maxs[2] );
				light.pad = 0;
				for ( int side = 0; side < 6; side++ )
				{
					VectorToColorRGBExp32( list[i].cube[side], light.cube.m_Color[side] );
				}
			}
		}
	}
	for ( int i = 0; i < numleafs; i++ )
	{
		// UNDONE: Do this dynamically in the engine instead.  This will allow us to sample across leaf
		// boundaries always which should improve the quality of lighting in general
		if ( g_pLeafAmbientIndex->Element(i).ambientSampleCount == 0 )
		{
			if ( !(dleafs[i].contents & CONTENTS_SOLID) )
			{
				Msg("Bad leaf ambient for leaf %d\n", i );
			}

			int refLeaf = NearestNeighborWithLight(i);
			g_pLeafAmbientIndex->Element(i).ambientSampleCount = 0;
			g_pLeafAmbientIndex->Element(i).firstAmbientSample = refLeaf;
		}
	}
	Msg("done\n");
}