/* Chaney.cc
** Markov Chains
*/

#include <iostream>
#include <fstream>
#include <string.h>
#include <string>
#include <set>
#include <map>
#include <pair.h>
#include <time.h>

using namespace std;

typedef pair <int, int> Target;

class Chain {
public:
    void Train(istream&);
    void Zap();
    string Nextword(string word);
    void Print(ostream&);

private:
    map<string, map<string, Target> > m_chains;
    map<string, bool> m_eos;
};

void Chain::Train(istream& cin)
{
    bool done = false;
    string cur, prev;

    cur = "";
    while (!done)
    {
	prev = cur;
	cin >> cur;
	if (cin.eof())
	{
	    done = true;
	    cur = "";
	}

	if ((prev != "") || (prev != cur))
	    m_chains[prev][cur].first++;
	if (m_eos[prev])
	    m_chains[""][cur].first++;

	if (strchr(".!?", cur[cur.length() - 1]))
	{
	    /*prev = cur;
	    cur = "";
	    m_chains[prev][cur].first++;*/
	    m_eos[cur] = true;
	}
    }
}

void Chain::Zap()
{
    map<string, map<string, Target> >::iterator iter1;
    map<string, Target>::iterator iter2;

    for (iter1 = m_chains.begin(); iter1 != m_chains.end(); iter1++)
	for (iter2 = (*iter1).second.begin();
	     iter2 != (*iter1).second.end(); iter2++)
	    (*iter2).second.second = (*iter2).second.first;
}

void Chain::Print(ostream& cout)
{
    map<string, map<string, Target> >::iterator iter1;
    map<string, Target>::iterator iter2;

    for (iter1 = m_chains.begin(); iter1 != m_chains.end(); iter1++)
    {
	cout << "Seed: " << (*iter1).first << endl;
	for (iter2 = (*iter1).second.begin();
	     iter2 != (*iter1).second.end(); iter2++)
	    cout << "    " << (*iter2).first << " " <<
		(*iter2).second.first << endl;
    }
}    

string Chain::Nextword(string word)
{
    int total = 0;

    // End-of-sentence has 1 in 8 chance of ending a paragraph
    if (m_eos[word])
	if (!(rand()&7))
	    return "";

    // If this word has no followers, we must end
    if (m_chains[word].begin() == m_chains[word].end())
	return "";
	

    map<string, Target>::iterator iter;
    for (iter = m_chains[word].begin(); iter != m_chains[word].end();
	 iter++)
	total += (*iter).second.second;

    total = rand()%(total + 1);
    iter = m_chains[word].begin();
    while (total > 0)
    {
	total -= (*iter).second.second;
	iter++;
	if (iter == m_chains[word].end())
	{
	    iter = m_chains[word].begin();
	    total--;
	}
    }

    return (*iter).first;
}

int main(int argc, char *argv[])
{
    Chain markov;
    string temp;
    bool done = false;
    int count = 1;

    srand(time(NULL));

    // Do the parameters
    for (int i = 1; i < argc; i++)
    {
	if (!strcmp(argv[i], "-c"))
	{
	    count = atoi(argv[i+1]);
	    i++;
	}
	else if (!strcmp(argv[i], "-h"))
	{
	    cout << "Usage: " << argv[0] << " [-c count] [file(s)]" << endl;
	    cout << "(A count of 0 will print out the database)" << endl;
	    return 0;
	} else
	{
	    done = true;

	    ifstream input(argv[i]);
	    markov.Train(input);
	    input.close();
	}
    }

    if (!done)
	markov.Train(cin);

    if (!count)
	markov.Print(cout);

    markov.Zap();
    for (int i = 0; i < count; i++)
    {
	bool link = false;
	temp = "";
	done = false;
	while (!done)
	{
	    temp = markov.Nextword(temp);
	    if (temp != "")
	    {
		if (!link && (((rand()>>3)&7) == 1))
		{
		    link = true;
		    //cout << "[";
		}

		cout << temp;

		if (link && (((rand()>>3)&1) == 1))
		{
		    link = false;
		    //cout << "]";
		}
	    } else
	    {
		done = true;
		if (link)
		{
		    link = false;
		    //cout << "]";
		}
		cout << endl << endl;
	    }

	    if (!done)
		cout << " ";
	}
    }

    return 0;
}