/*************************************************************************/
/*                                                                       */
/*                  Language Technologies Institute                      */
/*                     Carnegie Mellon University                        */
/*                      Copyright (c) 1999-2003                          */
/*                        All Rights Reserved.                           */
/*                                                                       */
/*  Permission is hereby granted, free of charge, to use and distribute  */
/*  this software and its documentation without restriction, including   */
/*  without limitation the rights to use, copy, modify, merge, publish,  */
/*  distribute, sublicense, and/or sell copies of this work, and to      */
/*  permit persons to whom this work is furnished to do so, subject to   */
/*  the following conditions:                                            */
/*   1. The code must retain the above copyright notice, this list of    */
/*      conditions and the following disclaimer.                         */
/*   2. Any modifications must be clearly marked as such.                */
/*   3. Original authors' names are not deleted.                         */
/*   4. The authors' names are not used to endorse or promote products   */
/*      derived from this software without specific prior written        */
/*      permission.                                                      */
/*                                                                       */
/*  CARNEGIE MELLON UNIVERSITY AND THE CONTRIBUTORS TO THIS WORK         */
/*  DISCLAIM ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING      */
/*  ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS, IN NO EVENT   */
/*  SHALL CARNEGIE MELLON UNIVERSITY NOR THE CONTRIBUTORS BE LIABLE      */
/*  FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES    */
/*  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN   */
/*  AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION,          */
/*  ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF       */
/*  THIS SOFTWARE.                                                       */
/*                                                                       */
/*************************************************************************/
/*                     Author :  Alan W Black                            */
/*                     Date   :  October 1999                            */
/*-----------------------------------------------------------------------*/
/*                                                                       */
/* Training method to split states of existing WFST based on data to     */
/* optimize entropy                                                      */
/*                                                                       */
/* Confusing as this has nothing to do with the modelling                */
/* technique known as "maximum entropy"                                  */
/*                                                                       */
/*=======================================================================*/
#include <iostream>
#include <cstdlib>
#include "EST_WFST.h"
#include "wfst_aux.h"
#include "EST_Token.h"
#include "EST_simplestats.h"

VAL_REGISTER_TYPE_NODEL(trans,EST_WFST_Transition)
SIOD_REGISTER_CLASS(trans,EST_WFST_Transition)
VAL_REGISTER_CLASS(pdf,EST_DiscreteProbDistribution)
SIOD_REGISTER_CLASS(pdf,EST_DiscreteProbDistribution)

static LISP *find_state_usage(EST_WFST &wfst, LISP data);
static double entropy(const EST_WFST_State *s);
static LISP *find_state_entropies(const EST_WFST &wfst, LISP *data);
EST_WFST_Transition *find_best_trans_split(EST_WFST &wfst,
					   int split_state,
					   LISP *data);
static LISP find_best_split(EST_WFST &wfst,	
			    int split_state_name,
			    LISP *data);
static double find_score_if_split(EST_WFST &wfst,
				  int fromstate,
				  EST_WFST_Transition *trans,
				  LISP *data);
static LISP find_split_pdfs(EST_WFST &wfst,
			    int split_state_name,
			    LISP *data,
			    EST_DiscreteProbDistribution &pdf_all);
static double score_pdf_combine(EST_DiscreteProbDistribution &a,
				EST_DiscreteProbDistribution &b,
				EST_DiscreteProbDistribution &all);
#if 0
static void split_state(EST_WFST &wfst, EST_WFST_Transition *trans);
#endif
static void split_state(EST_WFST &wfst, LISP trans_list, int ostate);

LISP load_string_data(EST_WFST &wfst,EST_String &filename)
{
    // Load in sentences into data table, assume sentence per line
    EST_TokenStream ts;
    LISP ss = NIL;
    EST_String t;
    int id;
    int i,j;
    
    if (ts.open(filename) == -1)
	EST_error("wfst_train: failed to read data from \"%s\"",
			  (const char *)filename);

    i = 0;
    j = 0;
    while (!ts.eof())
    {
	LISP s = NIL;
	do
	{
	    t = (EST_String)ts.get();
	    id = wfst.in_symbol(t);
	    if (id == -1)
	    {
		cerr << "wfst_train: data contains unknown symbol \"" <<
		    t << "\"" << endl;
	    }
	    s = cons(flocons(id),s);
	    j++;
	}
	while (!ts.eoln() && !ts.eof());
	i++;
	ss = cons(reverse(s),ss);
    }

    printf("wfst_train: loaded %d lines of %d tokens\n",
	   i,j);

    return reverse(ss);
}

static LISP *find_state_usage(EST_WFST &wfst, LISP data)
{
    // Builds list of states, and which data points the represent
    LISP *state_data = new LISP[wfst.num_states()];
    static LISP ddd = NIL;
    int s,i,id;
    LISP d,w;
    EST_WFST_Transition *trans;
//    EST_Litem *tp;

    if (ddd == NIL)
	gc_protect(&ddd);

    ddd = NIL;

    wfst.start_cumulate();   // zero existing weights

    for (i=0; i < wfst.num_states(); i++)
    {
	state_data[i] = NIL;
	ddd = cons(state_data[i],ddd);
//	// smoothing
//	for (tp=wfst.state(i)->transitions.head(); tp != 0; tp = tp->next())
//	    wfst.state(i)->transitions(tp)->set_weight(1);
    }

    for (i=0,d=data; d; d=cdr(d),i++)
    {
	s = wfst.start_state();
	for (w=car(d); w; w=cdr(w))
	{
	    state_data[s] = cons(w,state_data[s]);
	    id = get_c_int(car(w));
	    trans = wfst.find_transition(s,id,id);
	    if (!trans)
	    {
		printf("sentence %d not in language, skipping\n",i);
		continue;
	    }
	    else
	    {
		trans->set_weight(trans->weight()+1);
		s = trans->state();
	    }
	}
    }
	
    wfst.stop_cumulate();
    return state_data;
}

static double entropy(const EST_WFST_State *s)
{
    double sentropy,w;
    EST_Litem *tp;
    for (sentropy=0,tp=s->transitions.head(); tp != 0; tp = tp->next())
    {
	w = s->transitions(tp)->weight();  /* the probability */
	if (w > 0)
	    sentropy += w * log(w);
    }
    return -1 * sentropy;
}

void wfst_train(EST_WFST &wfst, LISP data)
{
    LISP *state_data;
    LISP *state_entropies;
    LISP best_trans_list = NIL;
    int c=0,i, max_entropy_state;
    gc_protect(&data);

    while (1)
    {
	// Build table of state to points in data, and cumulate transitions
	state_data = find_state_usage(wfst,data);

	/* find entropy for each state (sorted) */
	state_entropies = find_state_entropies(wfst,state_data);

	max_entropy_state = -1;
	for (i=0; i < wfst.num_states(); i++)
	{
//	    double me = (double)get_c_float(car(state_entropies[i]));
	    max_entropy_state = get_c_int(cdr(state_entropies[i]));
//	    printf("trying %d %g\n",max_entropy_state,me);

//	    best_trans = find_best_trans_split(wfst,max_entropy_state,
//					       state_data);
	    best_trans_list = find_best_split(wfst,max_entropy_state,
					      state_data);
	    if (best_trans_list != NIL)
		break;
//	    else
//		printf("No best trans\n");
	}
	delete [] state_entropies;

	if (max_entropy_state == -1)
	{
	    printf("No new max_entropy state\n");
	    break;
	}
	if (best_trans_list == NIL)
	{
	    printf("No best_trans in max_entropy state\n");
	    break;
	}

        /* for each transition *entering* max_entropy_state */
        /*     find entropy if it were split          */
        /*     find best split                      */

        /* print stats */
        /* some sort of stop check */
	c++;
	printf("c is %d\n",c);
	if (c > 5000)
	{
	    printf("reached cycle end %d\n",c);
	    break;
	}
        /* split on best split                      */
        split_state(wfst, best_trans_list, max_entropy_state);
	
	if ((c % 100) == 0)
	{
	    EST_String chkpntname = "chkpnt";
	    char bbb[7];
	    sprintf(bbb,"%03d",c);
	    wfst.save(chkpntname+bbb+".wfst");
	}

	delete [] state_data;
	user_gc(NIL);
    }
}

static int me_compare_function(const void *a, const void *b)
{
    LISP la;
    LISP lb;
    la = *(LISP *)a;
    lb = *(LISP *)b;

    float fa = get_c_float(car(la));
    float fb = get_c_float(car(lb));

    if (fa < fb)
	return 1;
    else if (fa == fb)
	return 0;
    else
	return -1;
}

static LISP *find_state_entropies(const EST_WFST &wfst, LISP *data)
{
    double all_entropy = 0;
    int i;
    double sentropy;
    LISP *slist = new LISP[wfst.num_states()];
    static LISP ddd = NIL;

    if (ddd == NIL)
	gc_protect(&ddd);
    ddd = NIL;
    
    for (i=0; i < wfst.num_states(); i++)
    {
	const EST_WFST_State *s = wfst.state(i);
	sentropy = entropy(s);
//	printf("dlength is %d %d\n",i,siod_llength(data[i]));
	all_entropy += sentropy * siod_llength(data[i]);
	slist[i] = cons(flocons(sentropy),flocons(i));
	ddd = cons(slist[i],ddd);
    }
    printf("average entropy is %g\n",all_entropy/i);

    qsort(slist,wfst.num_states(),sizeof(LISP),me_compare_function);

    return slist;
}

static LISP find_best_split(EST_WFST &wfst,	
			    int split_state_name,
			    LISP *data)
{
    // Find the best partition of incoming translations that
    // minimises entropy
    EST_DiscreteProbDistribution pdf_all(&wfst.in_symbols());
    EST_DiscreteProbDistribution *a_pdf, *b_pdf;
    LISP splits,s,dd,r;
    LISP *ssplits;
    gc_protect(&splits);
    EST_String sname;
    int b,best_b;
    EST_Litem *i;
    int num_pdfs;
    double best_score, score, sfreq;

    for (dd = data[split_state_name]; dd; dd = cdr(dd))
	pdf_all.cumulate(get_c_int(car(car(dd))));
    splits = find_split_pdfs(wfst,split_state_name,data,pdf_all);
    if (siod_llength(splits) < 2)
	return NIL;
    ssplits = new LISP[siod_llength(splits)];
    for (num_pdfs=0,s=splits; s != NIL; s=cdr(s),num_pdfs++)
	ssplits[num_pdfs] = car(s);

    qsort(ssplits,num_pdfs,sizeof(LISP),me_compare_function);
    // Combine trans pdfs in pdfs until more combination doesn't improve
    while (1)
    {

	best_score = get_c_float(car(ssplits[0]));
	best_b = -1;
	a_pdf = pdf(car(cdr(cdr(ssplits[0]))));
        for (b=1; b < num_pdfs; b++)
	{
	    if (ssplits[b] == NIL)
		continue;
	    score = score_pdf_combine(*a_pdf,*pdf(car(cdr(cdr(ssplits[b])))),
				      pdf_all);
	    if (score < best_score)
	    {
		best_score = score;
		best_b = b;
	    }
	}

	// combine a and b
	if (best_b == -1)
	    break;
	else
	{
	    // combine a and b
	    // Add trans to 0
	    setcar(cdr(ssplits[0]),
		   append(car(cdr(ssplits[0])),
			  car(cdr(ssplits[best_b]))));
	    setcar(ssplits[0], flocons(best_score));
	    // Update 0's pdf with values from best_b's
	    b_pdf = pdf(car(cdr(cdr(ssplits[best_b]))));
	    for (i=b_pdf->item_start(); !b_pdf->item_end(i);
		 i = b_pdf->item_next(i))
	    {
		b_pdf->item_freq(i,sname,sfreq);
		a_pdf->cumulate(i,sfreq);
	    }
	    ssplits[best_b] = NIL;
	}

    }
    
    printf("score %g ",(double)get_c_float(car(ssplits[0])));
    for (dd=car(cdr(ssplits[0])); dd; dd=cdr(dd))
	printf("%s ",(const char *)wfst.in_symbol(trans(car(dd))->in_symbol()));
    printf("\n");
    gc_unprotect(&splits);
    r = car(cdr(ssplits[0]));
    delete [] ssplits;
    return r;
}

static double score_pdf_combine(EST_DiscreteProbDistribution &a,
				EST_DiscreteProbDistribution &b,
				EST_DiscreteProbDistribution &all)
{
    // Find score of (a+b) vs (all-(a+b))
    EST_DiscreteProbDistribution ab(a);
    EST_DiscreteProbDistribution all_but_ab(all);
    EST_Litem *i;
    EST_String sname;
    double sfreq, score;
    for (i=b.item_start(); !b.item_end(i);
	 i = b.item_next(i))
    {
	b.item_freq(i,sname,sfreq);
	ab.cumulate(i,sfreq);
    }
    
    for (i=ab.item_start(); !ab.item_end(i);
	 i = ab.item_next(i))
    {
	ab.item_freq(i,sname,sfreq);
	all_but_ab.cumulate(i,-1*sfreq);
    }
    
    score = (ab.entropy() * ab.samples()) +
	(all_but_ab.entropy() * all_but_ab.samples());

    return score;

}

static LISP find_split_pdfs(EST_WFST &wfst,
			    int split_state_name,
			    LISP *data,
			    EST_DiscreteProbDistribution &pdf_all)
{
    // Find following pdfs for each incoming transition as if they where
    // split to a new state
    int i,id, in;
    EST_Litem *tp;
    LISP pdfs = NIL,dd,ttt,p,t;
    EST_DiscreteProbDistribution empty;
    double value;

    for (i=0; i < wfst.num_states(); i++)
    {
	const EST_WFST_State *s = wfst.state(i);
	for (tp=s->transitions.head(); tp != 0; tp = tp->next())
	{
	    if ((s->transitions(tp)->state() == split_state_name)
		&& (s->transitions(tp)->weight() > 0))
	    {
		in = s->transitions(tp)->in_symbol();
		EST_DiscreteProbDistribution *pdf = 
		    new EST_DiscreteProbDistribution(&wfst.in_symbols());
		for (dd = data[i]; dd; dd = cdr(dd))
		{
		    id = get_c_int(car(car(dd)));
		    if (id == in)
		    {   // This one would go to the new state so we count it
			if (cdr(car(dd))) // not end of data string
			    pdf->cumulate(get_c_int(car(cdr(car(dd)))));
		    }
		}
		// value, list of trans, pdf
		value = score_pdf_combine(*pdf,empty,pdf_all);
		if ((value > 0) && // ignore transitions with no data
		    (pdf->samples() > 10))// and those with only a few data pnts
		{
		    t = siod(s->transitions(tp));
		    p = siod(pdf);
		    ttt = cons(flocons(value),
			       cons(cons(t,NIL),
				    cons(p,NIL)));
		    pdfs = cons(ttt,pdfs);
		}
		else
		    delete pdf;
	    }
	}
    }
    return pdfs;
}

EST_WFST_Transition *find_best_trans_split(EST_WFST &wfst,
					   int split_state_name,
					   LISP *data)
{
    EST_Litem *tp;
    EST_WFST_Transition *best_trans = 0;
    const EST_WFST_State *split_state = wfst.state(split_state_name);
    double best_score,bb;
    int i;

    best_score = entropy(split_state)*siod_llength(data[split_state_name]);
//    printf("unsplit score %g\n",best_score);

    /* For each transition going to split_state */
    for (i=1; i < wfst.num_states(); i++)
    {
	const EST_WFST_State *s = wfst.state(i);
	for (tp=s->transitions.head(); tp != 0; tp = tp->next())
	{
	    if ((wfst.state(s->transitions(tp)->state()) == split_state) &&
		(s->transitions(tp)->weight() > 0))
	    {
		bb = find_score_if_split(wfst,i,s->transitions(tp),data);
//		cout << i << " " 
//		     << wfst.in_symbol(s->transitions(tp)->in_symbol()) << " "
//		     << s->transitions(tp)->state() << " " << bb << endl;
		if (bb == -1)  /* didn't find a split */
		    continue;
		if (bb < best_score)
		{
		    best_score = bb;
		    best_trans = s->transitions(tp);
		}
	    }
	}
    }

    if (best_trans)
	cout << "best " << wfst.in_symbol(best_trans->in_symbol()) << " "
	     << best_trans->weight() << " " 
 	     << best_trans->state() << " " << best_score << endl;
    return best_trans;
}

static double find_score_if_split(EST_WFST &wfst,
                                  int fromstate,
				  EST_WFST_Transition *trans,
				  LISP *data)
{
    double ent_split;
    double ent_remain;
    double score;
    EST_DiscreteProbDistribution pdf_split(&wfst.in_symbols());
    EST_DiscreteProbDistribution pdf_remain(&wfst.in_symbols());
    int in, tostate, id;
    EST_Litem *i;
    double sfreq;
    EST_String sname;

    ent_split = ent_remain = 32*32*32*32;
    LISP dd;

//    printf("considering %d %s %g %d\n",
//	   fromstate,
//	   (const char *)wfst.in_symbol(trans->in_symbol()),
//	   trans->weight(),
//	   trans->state());

    /* find entropy of possible new state */
    /* for each data point through fromstate */
    in = trans->in_symbol();
    for (dd = data[fromstate]; dd; dd = cdr(dd))
    {
	id = get_c_int(car(car(dd)));
	if (id == in)
	{   // This one would go to the new state so we count it
	    if (cdr(car(dd))) // not end of data string
		pdf_split.cumulate(get_c_int(car(cdr(car(dd)))));
	}
    }
    if (pdf_split.samples() > 0)
	ent_split = pdf_split.entropy();
    /* find entropy of old state minus trans into it */
    tostate = trans->state();
    // Actually only need to do this once per state
    for (dd = data[tostate]; dd; dd = cdr(dd))
	pdf_remain.cumulate(get_c_int(car(car(dd))));
    // Subtract the bit thats split
    for (i=pdf_split.item_start(); !pdf_split.item_end(i);
	 i = pdf_split.item_next(i))
    {
	pdf_split.item_freq(i,sname,sfreq);
	pdf_remain.cumulate(i,-1*sfreq);
    }
    if (pdf_remain.samples() > 0)
	ent_remain = pdf_remain.entropy();

    if ((pdf_remain.samples() == 0) ||
	(pdf_split.samples() == 0))
	return -1;

    score = (ent_remain * pdf_remain.samples()) +
	(ent_split * pdf_split.samples());
//    printf("tostate %d remain %g %d split %g %d score %g\n",
//	   tostate, ent_remain, (int)pdf_remain.samples(),
//	   ent_split, (int)pdf_split.samples(), score);

    return score;
}

#if 0
static void split_state(EST_WFST &wfst, EST_WFST_Transition *trans)
{
    /* Split off a new state for given trans.  Add transitions    */
    /* to this new state for all transitions in (old) state trans */
    /* goes to                                                    */
    EST_Litem *tp;
    int nstate = wfst.add_state(wfst_final);
    int ostate = trans->state();

//    printf("state %d entropy %g\n",ostate,entropy(wfst.state(ostate)));
    /* must be done before adding the new transitions to nstate */
    trans->set_state(nstate);

    for (tp=wfst.state(ostate)->transitions.head(); tp != 0; tp = tp->next())
    {
	wfst.state_non_const(nstate)->
	    add_transition(0.0,  /* weight will be filled in later*/
			   wfst.state(ostate)->transitions(tp)->state(),
			   wfst.state(ostate)->transitions(tp)->in_symbol(),
			   wfst.state(ostate)->transitions(tp)->out_symbol());

    }
//    printf(" nstate %d entropy %g\n",nstate,entropy(wfst.state(nstate)));
//    printf(" ostate %d entropy %g\n",ostate,entropy(wfst.state(ostate)));

}
#endif

static void split_state(EST_WFST &wfst, LISP trans_list, int ostate)
{
    /* Split off a new state for given trans.  Add transitions    */
    /* to this new state for all transitions in (old) state trans */
    /* goes to                                                    */
    EST_Litem *tp;
    int nstate = wfst.add_state(wfst_final);
    LISP t;

    /* must be done before adding the new transitions to nstate */
    for (t=trans_list; t; t=cdr(t))
	trans(car(t))->set_state(nstate);

    for (tp=wfst.state(ostate)->transitions.head(); tp != 0; tp = tp->next())
    {
	wfst.state_non_const(nstate)->
	    add_transition(0.0,  /* weight will be filled in later*/
			   wfst.state(ostate)->transitions(tp)->state(),
			   wfst.state(ostate)->transitions(tp)->in_symbol(),
			   wfst.state(ostate)->transitions(tp)->out_symbol());

    }
}