#include "MarkovSampler.h"

MarkovSampler::MarkovSampler(Analysis& ana, CRPLM& model, FST& misspeller,
							 LogProbArcSelector<FSTArc> selector,
							 int verbose):
	_ana(ana),
	_model(model),
	_misspeller(misspeller),
	_nSamples(10),
	_selector(selector),
	_verbose(verbose),
	_cutoff(0)
{
	_cutoff = getCutoff();
}

MarkovSampler::~MarkovSampler()
{
}

int MarkovSampler::sampleFrom(int jj, int kk)
{
	if(jj == kk)
	{
		assert(jj == 0);
		return 0;
	}

	LogProb max = 0;
	bool begin = true;

	intIntProbMap& itemsByPrev = _samples[jj][kk];

	// cerr<<"sampling from\n";
	// writeHashMap(items, cerr);

	foreach(intIntProbMap, byPrev, itemsByPrev)
	{
		foreach(intProbMap, item, byPrev->second)
		{
			LogProb term = item->second;
			if(begin || term > max)
			{
				max = term;
			}
			begin = false;
		}
	}
	assert(!begin);

	Prob total = 0;
	foreach(intIntProbMap, byPrev, itemsByPrev)
	{
		foreach(intProbMap, item, byPrev->second)
		{
			LogProb term = item->second;
			Prob px = expl(term - max);
			// cerr<<"term "<<term<<" normed to "<<px<<"\n";
			total += px;
		}
	}

	Prob sm = gsl_rng_uniform(RNG);
	int chosen = -1;
	Prob progress = 0;
	foreach(intIntProbMap, byPrev, itemsByPrev)
	{
		foreach(intProbMap, item, byPrev->second)
		{
			LogProb term = item->second;
			Prob px = expl(term - max);
			progress += px;
			if(progress / total > sm)
			{
				chosen = item->first;
				break;
			}
		}
		if(chosen != -1)
		{
			break;
		}
	}
	assert(chosen != -1);

	if(_verbose & TRANSITIONS)
	{
		int count = ifExists(_model.lexicon(), chosen, 0);
		cerr<<"sample on "<<jj<<" "<<kk<<" is "<<
			_model.repr(chosen)<<" ("<<count<<")\n";
	}

	return chosen;

}

void MarkovSampler::sampleIntended(int jj, int kk, 
								   SymbolString* targ,
								   intProbMap& samples)
{
	bool verbose = _verbose & SAMPLES;

	if(verbose)
	{
		cerr<<*targ<<" ---- target\n";
	}

	FSTKit fsts(targ->alphabet());
	// FST targAcc;
	// fsts.stringToFst(*targ, targAcc);

	for(int ii = 0; ii < _nSamples; ++ii)
	{
		FST* misspellVocab; //was vocab, and compositions in place
		if(_model._ngrams == 2)
		{
			int prevToken = sampleFrom(jj, kk);

			intLst ctxt;
			ctxt.push_back(prevToken);
			CRP* proc = _model.getProcessOrEmpty(ctxt);

			if(_verbose & FSTS)
			{
				proc->fst()->Write("vocab.fst");
			}

			misspellVocab = proc->misspellFst(_misspeller);
		}
		else
		{
			if(_verbose & FSTS)
			{
				_model._root->fst()->Write("vocab.fst");
			}

			misspellVocab = _model._root->misspellFst(_misspeller);
		}
		assert(misspellVocab != NULL);

		// ComposeFst<FSTArc> misspellVocab(*vocab, _misspeller);
		FST unigramPost;
//		Compose(*misspellVocab, targAcc, &unigramPost);
		// fsts.composeWithString(*misspellVocab, *targ, unigramPost);

		if(_verbose & FSTS)
		{
			misspellVocab->Write("misspellvoc.fst");
			// unigramPost.Write("unigrams.fst");
		}

		FST sample;
		LogProb pr = fsts.sampleComposedWithString(*misspellVocab, sample,
												   *targ, _verbose & FSTS);

		// Project(&sample, PROJECT_INPUT);

		// if(_verbose & FSTS)
		// {
		// 	sample.Write("sample.fst");
		// }

		// FST unipath;
		// Compose(sample, unigramPost, &unipath);
		// LogProb pr = fsts.fstProb(unipath);
		pr = expl(pr);

		SymbolString* str = fsts.fstToString(sample);
		int strSym = _model.intern(*str);

		if(verbose)
		{
			int count = ifExists(_model.lexicon(), strSym, 0);
			cerr<<"sampled "<<*str<<" (lex "<<count<<") pr "<<pr<<"\n";
		}

		samples[strSym] += pr;
		delete str;
	}
}

LogProb MarkovSampler::proposalTransProb(int prevToken, int nextToken,
										 SymbolString& targ)
{
	if(_verbose & PROPOSAL)
	{
		cerr<<"Proposal: compute transition from "<<_model.repr(prevToken)<<
			" to "<<_model.repr(nextToken)<<" spanning "<<targ<<"\n";
	}

	FST* misspellVocab;
	if(_model._ngrams == 2)
	{
		intLst ctxt;
		ctxt.push_back(prevToken);
		CRP* proc = _model.getProcessOrEmpty(ctxt);

		// if(_verbose & PROPOSAL)
		// {
		// 	proc->fst()->Write("vocab.fst");
		// }

		misspellVocab = proc->misspellFst(_misspeller);
	}
	else
	{
		// if(_verbose & PROPOSAL)
		// {
		// 	_model._root->fst()->Write("vocab.fst");
		// }

		misspellVocab = _model._root->misspellFst(_misspeller);
	}
	assert(misspellVocab != NULL);

	SymbolString& next = _model.repr(nextToken);
	FST underAcc;
	FSTKit fsts(targ.alphabet());
	fsts.stringToFst(next, underAcc);

	FST unigramPost;
	Compose(underAcc, *misspellVocab, &unigramPost);

	// if(_verbose & PROPOSAL)
	// {
	// 	misspellVocab->Write("misspellvoc.fst");
	// 	unigramPost.Write("unigrams.fst");
	// }

	LogProb samples[_nSamples];
	for(int ii = 0; ii < _nSamples; ++ii)
	{
		FST path;
		LogProb pr = fsts.sampleComposedWithString(unigramPost, path,
												   targ, 0);
		samples[ii] = pr;
	}

	return avgOfLogs(samples, _nSamples);
}

SymbolString* MarkovSampler::getString(int jj, int kk)
{
	SymbolString* str = new SymbolString(_ana.surface.alphabet());
	str->appendRange(_ana.surface, jj, kk);
	return str;
}

LogProb MarkovSampler::endTransition(int kk, int tt)
{
	LogProb samples[_nSamples];

	for(int ii = 0; ii < _nSamples; ++ii)
	{
		int prevToken = sampleFrom(kk, tt);
		intLst context;
		context.push_back(prevToken);
		CRP* proc = _model.getProcess(context);
		Prob res = proc->prob(0);

		if(_verbose & TRANSITIONS)
		{
			cerr<<"checking transition between "<<_model.repr(prevToken)<<
				" and [end] ("<<
				_model.lexicon()[prevToken]<<")\n";
			cerr<<"="<<res<<" log "<<logl(res)<<"\n";
		}

		samples[ii] = res;
	}

	return avgOfLogs(samples, _nSamples);
}

LogProb MarkovSampler::transition(int jj, int kk, int tt)
{
	SymbolString* curr = getString(kk, tt);

	if(_verbose & TRANSITIONS)
	{
		SymbolString* prev = getString(jj, kk);
		cerr<<"checking transition between "<<*prev<<
			" and "<<*curr<<" ("<<
			_model.lexicon()[_model.intern(*prev)]<<" "<<
			_model.lexicon()[_model.intern(*curr)]<<
			")\n";
		delete prev;
	}

	intProbMap currInts;
	sampleIntended(jj, kk, curr, currInts);
	assert(curr->size() > 0);

	Prob pAvg = 0;

	LogProb prevPr = _table[jj][kk];

	foreach(intProbMap, sample, currInts)
	{
		if(_verbose & TRANSITIONS)
		{
			cerr<<"Sample: "<<
				_model.repr(sample->first)<<" : "<<sample->second<<"\n";
		}

		pAvg += sample->second;

		LogProb term = logl(sample->second) + prevPr;

		if(inMap(_samples[kk][tt][jj], sample->first))
		{
			LogProb sum = logplus(_samples[kk][tt][jj][sample->first], term);
			_samples[kk][tt][jj][sample->first] = sum;
		}
		else
		{
			_samples[kk][tt][jj][sample->first] = term;
		}
	}

	delete curr;

	pAvg /= _nSamples;
	assertProb(pAvg);

	if(_verbose & TRANSITIONS)
	{
		cerr<<"avg prob "<<pAvg<<" log "<<logl(pAvg)<<"\n";
	}

	return logl(pAvg);
}

int MarkovSampler::sampleNext(ProbMat& table, int last)
{
	// cerr<<"Sampling backwards from "<<last<<"\n";
	LogProb maxT = 0;

	int beginT = max(0, last - _cutoff);
	for(int tt = beginT; tt < last; ++tt)
	{
		LogProb term = table[last][tt];
		if(tt == beginT || term > maxT)
		{
			maxT = term;
		}
	}

	probToInt probs;
	Prob total = 0;
	for(int tt = beginT; tt < last; ++tt)
	{
		LogProb term = table[last][tt];
		Prob px = expl(term - maxT);
		// cerr<<"term "<<term<<" normed to "<<px<<"\n";
		if(px > 0)
		{
			probs[px + total] = tt;
			total += px;
		}
	}

	// writeMap(probs, cerr);

	Prob sm = gsl_rng_uniform(RNG);
	int chosen = probs.lower_bound(sm)->second;
	return chosen;
}

int MarkovSampler::getCutoff()
{
	int cutoff = 1;
	intIntMap& lexicon = _model.lexicon();
	int longestWord = 0;
	foreach(intIntMap, word, lexicon)
	{
		int wordLen = _model.repr(word->first).size();
		if(wordLen + 1 > cutoff)
		{
			cutoff = wordLen + 1;
			longestWord = word->first;
		}
	}

	if(cutoff > MAX_WORD_LEN)
	{
		cutoff = MAX_WORD_LEN;
	}

	if(_verbose & CUTOFF)
	{
		cerr<<"Cutoff param: "<<cutoff<<" (longest word "<<
			_model.repr(longestWord)<<": "<<_model.lexicon()[longestWord]
			<<")"<<
			" utt length "<<_ana.nSyms()<<"\n";
	}

	return cutoff;
}

void MarkovSampler::proposeBreaks(ints& breakpoints)
{
	assert(_model._ngrams <= 2);

	int nn = _ana.nSyms();

	_table.resize(nn + 1);
	for(int tt = 0; tt <= nn; ++tt)
	{
		_table[tt].resize(nn + 1);
	}
	_table[0][0] = 0.0;

	for(int tt = 1; tt <= nn; ++tt)
	{
		if(_verbose & INDICES)
		{
			cerr<<"t="<<tt<<"...\n";
		}

		//word boundary of the word ending at t
		int beginK = max(0, tt - _cutoff);
		for(int kk = beginK; kk < tt; ++kk)
		{
			if(_verbose & INDICES)
			{
				cerr<<"\tkk="<<kk<<"\n";
			}

			int beginJ = max(0, kk - _cutoff);
			for(int jj = beginJ; jj < kk || jj == 0; ++jj)
			{
				if(_verbose & INDICES)
				{
					cerr<<"\t\tjj="<<jj<<"\n";
				}

				LogProb probBefore = _table[kk][jj];
				LogProb ptrans = transition(jj, kk, tt);

				LogProb term = probBefore + ptrans;
				assertLogProb(term);
				LogProb sum = term;
				if(jj > beginJ)
				{
					sum = logplus(_table[tt][kk], term);
				}
				assertLogProb(sum);

				_table[tt][kk] = sum;
			}
		}

		// for(int ii = 0; ii <= tt; ++ii)
		// {
		// 	cerr<<ii<<": ";
		// 	for(int jj = 0; jj < ii; ++jj)
		// 	{
		// 		cerr<<"("<<jj<<" "<<_table[ii][jj]<<") ";
		// 	}
		// 	cerr<<"\n";
		// }
	}

	//sample the end token
	if(_model._ngrams == 2)
	{
		int beginK = max(0, nn - _cutoff);
		for(int kk = beginK; kk < nn; ++kk)
		{
			LogProb probBefore = _table[nn][kk];
			LogProb ptrans = endTransition(kk, nn);
			LogProb term = probBefore + ptrans;
			assertLogProb(term);

			_table[nn][kk] = term;
		}
	}

	int last = nn;
	while(last != 0)
	{
		int next = sampleNext(_table, last);
		breakpoints.push_back(last);
		last = next;
	}

	if(_verbose & INDICES)
	{
		cerr<<"breakpoints: ";
		writeVector(breakpoints, cerr);
	}
}

LogProb MarkovSampler::pPropose(ints& breakpoints, ints& lex)
{
	LogProb propose = 0;
	int last = _ana.nSyms();
	int next;
	ints::reverse_iterator prevWord = lex.rbegin();
	int currWord;

	if(_verbose & PROPOSAL)
	{
		cerr<<"\n";
		cerr<<"Probability of proposing ";
		writeVector(breakpoints, cerr);
	}

	if(_model._ngrams == 2)
	{
		intLst context;
		context.push_back(*prevWord);
		CRP* proc = _model.getProcess(context);
		Prob end = proc->prob(0);
		LogProb pProp = logl(end);

		if(_verbose & PROPOSAL)
		{
			cerr<<"End of string marker = "<<pProp<<"\n";
		}

		propose += pProp;
	}

	int sz = breakpoints.size();
	//breakpoints[0] is the last symbol addr
	for(int ptr = 1; ptr < sz; ++ptr)
	{
		currWord = *prevWord;
		prevWord++;
		next = breakpoints[ptr];
		// int next2 = (ptr + 1 < sz) ? breakpoints[ptr + 1] : 0;
		LogProb pProp;
		// //XXX this won't work for _nSamples > 1
		// assert(_nSamples == 1);
		// if(inMap(_samples[next][last][next2], currWord))
		// {
		// 	pProp = _samples[next][last][next2][currWord];
		// }
		// else
		{
			SymbolString* targ = getString(next, last);
			pProp = proposalTransProb(*prevWord, currWord, *targ);
			delete targ;
		}

		if(_verbose & PROPOSAL)
		{
			cerr<<"Interval "<<next<<" "<<last<<"\n";
			cerr<<"Prob of interval "<<*getString(next, last)<<" "<<
				_model.repr(currWord)<<" = ";
			cerr<<pProp<<"\n";
		}

		propose += pProp;
		last = next;
	}
	{
		currWord = *prevWord;
		LogProb pProp;
		// if(inMap(_samples[0][last][0], currWord))
		// {
		// 	pProp = _samples[0][last][0][currWord];
		// }
		// else
		{
			SymbolString* targ = getString(0, last);
			pProp = proposalTransProb(0, currWord, *targ);
			delete targ;
		}
		if(_verbose & PROPOSAL)
		{
			cerr<<"Interval "<<0<<" "<<last<<"\n";
			cerr<<"Prob of interval "<<*getString(0, last)<<"\n";
			cerr<<pProp<<"\n";
		}
		propose += pProp;
	}

	return propose;
}

void MarkovSampler::sampleIntendedForms(ints& breakpoints, ints& newLex)
{
	int prev = 0;
	int prev2 = 0;
	int sz = breakpoints.size();
	for(int bp = sz - 1; bp >= 0; --bp)
	{
		int point = breakpoints[bp];

		int chosen = -1;
		{
			bool begin = true;
			LogProb mx = 0;
			foreach(intProbMap, item, _samples[prev][point][prev2])
			{
				LogProb term = item->second;
				if(begin || term > mx)
				{
					mx = term;
				}
				begin = false;
			}
			assert(!begin);

			Prob total = 0;
			foreach(intProbMap, item, _samples[prev][point][prev2])
			{
				total += expl(item->second - mx);
			}

			Prob sm = gsl_rng_uniform(RNG);			
			Prob progress = 0;
			foreach(intProbMap, item, _samples[prev][point][prev2])
			{
				LogProb term = item->second;
				Prob px = expl(term - mx);
				// cerr<<"term "<<term<<" normed to "<<px<<"\n";
				progress += px;
				if(progress / total >= sm)
				{
					chosen = item->first;
					break;
				}
			}
			assert(chosen != -1);
		}

		// if(_verbose & PROPOSAL)
		// {
		// 	cerr<<"Sampling from ";
		// 	writeHashMap(_samples[prev][point][prev2], cerr);
		// 	cerr<<"Chose "<<chosen<<" "<<_model.repr(chosen)<<"\n";
		// }

		newLex.push_back(chosen);

		prev2 = prev;
		prev = point;
	}
}

LogProb MarkovSampler::transducerLexProb(ints& intended, ints& breakpoints)
{
	LogProb res = 0;

	int sz = intended.size();
	int prev = 0;
	for(int ii = 0; ii < sz; ++ii)
	{
		int point = breakpoints[sz - (ii + 1)];
		SymbolString* surf = getString(prev, point);

		LogProb term = transducerProb(intended[ii], _model.intern(*surf));
		if(_verbose & PROPOSAL)
		{
			cerr<<"transducer prob between "<<_model.repr(intended[ii])<<
				" and "<<*surf<<" = "<<term<<"\n";
		}
		res += term;
		delete surf;

		prev = point;
	}

	assertLogProb(res);
	return res;
}

LogProb MarkovSampler::transducerProb(int inSym, int outSym)
{
	SymbolString& ins = _model.repr(inSym);
	SymbolString& outs = _model.repr(outSym);

	FSTKit fsts(ins.alphabet());

	FST inF;
	fsts.stringToFst(ins, inF);

	FST outF;
	fsts.stringToFst(outs, outF);

	ComposeFst<FSTArc> misspellIn(inF, _misspeller);
	FST through;
	Compose(misspellIn, outF, &through);

	if(through.NumStates() == 0)
	{
		cerr<<"Failed to compute prob between "<<ins<<" and "<<outs<<"\n";
		inF.Write("in.fst");
		outF.Write("out.fst");
		_misspeller.Write("miss.fst");
		assert(0);
	}

	return fsts.fstProb(through);
}

void MarkovSampler::sample(double temp)
{
	ints orig(_ana.nSegments());
	ints origLex;

	{
		int ctr = _ana.nSegments() - 1;

		Segment* curr = _ana.segAt(0);

		while(!curr->isStartOrEnd())
		{
			if(curr->begin != 0)
			{
				orig[ctr--] = curr->begin;
			}
			origLex.push_back(curr->strSym);

			curr = _ana.nextSegment(*curr);
		}

		assert(ctr == 0);
		orig[ctr] = _ana.nSyms();
	}

	if(_verbose & MHR)
	{
		cerr<<"Original breakpoints\n";
		writeVector(orig, cerr);
	}

	// cerr<<"removing old word string\n";
	LogProb postOld = _model.removeSeq(origLex, temp, _verbose & PROPOSAL);
	postOld += transducerLexProb(origLex, orig);

	ints breakpoints;
	proposeBreaks(breakpoints);

	ints newLex;
	sampleIntendedForms(breakpoints, newLex);

	LogProb propOld = pPropose(orig, origLex);
	LogProb pProp = pPropose(breakpoints, newLex);

	LogProb postNew = _model.updateSeq(newLex, temp, _verbose & PROPOSAL);
	postNew += transducerLexProb(newLex, breakpoints);

	LogProb mhr = (propOld - pProp) + (postNew - postOld);
	Prob mhrE = expl(mhr);
	Prob sm = gsl_rng_uniform(RNG);

	if(_verbose & MHR)
	{
		cerr<<"Propose (old "<<propOld<<" "<<pProp<<") "<<
			(propOld - pProp)<<"\n";
		cerr<<"Post (old "<<postOld<<" "<<postNew<<") "<<
			(postOld - postNew)<<"\n";
		cerr<<"MHR "<<mhr<<" "<<mhrE<<"\n";
	}

	if(sm > mhrE)
	{
		if(_verbose & MHR)
		{
			cerr<<"Reject\n";
		}

		_model.removeSeq(newLex, temp, _verbose & PROPOSAL);
		_model.updateSeq(origLex, temp, _verbose & PROPOSAL);
	}
	else
	{
		//update ana
		if(_verbose & MHR)
		{
			cerr<<"Accept\n";
		}

		//skip bpt 0, the last position
		int sz1 = orig.size();
		for(int ii = 1; ii < sz1; ++ii)
		{
			_ana.merge(orig[ii]);
		}

		int sz2 = breakpoints.size();
		for(int ii = 1; ii < sz2; ++ii)
		{
			_ana.split(breakpoints[ii]);
		}

		int sz3 = newLex.size();
		Segment* seg = _ana.segAt(0);
		for(int ii = 0; ii < sz3; ++ii, seg = _ana.nextSegment(*seg))
		{
			int smp = newLex[ii];
			seg->strSym = smp;
			seg->str = _model.repr(smp);
		}
	}

	if(_verbose & MHR)
	{
		cerr<<_ana<<"\n";
	}
}
