// rand_tree.cc   File for uniform random tree generator, with GPQUICK-2.1
// See ETL-TR-95-35 14 Nov 1995 Hitoshi Iba Random tree generator for GP
// W.B.Langdon cs.bham.ac.uk 9 August 1997
#define main_version "$Revision: 1.20 $"

//Modifications (in reverse order)

//WBL 20 May 98  BUGFIX lookup() following email from Christian Igel
//               <igel@neuroinformatik.ruhr-uni-bochum.de>

//WBL  7 Jul 98  Use RndOp

//WBL 12 Jun 98  Allow caller to specify tree shape to rand_tree

//WBL 11 Mar 98  Add code to create tables only one tree size

//WBL 10 Mar 98  Speed up by, eg using static structure for gtree and stack

//WBL 21 Jan 98  Use double precision (should take above 1e07 trees created)
//               remove code to remove trees less than FLT_EPSILON
//               Potentially changes GPQUICK_SEED, previous verison r1.4

//WBL 29 Dec 97  Optionally include number of different terminals and functions
//               of each arity when calculating tree shape probabilties
//               tree_count replaced by version in ntrees.cc r1.1
//               Also ensure scope of for loop variables is limited

//WBL 12 Aug 97  Use log of factorial, add remove i[4] loop

//WBL 11 Aug 97  Extend range of size before overflow
//               BUGFIX limit (efficiency saving only), 

//WBL  1 Aug 97  File created from rand_tree-test.cc

#include <assert.h>
#include "pch.h"
#include "chrome.h"

#define DEBUG true

//dont change!
#define max_arity 4

double* log_fact;
int log_fact_size = 0;
void new_log_fact(const int x) {
  log_fact = new double[x+1];
  log_fact_size = x+1;
  double f = 0;
  log_fact[0] = 0;
  for(int i=1;i<log_fact_size;i++) log_fact[i] = (f += log(double(i)));
}
inline
double log_factorial(const int x) {
  assert (x<log_fact_size);
  return log_fact[x];
}

double tree_count(const int size, int arity_count[], const int arity[]) {

  //actual count of distinct tree shapes with a particular choice of
  //node arities is given by Cn(s) (8). However we are interested in
  //the proportion of the total rather than the absolute number. This
  //espacially important since Cn(s) overflows.

  // value returned is log(actual number of trees)
  // arity_count[0] set to number of terminals

  int residue = size;
  {for(int i=1;i<=max_arity;i++) residue -= arity_count[i];}
  arity_count[0] = residue;

  double trees = log_factorial(size)-log(size);
  {for(int i=0;i<=max_arity;i++) trees -= log_factorial(arity_count[i]);}

  //now include number of ways of labeling tree with given terminal+functions
  {for(int i=0;i<=max_arity;i++) 
    if(arity[i] != 0)  trees += arity_count[i] * log(arity[i]);}

  return trees;
}

int limit(const int size, const int arity[], const int index) {
  return (arity[index]==0)? 0 : size / index;
}

class gtree {
//public:
  int arity;
  gtree* child[max_arity];
public:
  gtree() {};
//use static structures for speed
//   inline gtree(const int a, gtree* c[]) {
//     arity = a;
//     for(int i=0; i<max_arity;i++)
//       child[i] = (i<a)? c[i] : NULL;
//   }
//   inline ~gtree() {
//     for(int i=0; i<arity;i++)
//       delete child[i];
//   }
  inline int Arity() const               { return arity; };
#ifdef DEBUG
  inline void setarity(const int c)      { assert(c>=0&&c<max_arity);
					   arity = c; };
  inline gtree* Child(const int c) const { assert(c>=0&&c<arity);
					   return child[c]; };
  inline void setchild(const int c, /*const cc.64*/ gtree* t) { assert(c>=0&&c<arity);
						      child[c] = t; };
#else
  inline void setarity(const int c)      { arity = c; };
  inline gtree* Child(const int c) const { return child[c]; };
  inline void setchild(const int c, const gtree* t) { child[c] = t; };
#endif
friend ostream& operator << (ostream& os, const gtree& t);
};//end class gtree

ostream& operator << (ostream& os, const gtree& t) {
  os<<"("<<t.arity;
  for(int i=0; i<t.arity;i++)
    os<<*t.child[i];
  os<<")"<<flush;
  return os;
}

gtree fulllist[EXPRLEN];
gtree* fulltree(const int depth, const int arity, int& fip) {
  assert(fip<EXPRLEN);  //use static structure
  const int ip = fip++;
  fulllist[ip].setarity((depth>1)? arity:0); //function or terminal?
  for(int a=0; a<fulllist[ip].Arity(); a++) {
//  cout<<"Arity["<<ip<<"] "<<fulllist[ip].Arity()<<endl;
    fulllist[ip].setchild(a,fulltree(depth-1,arity,fip));
  }
  return &fulllist[ip];
}//end fulltree

class onedom {
  int offset;
  int size;
  int* data;
public:
  onedom(const int s, const int in[]) {
    size = s;
    data = new int[size];
    memcpy(data,in,size*sizeof(int));
    //find offset such that data is 1-dominated
    int c;
    for(offset=0;offset<size;offset++) {
      c = 0;
      for(int i=0;i<size;i++) {
	c += 1 - item(i); 
	//cout<<"i "<<i<<" count "<<c<<endl;
	if(c<=0) break;
      }
      //cout<<"offset "<<offset<<" count "<<c<<endl;
      if(c==1) break;
    }
    assert(c==1);
  }
  ~onedom() { delete[] data;}
  inline int item(const int index) const {return data[(index+offset)%size];}
};//endclass onedom


class stack {
public:
  int sp;
  int size;
//  gtree** data;
  gtree* data[EXPRLEN];
public:
  stack() : sp(0), size(EXPRLEN) {};
//use static structures for speed
//   stack(const int s) {
//     data = new gtree*[s];
//     size = s;
//     sp = 0;
//   }
#ifdef DEBUG
//  ~stack()            {assert(sp==0); delete[] data;}
  gtree* pop()        {assert(sp>0);   return data[--sp];  }
  void push(gtree* t) {assert(sp<size);       data[sp++]=t;}
#else
//  inline ~stack()            { delete[] data; }
  inline gtree* pop()        { return data[--sp];  }
  inline void push(gtree* t) { data[sp++]=t; }
#endif
};//endclass stack

gtree tlist[EXPRLEN];
stack sstack;

gtree* random_tree(const int size, const int arity_count[] ) {
//Given numbers of each function by arity, create a random tree
int* dyck = new int[size];
{for(int i=0; i<size; i++) dyck[i] = -1;}

//cout<<" random_tree "<<size<<flush;
//for(i=0; i<max_arity; i++) cout<<" "<<arity_count[i]<<flush;

int u=0;
{for(int i=1; i<=max_arity; i++) {
  for(int n=0; n<arity_count[i-1]; n++) {
    int r = rnd(size-u);
    int j, k;
    for(j=-1, k=0;;k++) {
      if(dyck[k]<=0) j++;
      if(j==r) {
	dyck[k] = i;
	u++;
	//cout<<u<<" "<<r<<" dyck["<<k<<"]="<<i<<endl;
	break;
      }
    }
  }
}}
{for(int i=0; i<size; i++) {if(dyck[i] == -1) dyck[i]=0;}}

//cout<<" Dyck ";
//for(i=0; i<size; i++) {
//  cout<<"x";
//  for(int j=0; j<dyck[i]; j++) cout<<"y";
//}
//cout<<flush;

//Before convert Dyck word to tree we need to convert it to
//1-dominated form. cf ETL-TR-95-35 

onedom* d = new onedom(size,dyck);
//onedom* d2 = new onedom(size,dyck); delete d2;

//cout<<" 1-dominated ";
//for(i=0; i<size; i++) {
//  cout<<"x";
//  for(int j=0; j<d->item(i); j++) cout<<"y";
//}
//cout<<flush;

//stack* s = new stack(size);

{for(int i=0; i<size; i++) {
  tlist[i].setarity(d->item(i)); //use static structure for speed
  for(int j=d->item(i); j>0; --j) tlist[i].setchild(j-1,sstack.pop());
  sstack.push(&tlist[i]);
//   gtree* t[max_arity];          original code
//   //cout<<"Popping "<<d->item(i)<<flush;
//   for(int j=d->item(i); j>0; --j) t[j-1] = s->pop();
//   gtree* n =  new gtree(d->item(i),t);
//   //cout<<" pushing ("<<n->arity<<")"<<flush;
//   s->push(n);
//   //cout<<" stack "<<s->sp<<" deep"<<endl;
}}

//cout<<" tree "<<*final<<endl;
gtree* final = sstack.pop();


//delete final;
//delete s;
delete d;
delete[] dyck;
return final;
}//end random_tree

class tree_chain {
public:
  double count;
  int* tree_arity_counts;
  tree_chain* next;
  tree_chain(const double c, int* const tac, tree_chain* const n):
    count(c), tree_arity_counts(tac), next(n) {};
};

class tree_table {
  int** tac;
  double* tac_index;
  int tac_size;
  double tac_total;
  void couttac(const int j) const {
    cout<<"tac_index["<<j<<"] "<<tac_index[j]<<"\t["<<flush;
    for(int k=0;k<max_arity-1;k++) cout<<tac[j][k]<<","<<flush;
    cout<<tac[j][max_arity-1]<<"]"<<endl;
  }
public:
  tree_table(const int size, const int arity[]) {
//all_arity
//Given tree size and Arities of functions
//return all possible combinations of arities 

  //max_arity = 4;
  double max = -FLT_MAX;
  tree_chain* head = NULL;
//cout<<"\nTrees of size "<<size<<flush;
//  for(int z=1;z<=4;z++) cout<<" "<<arity[z]; cout<<endl;
//cout<<"All possible combinations of arities for tree of size "<<size<<endl;
int i[5];
for(i[1]=0;i[1]<=limit((size-1),                   arity,1);i[1]++)
for(i[2]=0;i[2]<=limit((size-1)-i[1],              arity,2);i[2]++)
for(i[3]=0;i[3]<=limit((size-1)-i[1]-i[2]*2,       arity,3);i[3]++) {
  const int r3 = (size-1)-i[1]-i[2]*2-i[3]*3;
  i[4] = r3/4; 
  if(r3%4==0 && (arity[4]>0 || i[4]==0)) {
    //cout<<"("<<i[1]<<","<<i[2]<<","<<i[3]<<","<<i[4]<<") "<<flush;
    const double tree = tree_count(size,i,arity);
    //cout<<tree<<" "<<endl;
    if(tree>max) max = tree;
    int* savei = new int[max_arity];
    memcpy(savei,&i[1],sizeof(int)*max_arity);
    head = new tree_chain(tree,savei,head);
  }
}
//end all_arity

  //We rescale enties, by normalising we can convert back from logs
  //and discard many entries.  Set normalise count relative to best
  //inividual

    tree_chain* curr;
    double total = 0;
    for(curr=head; curr!=NULL; curr = curr->next) {
      curr->count = exp((curr->count)-max);
      total += curr->count;
    }

  //Normalise a second time so total is unity. We now discard small entries
    tree_chain* next;
    tree_chain* prev = NULL;
    tac_size = 0;
    double running = 0;
    for(curr=head; curr!=NULL; curr = next) {
      next = curr->next;
      const double normalise = curr->count / total;
/*    if(normalise < FLT_EPSILON) {//could be bigger to trim tables more
*	if(curr==head) head=next;
*	delete curr->tree_arity_counts;
*	delete curr;
*	if(prev!=NULL) prev->next = next;
*     }
*     else{ */
	prev = curr;
	tac_size++;
	running += normalise;
	curr->count = running;
/*    } */
    }
    if((running<.9999999||running>1.0000001) && tac_size>0)
      cout<<"ERROR "<<total<<" != "<<running;
    tac_total = running;

  //Now create tables and copy data into them
    tac       = new  int*[tac_size];
    tac_index = new double[tac_size];
    for(int j=0;j<tac_size;j++) {
      tac[j]       = head->tree_arity_counts;
      tac_index[j] = head->count;
      //couttac(j);
      next = head->next;
      delete head;
      head = next;
    }
    assert(head==NULL);
    //cout<<"\t"<<tac_size<<" table entries\n";
  }//end tree_table

  ~tree_table() {
    for(int j=0;j<tac_size;j++) delete[] tac[j];
    delete[] tac;
    delete[] tac_index;
  }

  void lookuperror(const char* s,const float f, const int i) {
    cout<<"ERROR "<<f<<s<<flush;
    couttac(i);
  }//end lookuperror

  int* lookup() {
    if(tac_size<=0) return NULL;
    const double f = drand_0to1()*tac_total;
    //based on time.h
    int bot = 0; //binary chop search
    int top = tac_size - 1;
    //cout<<"choping for "<<f<<" "<<flush;
    do {
      if(bot>top) lookuperror(" bot>top ",f,bot);
      const int i = (top+bot)/2;
      //cout<<tac_index[i]<<" "<<i<<" "<<bot<<" "<<top<<","<<flush;
      if     (i>0 && f < tac_index[i-1] ) top = i - 1;
      else if(       f > tac_index[i]   ) bot = i + 1;
      else {
	//cout<<" found "<<i<<" "<<tac_index[i]<<endl;
	return tac[i];
      }
    } while(TRUE);
  }//end lookup
};//end tree_table


class rand_tree_tables {
  tree_table** table;
  int lower_limit;
  int limit;
public:
  rand_tree_tables(const int maxlength, const int arity[]) {
    lower_limit = 1;
    limit = maxlength;
    table = new tree_table*[limit+1];
    for(int i=1;i<=limit;i++)
      table[i-1] = new tree_table(i,arity);
  }//end rand_tree_tables
  rand_tree_tables(const int arity[], const int len) { //one length only
    lower_limit = len;
    limit = len;
    table = new tree_table*[1];
    table[0] = new tree_table(len,arity);
  }//end rand_tree_tables

  ~rand_tree_tables() {
    for(int i=lower_limit;i<=limit;i++) delete table[i-lower_limit];
    delete[] table;
  }

  int* rand_tree_lookup(const int length) {
    return (length<lower_limit || length>limit)? NULL : 
                                          table[length-lower_limit]->lookup();
  }
};//end rand_tree_tables


//assumes ip set up  //based on SubInit
void Chrome::lable_tree(const gtree* tp, CSelector* const funcbag)
{
//assert(ip>=0&&ip<params->params[pMaxExpr]);
  const int a = tp->Arity();
  int i;

//  do {i=funcbag->roulette();} while (funclist[i]->argnum != a);

//  if (funclist[i]->varnum == 0)
      {SETNODE(expr[ip],probl->RndOp(0,a),0);} //faster!
//  else
//    {SETNODE(expr[ip],i,rnd(funclist[i]->varnum));}
  ip++;
  for(i=0;i<a;i++) lable_tree(tp->Child(i), funcbag);
}//end lable_tree

//assumes ip set up
BOOL Chrome::rand_tree(const int length, const int chrometree, const gtree* tp)
{
  //cout<<"rand_tree "<<length<<" "<<flush;
if(ip+length>params->params[pMaxExpr]) return FALSE; //no room in expr
if(tp!=NULL) 
  lable_tree(tp, probl->funcbag[chrometree]);
else {
  const int* tac = probl->rand_tree_lookup(length, chrometree);
  if(tac==NULL) return FALSE;
  gtree* tp = random_tree(length,tac);
  lable_tree(tp, probl->funcbag[chrometree]);
//  delete tp; now static structure for speedup
}
  return TRUE;
}//end rand_tree;

//must be called after all functions have been added
//call this OR init_rand_tree(int size_limit) NOT BOTH
//create tables for one length only
void Problem::init_1rand_tree(const int size) {
  new_log_fact(size); 
  assert(NUM_TREES==1);
//for (int t=0;t<NUM_TREES;t++) { -- doesnt handle multiple trees!
    int arity[max_arity+1];
    memset(arity,0,sizeof(arity));
    for (int i=0;i<funcbag[0]->count();i++) //t
      {
	const FHANDLE f = funcbag[0]->GetHandle(i); //t
	const int a = funclist[f]->argnum;
	assert(a<=max_arity);
#ifdef RAND_TREE_ARITIES
	arity[a]++;
#else
	arity[a] = 1;
#endif
      } 
    rand_tree[0] = new rand_tree_tables(arity,size); //t
//  }
}

//must be called after all functions have been added
void Problem::init_rand_tree(int size_limit) {
  if(size_limit<0) size_limit = -size_limit;//use absolute value
  new_log_fact(size_limit);
  for (int t=0;t<NUM_TREES;t++) {
    //time_t now;
    //time (&now);
    //cout<<"Creating Random tree lookup table for tree "<<t<<" at ";
    //cout<<ctime(&now)<<flush;

    int arity[max_arity+1];
    memset(arity,0,sizeof(arity));
    for (int i=0;i<funcbag[t]->count();i++)
      {
	const FHANDLE f = funcbag[t]->GetHandle(i);
	const int a = funclist[f]->argnum;
	assert(a<=max_arity);
#ifdef RAND_TREE_ARITIES
	arity[a]++;
#else
	arity[a] = 1;
#endif
      } 
    //could look for common arity masks between trees, but simplier to
    //assume they are different and so no saving can be made. 
    rand_tree[t] = new rand_tree_tables(size_limit,arity);

    //time (&now);
    //cout<<" "<<size_limit<<" done "<<ctime(&now)<<endl;
  }
}//end init_rand_tree

//cannot be called before init_rand_tree
int* Problem::rand_tree_lookup(const int length, const int t) const
{
  return rand_tree[t]->rand_tree_lookup(length);
}