#include "LCRP.h"

LCRP::LCRP(double& alpha, Generator* labels, double& discount):
	CRP(alpha, labels, discount),
	_discountFactor(0)
{
}

LCRP::~LCRP()
{
}

int LCRP::sample()
{
	assert(0);
}

Prob LCRP::update(int sample, double temp)
{
	// cerr<<"Updating at "<<sample<<"\n";
	// print(cerr);

	Prob res = 0;

	Prob total = 0;

	{
		Prob term = _prior->prob(sample) * (_alpha + _discountFactor);
		term = powl(term, temp);
		total += term;
	}

	if(total == 0)
	{
		// cerr<<"eject for impossibility\n";
		//no change for impossible update
		return 0;
	}

	TableLst& tables = _tables[sample];
	foreach(TableLst, tab, tables)
	{
		Prob term = (*tab)->count;
		term = powl(term, temp);
		total += term;
	}

	Prob sm = gsl_rng_uniform(RNG);
	Prob progress = 0;

	foreach(TableLst, tab, tables)
	{
		Prob term = (*tab)->count;
		term = powl(term, temp);
		progress += term;

		if(sm <= progress / total)
		{
			res = ((*tab)->count) / (_alpha + _total);

			// cerr<<"num "<<((*tab)->count - _discount)<<" den "<<
			// 	(_alpha + _total)<<" res "<<res<<"\n";

			(*tab)->count += 1;
			progress = -1;
			break;
		}
	}

	if(progress != -1)
	{
		res = (_alpha + _discountFactor) / (_alpha + _total);
		Prob term = _prior->update(sample, temp);
		// cerr<<"Sample from prior: "<<term<<"\n";
		res *= term;

		if(_initial[sample] < _discount)
		{
			_initial[sample] += 1;
			_discountFactor += 1; //xxx
		}
		else
		{
			//initial pool reached discount
			Table* nt = new Table(sample); //auto-inits count=1
			tables.push_back(nt);
			_nTablesFor[sample] += 1;
			_nTables += 1;
 
			assert(_tables[sample].size() == _nTablesFor[sample]);
		}
	}

	_total += 1;
	_counts[sample] += 1; //items with label X

	// print(cerr);
	// cerr<<"Discount factor "<<_discountFactor<<"\n";

	return res;
}

void LCRP::addPriorCount(int sample, int count)
{
	_total += count;
	_counts[sample] += count;
	_fakeCounts += count;

	assert(count <= _discount);
	_initial[sample] += count;
	_discountFactor += count;//xxx
}

Prob LCRP::remove(int sample, double temp)
{
	// cerr<<"Removing at "<<sample<<"\n";
	// print(cerr);
	// cerr<<"Discount factor "<<_discountFactor<<"\n";
	// cerr<<"Total "<<_total<<"\n";
	int count = _counts[sample];
	assert(count > 0);

	_counts[sample] -= 1;
	_total -= 1;

	Prob res = 0;

	TableLst& tables = _tables[sample];
	if(tables.empty())
	{
		if(ifExists(_initial, sample, 0) == 0)
		{
			//rolling back an impossibility
			return 0;
		}

		_initial[sample] -= 1;
		_discountFactor -= 1;//xxx
		res = (_alpha + _discountFactor) / (_alpha + _total);
		assert(isProb(res) && res > 0);
		res *= _prior->remove(sample, temp);
		assert(isProb(res) && res > 0);
	}
	else
	{
		Prob total = 0;
		Prob priorProb = _prior->prob(sample);
		assert(priorProb > 0);

		foreach(TableLst, tab, tables)
		{
			if((*tab)->count == 1)
			{
				Prob term = (_alpha + _discountFactor) * priorProb;
				term = powl(term, temp);
				total += term;
			}
			else
			{
				Prob term = (*tab)->count - 1;
				term = powl(term, temp);
				total += term;
			}
		}

		Prob sm = gsl_rng_uniform(RNG);

		//prev table-- which one? (no prior sampling here)
		Prob progress = 0;
		foreach(TableLst, tab, tables)
		{
			Prob term = 0;
			if((*tab)->count == 1)
			{
				term = (_alpha + _discountFactor) * priorProb;
				term = powl(term, temp);
				progress += term;
			}
			else
			{
				term = (*tab)->count - 1;
				term = powl(term, temp);
				progress += term;
			}

			// cerr<<term<<" "<<progress<<" "<<(progress/total)<<"\n";
			assert(term > 0);

			if(sm <= progress / total)
			{
				// cerr<<"Found table "<<**tab<<"\n";
				(*tab)->count -= 1;
				if((*tab)->count == 0)
				{
					// cerr<<"destroying\n";
					delete *tab;
					tables.erase(tab);
					_nTablesFor[sample] -= 1;
					_nTables -= 1;

					res = (_alpha + _discountFactor) / (_alpha + _total);
					assert(isProb(res) && res > 0);
					res *= _prior->remove(sample, temp);
					assert(isProb(res) && res > 0);
				}
				else
				{
					res = ((*tab)->count) / (_alpha + _total);
					assert(isProb(res) && res > 0);
				}

				assert(_tables[sample].size() == _nTablesFor[sample]);

				progress = -1;
				break;
			}
		}

		if(progress != -1)
		{
			cerr<<"Removal: failed to find a table!\n";
			assert(0);
		}
	}

	if(_counts[sample] == 0)
	{
		_counts.erase(sample);
		_tables.erase(sample);
		_nTablesFor.erase(sample);
		assert(_initial[sample] == 0);
		_initial.erase(sample);
	}

	// print(cerr);

	assert(isProb(res) && res > 0);
	return res;
}

Prob LCRP::prob(int sample)
{
	Prob denom = _total + _alpha;

	Prob fromPrior =  _prior->prob(sample) * 
		(_alpha + _discountFactor) / denom;
	assertProb(fromPrior);
	if(fromPrior == 0)
	{
		return 0;
	}

	int count = ifExists(_counts, sample, 0);
	double disc = (count > _discount) ? _discount : count;

	Prob fromCounts = (count - disc) / (_total + _alpha);
	assertProb(fromCounts);

//	cerr<<"count for this item "<<count<<" total "<<_total<<"\n";

	Prob res = fromCounts + fromPrior;
	assert(res > 0);

	return res;
}

LogProb LCRP::ll(bool verbose)
{
	//configuration for 0 or 1 is determined
	//this special-casing is needed b/c floating-point prec. produces
	//invalid results
	if(_total <= 1 || _nTables == 0)
	{
		return 0;
	}

	LogProb res = 0;
	Prob currDiscFactor = 0;
	bool overOne = false;

	// cerr<<"---- "<<_nTables<<"\n";

	foreach(intIntMap, init, _initial)
	{
		assert(init->second > 0);

		res += gsl_sf_lngamma(_alpha + currDiscFactor + init->second);
		res -= gsl_sf_lngamma(_alpha + currDiscFactor);
		currDiscFactor += init->second;//xxx

		// cerr<<"\t*"<<init->second<<" "<<res<<"\n";
	}

	foreach(intToTableLst, tabLst, _tables)
	{
		foreach(TableLst, tab, tabLst->second)
		{
			if((*tab)->count > 1)
			{
				overOne = true;
			}

			//numerator terms 1..n associated with seating at table
			res += gsl_sf_lngamma((*tab)->count);

			//term for table creation (probably more succinct way to do this)
			res += logl(_alpha + currDiscFactor);
			// cerr<<"\t"<<(*tab)->count<<" "<<res<<"\n";
		}
	}

	if(!overOne)
	{
		return 0;
	}

	//denom terms 1..N associated with all customers
	res -= gsl_sf_lngamma(_alpha + _total);
	res += gsl_sf_lngamma(_alpha);
	
	// cerr<<"------"<<_total<<" "<<res<<"\n";

	//remember to compute prior terms associated with labels somewhere else!
	assertLogProb(res);

	return res;
}

void LCRP::print(ostream& os)
{
	foreach(intToTableLst, tabLst, _tables)
	{
		os<<":: "<<tabLst->first<<" ;"<<_initial[tabLst->first]<<
			" x"<<tabLst->second.size()<<" :: ";
		foreach(TableLst, tab, tabLst->second)
		{
			os<<**tab<<" ";
		}
		os<<"\n";
	}
	os<<"Alpha: "<<_alpha<<" discount "<<_discount<<"\n";
	os<<"Total: "<<_total<<" Config LL: "<<ll()<<"\n";
}
