#include "SamplingStrategy.h"
#include "Ray.h"
#include "Intersectable.h"
#include "Preprocessor.h"
#include "ViewCellsManager.h"
#include "AxisAlignedBox3.h"
#include "RssTree.h"

namespace GtpVisibilityPreprocessor {


//HaltonSequence SamplingStrategy::sHalton;

HaltonSequence ObjectBasedDistribution::sHalton;
HaltonSequence MixtureDistribution::sHalton;
HaltonSequence GlobalLinesDistribution::sHalton;
HaltonSequence SpatialBoxBasedDistribution::sHalton;
HaltonSequence ObjectDirectionBasedDistribution::sHalton;
HaltonSequence DirectionBasedDistribution::sHalton;
HaltonSequence MutationBasedDistribution::sHalton;


SamplingStrategy::SamplingStrategy(Preprocessor &preprocessor): 
  mPreprocessor(preprocessor), mRatio(1.0f),  mTotalRays(0), mTotalContribution(0.0f)
{
  
}


SamplingStrategy::~SamplingStrategy() 
{
}

int
SamplingStrategy::GenerateSamples(const int number,
								  SimpleRayContainer &rays)
{
	SimpleRay ray;
	int samples = 0;
	int i = 0;
	const int maxTries = 20;
	// tmp changed matt. Q: should one rejected sample 
	// terminate the whole method?
	if (0)
	{
		for (; i < number; i++) {
			if (!GenerateSample(ray))
				return i;
			rays.push_back(ray);
		}
	}
	else
	{
		for (; i < number; i++) 
		{
			int j = 0;
			bool sampleGenerated = false;

			for (j = 0; !sampleGenerated && (j < maxTries); ++ j)
			{
				sampleGenerated = GenerateSample(ray);

				if (sampleGenerated)
				{		
					++ samples;
					rays.push_back(ray);
				}
			}
		}
	}
	
	return samples;
}


/*********************************************************************/
/*            Individual sampling strategies implementation          */
/*********************************************************************/


bool ObjectBasedDistribution::GenerateSample(SimpleRay &ray)
{
  Vector3 origin, direction; 
  
  float r[5];
  sHalton.GetNext(5, r);
  
  mPreprocessor.mViewCellsManager->GetViewPoint(origin,
												Vector3(r[2],r[3],r[4]));
  

  Vector3 point, normal;
  
  r[0] *= (float)mPreprocessor.mObjects.size() - 1;
  const int i = (int)r[0];
  
  Intersectable *object = mPreprocessor.mObjects[i];
  
  // take the remainder as a parameter over objects surface
  r[0] -= (float)i;
  
  object->GetRandomSurfacePoint(r[0], r[1], point, normal);
  
  direction = point - origin;
  
  const float c = Magnitude(direction);
  
  if (c <= Limits::Small) 
	return false;
  
  // $$ jb the pdf is yet not correct for all sampling methods!
  const float pdf = 1.0f;
  
  direction *= 1.0f / c;
  ray = SimpleRay(origin, direction, OBJECT_BASED_DISTRIBUTION, pdf);
  
  return true;
}


bool
ObjectDirectionBasedDistribution::GenerateSample(SimpleRay &ray)
{	
  Vector3 origin, direction; 


  float r[4];
  sHalton.GetNext(4, r);

  r[0] *= (float)mPreprocessor.mObjects.size() - 1;
  const int i = (int)r[0];
  
  Intersectable *object = mPreprocessor.mObjects[i];
  
  // take the remainder as a parameter over objects surface
  r[0] -= (float)i;

  Vector3 normal;

  object->GetRandomSurfacePoint(r[0], r[1], origin, normal);
  
  direction = Normalize(CosineRandomVector(r[2], r[3], normal));
  
  origin += 1e-2f*direction;
  
  // $$ jb the pdf is yet not correct for all sampling methods!
  const float pdf = 1.0f;
  
  ray = SimpleRay(origin, direction, OBJECT_DIRECTION_BASED_DISTRIBUTION, pdf);
  
  return true;
}


bool DirectionBasedDistribution::GenerateSample(SimpleRay &ray)
{	

  float r[5];
  sHalton.GetNext(5, r);

  Vector3 origin, direction; 
  mPreprocessor.mViewCellsManager->GetViewPoint(origin,
												Vector3(r[2],r[3],r[4])
												);
  
  direction = UniformRandomVector(r[0], r[1]);
  const float c = Magnitude(direction);
  
  if (c <= Limits::Small) 
	return false;
  
  const float pdf = 1.0f;
  
  direction *= 1.0f / c;
  ray = SimpleRay(origin, direction, DIRECTION_BASED_DISTRIBUTION, pdf);
  
  return true;
}


bool DirectionBoxBasedDistribution::GenerateSample(SimpleRay &ray)
{
	Vector3 origin, direction; 
	mPreprocessor.mViewCellsManager->GetViewPoint(origin);

	const float alpha = RandomValue(0.0f, 2.0f * (float)M_PI);
	const float beta = RandomValue((float)-M_PI * 0.5f, (float)M_PI * 0.5f);
	
	direction = VssRay::GetDirection(alpha, beta);
	
	const float c = Magnitude(direction);

	if (c <= Limits::Small) 
		return false;

	const float pdf = 1.0f;

	direction *= 1.0f / c;
	ray = SimpleRay(origin, direction, DIRECTION_BOX_BASED_DISTRIBUTION, pdf);

	return true;
}


bool SpatialBoxBasedDistribution::GenerateSample(SimpleRay &ray)
{
  Vector3 origin, direction; 

  float r[6];
  sHalton.GetNext(6, r);
  mPreprocessor.mViewCellsManager->GetViewPoint(origin, Vector3(r[0],r[1],r[2]));
  
  direction = mPreprocessor.mKdTree->GetBox().GetRandomPoint(Vector3(r[3],
																	 r[4],
																	 r[5])
																	 ) - origin;
  //cout << "z";
  const float c = Magnitude(direction);
  
  if (c <= Limits::Small) 
	return false;
  
  const float pdf = 1.0f;
  
  direction *= 1.0f / c;
  ray = SimpleRay(origin, direction, SPATIAL_BOX_BASED_DISTRIBUTION, pdf);
  
  return true;
}


bool ReverseObjectBasedDistribution::GenerateSample(SimpleRay &ray)
{
	Vector3 origin, direction; 

	mPreprocessor.mViewCellsManager->GetViewPoint(origin);

	Vector3 point;
	Vector3 normal;
	//cout << "y";
	const int i = (int)RandomValue(0, (float)mPreprocessor.mObjects.size() - 0.5f);

	Intersectable *object = mPreprocessor.mObjects[i];

	object->GetRandomSurfacePoint(point, normal);
	
	direction = origin - point;
	
	// $$ jb the pdf is yet not correct for all sampling methods!
	const float c = Magnitude(direction);
	
	if ((c <= Limits::Small) || (DotProd(direction, normal) < 0))
	{
		return false;
	}

	// $$ jb the pdf is yet not correct for all sampling methods!
	const float pdf = 1.0f;
	//cout << "p: " << point << " ";
	direction *= 1.0f / c;
	// a little offset
	point += direction * 0.001f;

	ray = SimpleRay(point, direction, REVERSE_OBJECT_BASED_DISTRIBUTION, pdf);
	
	return true;
}


bool ViewCellBorderBasedDistribution::GenerateSample(SimpleRay &ray)
{
	Vector3 origin, direction; 

	ViewCellContainer &viewCells = mPreprocessor.mViewCellsManager->GetViewCells();

	Vector3 point;
	Vector3 normal, normal2;
	
	const int vcIdx = (int)RandomValue(0, (float)viewCells.size() - 0.5f);
	const int objIdx = (int)RandomValue(0, (float)mPreprocessor.mObjects.size() - 0.5f);

	Intersectable *object = mPreprocessor.mObjects[objIdx];
	ViewCell *viewCell = viewCells[vcIdx];

	//cout << "vc: " << vcIdx << endl;
	//cout << "obj: " << objIdx << endl;

	object->GetRandomSurfacePoint(point, normal);
	viewCell->GetRandomEdgePoint(origin, normal2);

	direction = point - origin;

	// $$ jb the pdf is yet not correct for all sampling methods!
	const float c = Magnitude(direction);

	if ((c <= Limits::Small) /*|| (DotProd(direction, normal) < 0)*/)
	{
		return false;
	}

	// $$ jb the pdf is yet not correct for all sampling methods!
	const float pdf = 1.0f;
	//cout << "p: " << point << " ";
	direction *= 1.0f / c;
	ray = SimpleRay(origin, direction, VIEWCELL_BORDER_BASED_DISTRIBUTION, pdf);

	//cout << "ray: " << ray.mOrigin << " " << ray.mDirection << endl;

	return true;
}


#if 0
bool ObjectsInteriorDistribution::GenerateSample(SimpleRay &ray)
{
	Vector3 origin, direction; 

	// get random object
	const int i = RandomValue(0, mPreprocessor.mObjects.size() - 1);

	const Intersectable *obj = mPreprocessor.mObjects[i];

	// note: if we load the polygons as meshes, 
	// asymtotically every second sample is lost!
	origin = obj->GetBox().GetRandomPoint();

	// uniformly distributed direction
	direction = UniformRandomVector();

	const float c = Magnitude(direction);

	if (c <= Limits::Small) 
		return false;

	const float pdf = 1.0f;

	direction *= 1.0f / c;
	ray = SimpleRay(origin, direction, pdf);

	return true;
}

#endif


bool ReverseViewSpaceBorderBasedDistribution::GenerateSample(SimpleRay &ray)
{
	Vector3 origin, direction; 

	origin = mPreprocessor.mViewCellsManager->GetViewSpaceBox().GetRandomSurfacePoint();

	Vector3 point;
	Vector3 normal;
	//cout << "y";
	const int i = (int)RandomValue(0, (float)mPreprocessor.mObjects.size() - 0.5f);

	Intersectable *object = mPreprocessor.mObjects[i];

	object->GetRandomSurfacePoint(point, normal);
	
	direction = origin - point;
	
	// $$ jb the pdf is yet not correct for all sampling methods!
	const float c = Magnitude(direction);
	
	if ((c <= Limits::Small) || (DotProd(direction, normal) < 0))
	{
		return false;
	}

	// $$ jb the pdf is yet not correct for all sampling methods!
	const float pdf = 1.0f;
	//cout << "p: " << point << " ";
	direction *= 1.0f / c;
	// a little offset
	point += direction * 0.001f;

	ray = SimpleRay(point, direction, REVERSE_VIEWSPACE_BORDER_BASED_DISTRIBUTION, pdf);
	
	return true;
}


bool ViewSpaceBorderBasedDistribution::GenerateSample(SimpleRay &ray)
{
	Vector3 origin, direction; 

	origin = mPreprocessor.mViewCellsManager->GetViewSpaceBox().GetRandomSurfacePoint();

	Vector3 point;
	Vector3 normal;
	//cout << "w";
	const int i = (int)RandomValue(0, (float)mPreprocessor.mObjects.size() - 0.5f);

	Intersectable *object = mPreprocessor.mObjects[i];

	object->GetRandomSurfacePoint(point, normal);
	direction = point - origin;

	// $$ jb the pdf is yet not correct for all sampling methods!
	const float c = Magnitude(direction);

	if (c <= Limits::Small) 
		return false;

	// $$ jb the pdf is yet not correct for all sampling methods!
	const float pdf = 1.0f;
	
	direction *= 1.0f / c;

	// a little offset
	origin += direction * 0.001f;

	ray = SimpleRay(origin, direction, VIEWSPACE_BORDER_BASED_DISTRIBUTION, pdf);

	return true;
}




bool
GlobalLinesDistribution::GenerateSample(SimpleRay &ray) 
{
  Vector3 origin, termination, direction; 

  float radius = 0.5f*Magnitude(mPreprocessor.mViewCellsManager->GetViewSpaceBox().Size());
  Vector3 center = mPreprocessor.mViewCellsManager->GetViewSpaceBox().Center();

  const int tries = 1000;
  int i;
  for (i=0; i < tries; i++) {
	float r[4];
	sHalton.GetNext(4, r);
	
	origin = center + (radius*UniformRandomVector(r[0], r[1]));
	termination = center + (radius*UniformRandomVector(r[2], r[3]));
	
	direction = termination - origin;
	
	
	// $$ jb the pdf is yet not correct for all sampling methods!
	const float c = Magnitude(direction);
	if (c <= Limits::Small) 
	  return false;
	
	direction *= 1.0f / c;

	// check if the ray intersects the view space box
	static Ray ray;
	ray.Init(origin, direction, Ray::LOCAL_RAY);	
	
	float tmin, tmax;
	if (mPreprocessor.mViewCellsManager->
		GetViewSpaceBox().ComputeMinMaxT(ray, &tmin, &tmax) && (tmin < tmax))
	  break;
  }
  
  if (i!=tries) {
	// $$ jb the pdf is yet not correct for all sampling methods!
	const float pdf = 1.0f;
	
	
	ray = SimpleRay(origin, direction, GLOBAL_LINES_DISTRIBUTION, pdf);
	ray.mType = Ray::GLOBAL_RAY;
	return true;
  }
  
  return false;
}


  // has to called before first usage
void
MixtureDistribution::Init()
{
  for (int i=0; i < mDistributions.size(); i++) {
	// small non-zero value
	mDistributions[i]->mRays = 1;
	// unit contribution per ray
	if (1 || mDistributions[i]->mType != RSS_BASED_DISTRIBUTION)
	  mDistributions[i]->mContribution = 1.0f;
	else
	  mDistributions[i]->mContribution = 0.0f;
  }
  UpdateRatios();
}

void
MixtureDistribution::Reset()
{
  for (int i=0; i < mDistributions.size(); i++) {
	// small non-zero value
	mDistributions[i]->mTotalRays = 0;
	// unit contribution per ray
	mDistributions[i]->mTotalContribution = 0.0f;
  }
  UpdateRatios();
}

// Generate a new sample according to a mixture distribution
bool
MixtureDistribution::GenerateSample(SimpleRay &ray)
{
  float r;
  sHalton.GetNext(1, &r);
  int i;
  // pickup a distribution
  for (i=0; i < mDistributions.size()-1; i++)
	if (r < mDistributions[i]->mRatio)
	  break;

  return mDistributions[i]->GenerateSample(ray);
}

  // add contributions of the sample to the strategies
void
MixtureDistribution::ComputeContributions(VssRayContainer &vssRays)
{
  int i;
  
  VssRayContainer::iterator it = vssRays.begin();

  for (i=0; i < mDistributions.size(); i++) {
	mDistributions[i]->mContribution = 0;
	mDistributions[i]->mRays = 0;
  }

  for(; it != vssRays.end(); ++it) {
	VssRay *ray = *it;
	for (i=0; i < mDistributions.size()-1; i++) {
	  if (mDistributions[i]->mType == ray->mDistribution)
		break;
	}
 
	float contribution =
	  mPreprocessor.mViewCellsManager->ComputeSampleContribution(*ray,
																 true,
																 false);

	mDistributions[i]->mContribution += contribution;
	mDistributions[i]->mRays ++;
	
	mDistributions[i]->mTotalContribution += contribution;
	mDistributions[i]->mTotalRays ++;
  }

  
  UpdateRatios();
}

void
MixtureDistribution::UpdateDistributions(VssRayContainer &vssRays)
{
  // now update the distributions with all the rays
  for (int i=0; i < mDistributions.size(); i++) {
	mDistributions[i]->Update(vssRays);
  }
}
  
void
MixtureDistribution::UpdateRatios()
{
  // now compute importance (ratio) of all distributions
  float sum = 0.0f;
  int i;
  for (i=0; i < mDistributions.size(); i++) {
	cout<<i<<": c="<<mDistributions[i]->mContribution<<
	  " rays="<<mDistributions[i]->mRays<<endl;
	float importance = 0.0f;
	if (mDistributions[i]->mRays != 0) {
	  //	  importance = pow(mDistributions[i]->mContribution/mDistributions[i]->mRays, 3);
	  importance = mDistributions[i]->mContribution/mDistributions[i]->mRays;
	}
	mDistributions[i]->mRatio = importance;
	sum += importance;
  }

  const float minratio = 0.02f;
  float threshold = minratio*sum;
  for (i=0; i < mDistributions.size(); i++) {
	if (mDistributions[i]->mRatio < threshold) {
	  sum += threshold - mDistributions[i]->mRatio;
	  mDistributions[i]->mRatio = threshold;
	}
  }
  
  mDistributions[0]->mRatio /= sum;
	
  for (i=1; i < mDistributions.size(); i++) {
	float r = mDistributions[i]->mRatio / sum;
	mDistributions[i]->mRatio = mDistributions[i-1]->mRatio + r;
  }
  
  cout<<"ratios: ";
  float last = 0.0f;
  for (i=0; i < mDistributions.size(); i++) {
	cout<<mDistributions[i]->mRatio-last<<" ";
	last = mDistributions[i]->mRatio;
  }
  cout<<endl;
}



bool
MixtureDistribution::Construct(char *str)
{
  char *curr = str;

  while (1) {
	char *e = strchr(curr,'+');
	if (e!=NULL) {
	  *e=0;
	}
	
	if (strcmp(curr, "rss")==0) {
	  mDistributions.push_back(new RssBasedDistribution(mPreprocessor));
	} else
	  if (strcmp(curr, "object")==0) {
		mDistributions.push_back(new ObjectBasedDistribution(mPreprocessor));
	  } else
		if (strcmp(curr, "spatial")==0) {
		  mDistributions.push_back(new SpatialBoxBasedDistribution(mPreprocessor));
		} else
		  if (strcmp(curr, "global")==0) {
			mDistributions.push_back(new GlobalLinesDistribution(mPreprocessor));
		  } else
			if (strcmp(curr, "direction")==0) {
			  mDistributions.push_back(new DirectionBasedDistribution(mPreprocessor));
			} else
			  if (strcmp(curr, "object_direction")==0) {
				mDistributions.push_back(new ObjectDirectionBasedDistribution(mPreprocessor));
			  } else
				if (strcmp(curr, "reverse_object")==0) {
				  mDistributions.push_back(new ReverseObjectBasedDistribution(mPreprocessor));
				} else
				  if (strcmp(curr, "reverse_viewspace_border")==0) {
					mDistributions.push_back(new ReverseViewSpaceBorderBasedDistribution(mPreprocessor));
				  } else
					if (strcmp(curr, "mutation")==0) {
					  mDistributions.push_back(new MutationBasedDistribution(mPreprocessor));
					} 
	
	
	if (e==NULL)
	  break;
	curr = e+1;
  }

  Init();
  return true;
}



void
MutationBasedDistribution::Update(VssRayContainer &vssRays)
{
  //  for (int i=0; i < mRays.size(); i++)
  //	cout<<mRays[i].mSamples<<" ";
  //  cout<<endl;
  
  for (int i=0; i < vssRays.size(); i++) {
	VssRay *newRay = vssRays[i];
	if (newRay->mPvsContribution) {
	  // add this ray
	  newRay->Ref();
	  if (mRays.size() < mMaxRays)
		mRays.push_back(RayEntry(newRay));
	  else {
		// unref the old ray
		VssRay *oldRay = mRays[mBufferStart].mRay;
		oldRay->Unref();
		if (oldRay->RefCount() == 0)
		  delete oldRay;
		mRays[mBufferStart] = RayEntry(newRay);
		mBufferStart++;
		if (mBufferStart >= mMaxRays)
		  mBufferStart = 0;
	  }
	}
  }
}

bool
MutationBasedDistribution::GenerateSample(SimpleRay &sray)
{
  float r[5];
  sHalton.GetNext(5, r);

  if (mRays.size() == 0) {
	// use direction based distribution
	Vector3 origin, direction;
	mPreprocessor.mViewCellsManager->GetViewPoint(origin,
												  Vector3(r[0], r[1], r[2]));
	
	direction = UniformRandomVector(r[3], r[4]);
	
	const float pdf = 1.0f;
	sray = SimpleRay(origin, direction, MUTATION_BASED_DISTRIBUTION, pdf);

	return true;
  } 
  
  //  int index = (int) (r[0]*mRays.size());
  //  if (index >= mRays.size()) {
  //	cerr<<"mutation: index out of bounds\n";
  //	exit(1);
  //  }

  // get tail of the buffer
  int index = (mLastIndex+1)%mRays.size();
  if (mRays[index].mSamples > mRays[mLastIndex].mSamples) {
	// search back for index where this is valid
	index = (mLastIndex - 1 + mRays.size())%mRays.size();
	for (int i=0; i < mRays.size(); i++) {
	  if (mRays[index].mSamples > mRays[mLastIndex].mSamples)
		break;
	  index = (index - 1 + mRays.size())%mRays.size();
	}
	// go one step back
	index = (index+1)%mRays.size();
  }
  
  VssRay *ray = mRays[index].mRay;
  mRays[index].mSamples++;
  mLastIndex = index;
  
  Vector3 v;
  // mutate the origin
  Vector3 d = ray->GetDir();
  if (d.DrivingAxis() == 0)
	v = Vector3(0, r[1]-0.5f, r[2]-0.5f);
  else
	if (d.DrivingAxis() == 1)
	  v = Vector3(r[1]-0.5f, 0, r[2]-0.5f);
	else
	  v = Vector3(r[1]-0.5f, r[2]-0.5f, 0);

  v=v*mOriginMutationSize;

  Vector3 origin = ray->mOrigin + v;
  
  // mutate the termination
  if (d.DrivingAxis() == 0)
	v = Vector3(0, r[3]-0.5f, r[4]-0.5f);
  else
	if (d.DrivingAxis() == 1)
	  v = Vector3(r[3]-0.5f, 0, r[4]-0.5f);
	else
	  v = Vector3(r[3]-0.5f, r[4]-0.5f, 0);

  float size = 1.5f*Magnitude(ray->mTerminationObject->GetBox().Diagonal());
  if (size < Limits::Small)
	return false;
  
//   Vector3 nv;
  
//   if (Magnitude(v) > Limits::Small)
// 	nv = Normalize(v);
//   else
// 	nv = v;
  
//  v = nv*size + v*size;

  v = v*size;

  Vector3 termination = ray->mTermination + v;
  
  Vector3 direction = termination - origin;

  if (Magnitude(direction) < Limits::Small)
	return false;

  // shift the origin a little bit
  origin += direction*0.05f;

  direction.Normalize();
  
  // $$ jb the pdf is yet not correct for all sampling methods!
  const float pdf = 1.0f;
  
  sray = SimpleRay(origin, direction, MUTATION_BASED_DISTRIBUTION, pdf);
  return true;
}

MutationBasedDistribution::MutationBasedDistribution(Preprocessor &preprocessor) :
  SamplingStrategy(preprocessor)
{
  mType = MUTATION_BASED_DISTRIBUTION;
  mBufferStart = 0;
  mMaxRays = 1000000;

  mOriginMutationSize = 2.0f;
  mLastIndex = 0;
  //  mOriginMutationSize = Magnitude(preprocessor.mViewCellsManager->
  //								  GetViewSpaceBox().Diagonal())*1e-3;
  
}

}