#ifndef __PVS_H
#define __PVS_H

#include <map>
#include <vector>
#include "common.h"

namespace GtpVisibilityPreprocessor {

class KdNode;
class BspNode;
class Ray;
class Intersectable;
class ViewCell;


/** Information stored with a PVS entry. Consists of the number
	the object was seen from the view cell.
*/
template<typename T, typename S>
class PvsEntry 
{
public:

	PvsEntry() {}

	PvsEntry(T sample, const S &data): mObject(sample), mData(data) {}

	T mObject;
	S mData;

	template<typename T, typename S>
	friend int operator< (const PvsEntry<T, S> &a, const PvsEntry<T, S> &b);
	template<typename T, typename S>
	friend int operator== (const PvsEntry<T, S> &a, const PvsEntry<T, S> &b);
};


template<typename T, typename S>
int operator< (const PvsEntry<T, S> &a, const PvsEntry<T, S> &b)
{
	return a.mObject < b.mObject;
} 

template<typename T, typename S>
int operator== (const PvsEntry<T, S> &a, const PvsEntry<T, S> &b)
{
	return a.mObject == b.mObject;
} 


template<typename T, typename S>
struct LtSample
{
    bool operator()(const PvsEntry<T, S> &a, const PvsEntry<T, S> &b) const
    {
		return a.mObject < b.mObject;
	} 
};

template<typename T, typename S>
int equalSample (const PvsEntry<T, S> &a, const PvsEntry<T, S> &b)
{
	return a.mObject == b.mObject;
} 

/** Information stored with a PVS entry. Consists of the number
	the object was seen from the view cell.
*/
class PvsData {
public:
	PvsData() {}
	PvsData(const float sumPdf):
	mSumPdf(sumPdf) {}
	
	// $$JB in order to return meaningfull values
	// it assumes that the sum pdf has been normalized somehow!!!
	inline float GetVisibility() 
	{ 
		return mSumPdf; 
	}

	/// sum of probability density of visible sample rays
	float mSumPdf;
};


class MailablePvsData 
{
public:
	// sum of probability density of visible sample rays
	float mSumPdf;
	int mCounter;

	MailablePvsData() {}
	MailablePvsData(const float sumPdf):
	mSumPdf(sumPdf) {}

	// $$JB in order to return meaningfull values
	// it assumes that the sum pdf has been normalized somehow!!!
	float GetVisibility() 
	{ 
		return mSumPdf; 
	}

	////////////////////////////
	//  Mailing stuff

	// last mail id -> warning not thread safe!
	// both mailId and mailbox should be unique for each thread!!!
	static int sMailId;
	static int sReservedMailboxes;

	static void NewMail(const int reserve = 1) {
		sMailId += sReservedMailboxes;
		sReservedMailboxes = reserve;
	}

	void Mail() { mMailbox = sMailId; }
	bool Mailed() const { return mMailbox == sMailId; }

	void Mail(const int mailbox) { mMailbox = sMailId + mailbox; }
	bool Mailed(const int mailbox) const { return mMailbox == sMailId + mailbox; }

	int IncMail() { return ++ mMailbox - sMailId; }
	
	//////////////////////////////////////////

protected:

	int mMailbox;

};


template<typename T, typename S>
class PvsIterator
{
public:
PvsIterator<T, S>(){}
	PvsIterator<T, S>(const typename vector<PvsEntry<T, S> >::const_iterator &itCurrent,
					  const typename vector<PvsEntry<T, S> >::const_iterator &itEnd):
	mItCurrent(itCurrent), mItEnd(itEnd)
	{
	}

	bool HasMoreEntries() const 
	{
		return (mItCurrent != mItEnd);
	}

	const PvsEntry<T, S> &Next()
	{
		return *(mItCurrent ++);
	}
	
private:
	typename vector<PvsEntry<T, S> >::const_iterator mItCurrent;
	typename vector<PvsEntry<T, S> >::const_iterator mItEnd;
};


/** Template class representing the Potentially Visible Set (PVS) 
	mainly from a view cell, but also e.g., from objects.
*/
template<typename T, typename S>
class Pvs
{
	template<typename T, typename S>
	friend class PvsIterator;

public:

	Pvs(): mSamples(0), mEntries(), mLastSorted(0) {}

	/** creates pvs and initializes it with the given entries. 
		Assumes that entries are sorted-
	*/
	Pvs(const vector<PvsEntry<T, S> > &samples);
	virtual ~Pvs() {};

	/** Compresses PVS lossless or lossy.
	*/
	int Compress() {return 0;}
	int GetSize() const {return (int)mEntries.size();}
	bool Empty() const {return mEntries.empty();}

	/** Normalize the visibility of entries in order to get 
		comparable results.
	*/
	void NormalizeMaximum();

	/** Merges pvs of a into this pvs.
		Warning: very slow!
	*/
	void MergeInPlace(const Pvs<T, S> &a);

	/** Difference of pvs to pvs b.
		@returns number of different entries.
	*/
	int Diff(const Pvs<T, S> &b);

	/** Finds sample in PVS.
		@param checkDirty if dirty part of the pvs should be checked for entry 
			(warning: linear runtime in dirty part)
		@returns iterator on the sample.
	*/
	bool Find(T sample,
			  typename vector<PvsEntry<T, S> >::iterator &it,
			  const bool checkDirty = true);

	bool GetSampleContribution(T sample, const float pdf, float &contribution);

	/** Adds sample to PVS. 
		@contribution contribution of sample (0 or 1)
		@returns true if sample was not already in PVS.
	*/
	//bool AddSample(T sample, const float pdf, float &contribution);

	/** Adds sample to PVS.
		@returns contribution of sample (0 or 1)
	*/
	float AddSample(T sample, const float pdf);

	/** Adds sample to PVS without checking for presence of the sample
		pvs remains unsorted!
	*/
	void AddSampleDirty(T sample, const float pdf);

	/** Adds sample dirty (on the end of the vector) but
		first checks if sample is already in clean part of the pvs.
	*/
	bool AddSampleDirtyCheck(T sample, const float pdf);//, float &contribution);

	/** Sort pvs entries - this should always be called after a
		sequence of AddSampleDirty calls 
	*/
	void Sort();

	/** Adds sample to PVS. Assumes that the pvs is sorted
		@returns contribution of sample (0 or 1)
	*/
	float AddSamples(const vector<PvsEntry<T, S> > &samples);

	/** Adds sample to PVS.
		@returns PvsData
	*/
	typename std::vector<PvsEntry<T, S> >::iterator AddSample2(T sample, const float pdf);

	/** Subtracts one pvs from another one.
		WARNING: could contains bugs
		@returns new pvs size
	*/
	int SubtractPvs(const Pvs<T, S> &pvs);

	/** Returns PVS data, i.e., how often it was seen from the view cell, 
		and the object itsef.
	*/
	void GetData(const int index, T &entry, S &data);

	/** Collects the PVS entries and returns them in the vector.
	*/
	void CollectEntries(std::vector<T> &entries);

	/** Removes sample from PVS if reference count is zero.
		@param visibleSamples number of references to be removed
	*/
	bool RemoveSample(T sample, const float pdf);

	/** Compute continuous PVS difference 
	*/
	void ComputeContinuousPvsDifference(Pvs<T, S> &pvs, 
										float &pvsReduction,
										float &pvsEnlargement);

	/** Clears the pvs.
	*/
	void Clear(const bool trim = true);

	void Trim();

	static int GetEntrySizeByte(); 
	static float GetEntrySize();

	/** Compute continuous PVS difference 
	*/
	float GetPvsHomogenity(Pvs<T, S> &pvs);

	static void Merge(Pvs<T, S> &mergedPvs, const Pvs<T, S> &a, const Pvs<T, S> &b);

	static void Merge(Pvs<T, S> &mergedPvs, 
					  const typename std::vector<PvsEntry<T, S> >::const_iterator &aBegin,
					  const typename std::vector<PvsEntry<T, S> >::const_iterator &aEnd,
					  const typename std::vector<PvsEntry<T, S> >::const_iterator &bBegin,
					  const typename std::vector<PvsEntry<T, S> >::const_iterator &bEnd,
					  const int aSamples, 
					  const int bSamples);

	int GetSamples() const
	{
		return mSamples;
	}


	bool IsDirty() const
	{
		return mLastSorted < mEntries.size();
	}

	bool RequiresResort() const
	{
		// the last part should not be more than log of the sorted part. this
		// way we can achieve logarithmic behaviour for insertion and find
		const int dirtySize = (int)mEntries.size() - mLastSorted;
		return dirtySize > (int)(log((double)mEntries.size()) / log(2.0));
	}


	int GetLastSorted() const
	{
		return mLastSorted;
	}

	typename PvsIterator<T, S> GetIterator() const;

protected:

	/// vector of PVS entries
	vector<PvsEntry<T, S> > mEntries; 
	
	/// Number of samples used to create the PVS
	int mSamples;
  
	/// Last sorted entry in the pvs (important for find and merge)
	int mLastSorted;  
};


template <typename T, typename S>
Pvs<T, S>::Pvs(const vector<PvsEntry<T, S> > &samples)
{
	mEntries.reserve(samples.size());
	mEntries = samples;
	mLastSorted = 0;
	mSamples = samples.size();
}


template <typename T, typename S>
void Pvs<T, S>::Sort()
{
	std::vector<PvsEntry<T, S> >::iterator it = mEntries.begin() + mLastSorted;
	//std::vector<PvsEntry<T, S> >::const_iterator it = mEntries.begin() + mLastSorted;
	std::vector<PvsEntry<T, S> >::iterator it_end = mEntries.end();

	// throw out double entries
	std::vector<PvsEntry<T, S> >::iterator newEnd = unique(it, it_end);
	sort(it, newEnd);
	//sort(mEntries.begin(), mEntries.end());

	// now merge sorted ranges
	ObjectPvs newPvs;
	Merge(newPvs, 
		  mEntries.begin(), it, 
		  it, newEnd, 
		  mSamples, 0);
	
	mEntries = newPvs.mEntries;
	mLastSorted = (int)mEntries.size();
}


/**
   Compute continuous PVS difference of 'b' with respect to the base PVS (*this).
   Provides separatelly PVS reduction from PVS enlargement.

*/
template <typename T, typename S>
void
Pvs<T, S>::ComputeContinuousPvsDifference(Pvs<T, S> &b,
										  float &pvsReduction,
										  float &pvsEnlargement)
{
	pvsReduction = 0.0f;
	pvsEnlargement = 0.0f;

	// Uses sum of log differences, which corresponds to entropy
	std::vector<PvsEntry<T, S> >::iterator it;

	for (it = b.mEntries.begin(); it != b.mEntries.end(); ++ it) 
	{
		float bSumPdf = (*it).mData.mSumPdf;
		float aSumPdf = 0.0f;

		vector<PvsEntry<T, S> >::iterator oit;
		const bool entryFound = Find((*it).mObject, oit);		

		if (entryFound)
		{
			aSumPdf = (*it).mData.mSumPdf;

			// mark this entry as processed to avoid double counting
			(*it).mData.mSumPdf = -aSumPdf;
		}

#if 0
		const float diff = bSumPdf - aSumPdf;

		if (diff > 0.0f) {
			pvsEnlargement += diff;
		} else {
			pvsReduction += -diff;
		}
#else
		if (!entryFound)
			pvsEnlargement += 1.0f;
#endif
	}

	for (it = mEntries.begin(); it != mEntries.end(); ++ it) 
	{
		float aSumPdf = (*it).mData.mSumPdf;
		float bSumPdf = 0.0f;
		if (aSumPdf < 0.0f) {
		
			// this entry was already accounted for!
			// just revert it back
			(*it).mData.mSumPdf = -aSumPdf;
		} else {
			vector<PvsEntry<T, S> >::iterator oit;
		
			const bool entryFound = b.Find((*it).mObject, oit);
						
			if (entryFound) {
				bSumPdf = (*oit).mData.mSumPdf;
			}
#if 0
			const float diff = bSumPdf - aSumPdf;

			if (diff > 0.0f) {
				pvsEnlargement += diff;
			} else {
				pvsReduction += -diff;
			}

#else
			if (!entryFound) 
				pvsReduction += 1.0f;
#endif
		}
	}
}


template <typename T, typename S>
int Pvs<T, S>::Diff(const Pvs<T, S> &b)
{
	int dif = 0;

	std::vector<PvsEntry<T, S> >::const_iterator it;

	for (it = b.mEntries.begin(); it != b.mEntries.end(); ++ it)
	{
		vector<PvsEntry<T, S> >::iterator bit;
		const bool entryFound = Find((*it).first, bit);

		if (!entryFound) ++ dif;
	}

	return dif;
}


template <typename T, typename S> 
void Pvs<T, S>::MergeInPlace(const Pvs<T, S> &a)
{
	// early exit
	if (a.Empty())
	{
		return;
	}
	else if (Empty())
	{
		mEntries.reserve(a.GetSize());
		mEntries = a.mEntries;
		mSamples = a.mSamples;
		return;
	}

	ObjectPvs interPvs;
	
	Merge(interPvs, *this, a);
	
	mEntries.reserve(interPvs.GetSize());
	mEntries = interPvs.mEntries;
	mSamples = interPvs.mSamples;
}


template <typename T, typename S>
void Pvs<T, S>::Merge(Pvs<T, S> &mergedPvs, const Pvs<T, S> &a, const Pvs<T, S> &b)
{
	std::vector<PvsEntry<T, S> >::const_iterator ait = a.mEntries.begin(), ait_end = a.mEntries.end();
	std::vector<PvsEntry<T, S> >::const_iterator bit = b.mEntries.begin(), bit_end = b.mEntries.end();
	
	Merge(mergedPvs, 
		  ait, ait_end,
		  bit, bit_end,
		  a.mSamples,
		  b.mSamples);
}


template <typename T, typename S>
void Pvs<T, S>::Merge(Pvs<T, S> &mergedPvs, 
					  const typename std::vector<PvsEntry<T, S> >::const_iterator &aBegin,
					  const typename std::vector<PvsEntry<T, S> >::const_iterator &aEnd,
					  const typename std::vector<PvsEntry<T, S> >::const_iterator &bBegin,
					  const typename std::vector<PvsEntry<T, S> >::const_iterator &bEnd,
					  const int aSamples, 
					  const int bSamples)
{
	std::vector<PvsEntry<T, S> >::const_iterator ait = aBegin;
	std::vector<PvsEntry<T, S> >::const_iterator bit = bBegin;
	
	for (; (ait != aEnd); ++ ait)
	{
		Intersectable *aObj = (*ait).mObject;
		Intersectable *bObj = NULL;
		//Intersectable *bObjOld = NULL;
	
		const PvsEntry<T, S> &aEntry = (*ait);

		for (; (bit != bEnd) && ((*bit).mObject <= (*ait).mObject); ++ bit)
		{
			bObj = (*bit).mObject;

			// object found => add up probabilities
			if (bObj == aEntry.mObject)
			{
				PvsData newData(aEntry.mData.mSumPdf + (*bit).mData.mSumPdf);
				PvsEntry<T, S> entry(bObj, newData);
				mergedPvs.mEntries.push_back(entry);
			}
			else
			{
				mergedPvs.mEntries.push_back(*bit);
			}
			
			//bObjOld = bObj;
		}

		// only push back if objects different 
		// (equal case is handled by second loop)
		if (aObj != bObj)
		{
			mergedPvs.mEntries.push_back(*ait);
		}
	}

	// add the rest
	for (; (bit != bEnd); ++ bit)
	{
		mergedPvs.mEntries.push_back(*bit);
	}

	mergedPvs.mSamples = aSamples + bSamples;
}


template <typename T, typename S> void Pvs<T, S>::Clear(const bool trim = true)
{
	mEntries.clear();
	mSamples = 0;
	mLastSorted = 0;

	if (trim)
	{
		vector<PvsEntry<T,S> >().swap(mEntries);
	}
}


template <typename T, typename S> void Pvs<T, S>::Trim()
{
	vector<PvsEntry<T,S> >(mEntries).swap(mEntries);
}


template <typename T, typename S> 
bool Pvs<T, S>::Find(T sample,
					 typename vector<PvsEntry<T, S> >::iterator &it,
					 const bool checkDirty)
{
	bool found = false;

	PvsEntry<T, S> dummy(sample, PvsData());

	// only check clean part
	//mLastSorted = 0;
	vector<PvsEntry<T, S> >::iterator sorted_end = mEntries.begin() + mLastSorted;

	// binary search
	it = lower_bound(mEntries.begin(), sorted_end, dummy);

	if ((it != mEntries.end()) && ((*it).mObject == sample))
		found = true;

	// sample not found yet => search further in the unsorted part
	if (!found && checkDirty)
	{
		vector<PvsEntry<T, S> >::const_iterator dit, dit_end = mEntries.end();

		for (dit = sorted_end; (dit != dit_end) && ((*dit).mObject != sample); ++ dit);

		if ((dit != mEntries.end()) && ((*dit).mObject == sample))
			found = true;
	}
	
	return found;
}


template <typename T, typename S>
void Pvs<T, S>::GetData(const int index, T &entry, S &data)
{
	std::vector<PvsEntry<T, S> >::iterator i = mEntries.begin();
	for (int k = 0; k != index && i != mEntries.end(); ++ i, ++ k);

	entry = (*i).first;
	data = (*i).second;
}


template <typename T, typename S>
float Pvs<T, S>::AddSample(T sample, const float pdf)
{
	++ mSamples;
	
	vector<PvsEntry<T, S> >::iterator it;
	const bool entryFound = Find(sample, it);		

	if (entryFound)
	{	
		S &data = (*it).mData;
		data.mSumPdf += pdf;
		return data.mSumPdf;
	}
	else 
	{
		PvsEntry<T, S> entry(sample, pdf);
		mEntries.insert(it, entry);
		++ mLastSorted;
		return pdf;
	}
}


template <typename T, typename S>
void Pvs<T, S>::AddSampleDirty(T sample, const float pdf)
{
	++ mSamples;
	mEntries.push_back(PvsEntry<T, S>(sample, pdf));
}
					 

template <typename T, typename S>
typename vector< PvsEntry<T, S> >::iterator Pvs<T, S>::AddSample2(T sample, 
																  const float pdf)
{
	++ mSamples;
	
	vector<PvsEntry<T, S> >::iterator it;
	const bool entryFound == Find(sample, it);

	if (entryFound)
	{
		S &data = (*it).second;
		data->mSumPdf += pdf;
	}
	else 
	{
		PvsEntry<T, S> entry(sample, pdf);
		mEntries.insert(it, entry);
		++ mLastSorted;
	}

	return it;
}


/** Adds sample dirty (on the end of the vector) but
	first checks if sample is already in clean part of the pvs.
*/
template <typename T, typename S> 
bool Pvs<T, S>::AddSampleDirtyCheck(T sample,
									const float pdf)
									//,float &contribution)
{
	++ mSamples;

	vector<PvsEntry<T, S> >::iterator it;
	const bool entryFound = Find(sample, it);

	if (entryFound)
	{
		S &data = (*it).mData;

		data.mSumPdf += pdf;
		//contribution = pdf / data.mSumPdf;

		return false;
	}
	else 
	{
		AddSampleDirty(sample, pdf);
		//contribution = 1.0f;

		return true;
	}
}


template <typename T, typename S>
bool Pvs<T, S>::GetSampleContribution(T sample,
									  const float pdf,
									  float &contribution) 
{
	vector<PvsEntry<T, S> >::iterator it;
	const bool entryFound = Find(sample, it);

	if (entryFound)  
	{
		S &data = (*it).mData;
		contribution = pdf / (data.mSumPdf + pdf);
		return false;
	}
	else 
	{
		contribution = 1.0f;
		return true;
	}
}


template <typename T, typename S>
bool Pvs<T, S>::RemoveSample(T sample, const float pdf)
{
	-- mSamples;
	
	vector<PvsEntry<T, S> >::iterator it;
	const bool entryFound = Find(sample, it);

	if (!entryFound)
		return false;

	S &data = (*it).mData;

	data.mSumPdf -= pdf;

	if (data.mSumPdf <= 0.0f)
	{
		mEntries.erase(it);
		-- mLastSorted; // wrong if sample was in tail!!
	}

	return true;
}


template <typename T, typename S>
int Pvs<T, S>::SubtractPvs(const Pvs<T, S> &pvs)
{
	const int samples = mSamples - pvs.mSamples;

	std::vector<PvsEntry<T, S> >::
		const_iterator it, it_end = pvs.mEntries.end();

	// output PVS of view cell
	for (it = pvs.mEntries.begin(); it != it_end; ++ it) 
		RemoveSample((*it).mObject, (*it).mData.mSumPdf);

	mSamples = samples;

	return GetSize();
}


template <typename T, typename S>
void Pvs<T, S>::CollectEntries(std::vector<T> &entries)
{
	std::vector<PvsEntry<T, S> >::
		const_iterator it, it_end = mEntries.end();

	// output PVS of view cell
	for (it = mEntries.begin(); it != it_end; ++ it) 
		entries.push_back((*it)->first);
}


template <typename T, typename S>
void Pvs<T, S>::NormalizeMaximum()
{
	std::vector<PvsEntry<T, S> >::
		const_iterator it, it_end = mEntries.end();

	float maxPdfSum = -1.0f;

	// output PVS of view cell
	for (it = mEntries.begin(); it != it_end; ++ it) {
		float sum = (*it)->second.sumPdf;
		if (sum > maxSum)
			maxSum = sum;
	}

	maxSum = 1.0f / maxSum;

	for (it = mEntries.begin(); it != it_end; ++ it) {
		(*it)->second.sumPdf *= maxSum;
	}

}


template <typename T, typename S>
float Pvs<T, S>::GetEntrySize()
{
	return (float)(sizeof(T) + sizeof(S)) / float(1024 * 1024);
}


template <typename T, typename S>
int Pvs<T, S>::GetEntrySizeByte()
{
	return sizeof(T) + sizeof(S);
}


template <typename T, typename S>
float Pvs<T, S>::GetPvsHomogenity(Pvs<T, S> &pvs) 
{
	float pvsReduction, pvsEnlargement;

	ComputeContinuousPvsDifference(pvs,	pvsReduction, pvsEnlargement);

	return pvsReduction + pvsEnlargement;
}


template <typename T, typename S>
typename PvsIterator<T, S> Pvs<T, S>::GetIterator() const
{
	PvsIterator<T, S> pit(mEntries.begin(), mEntries.end());

	return pit;
}


///////////////////////////////////////

/** Class instantiating the Pvs template for kd tree nodes.
*/
class KdPvs: public Pvs<KdNode *, PvsData>
{
public:
	int Compress();
};


////////////
//-- typedefs

typedef PvsEntry<Intersectable *, PvsData> ObjectPvsEntry;
typedef std::vector<ObjectPvsEntry> ObjectPvsEntries;
typedef Pvs<ViewCell *, MailablePvsData> ViewCellPvs;
typedef PvsIterator<Intersectable *, PvsData> ObjectPvsIterator;


class ObjectPvs: public Pvs<Intersectable *, PvsData>
{
public:
	/** Counts object int the pvs. Different to method "GetSize", not
		only the raw container size is returned,
		but the individual contributions of the entries are summed up.
	*/
	float EvalPvsCost() const;

	friend ostream &operator<<(ostream &s, const ObjectPvs &p) 
	{
		ObjectPvsIterator pit = p.GetIterator();

		while (pit.HasMoreEntries())
		{		
			const ObjectPvsEntry &entry = pit.Next();
			Intersectable *obj = entry.mObject;

			cout << (int)obj << " ";
		}
		
		return s;
	}

};



}

#endif