#include <string>
#include <iostream>
#include <fstream>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <cstdlib>

#include "libherd.h"
#include "util.h"

using namespace std;

void HerdGen::consumeComments(ifstream& file, string& line) {
  while (file.is_open()) {
    getline(file, line);

    if (line.c_str()[0] != '#' && line.c_str()[0] != '\0') {
      return;
    }
  }
}

HerdGen::HerdGen(char* filename) {

  string line;
  ifstream file (filename);
  int i;

  if (! file.is_open()) {
    cerr << "HerdGen::HerdGen(): Unable to open file `" << filename << "'" << endl;
    exit(EXIT_FAILURE);
  }

  /* Format:
   * Number of phases
   * Initial probabilities (one per line)
   * Rates (one per line)
   */

  // Read number of branches
  consumeComments(file, line);
  n = strtol(line.c_str(), NULL, 10);

  initialProbabilities = new double[n];
  lengths = new int[n];
  rates = new double[n];

  consumeComments(file, line);

  i = 0;
  while (i < n && !file.eof()) {
    lengths[i] = strtol(line.c_str(), NULL, 10);
    i++;
    getline(file, line);
  }
  if (i < n) {
    cerr << "HerdGen::HerdGen(): Could not read " << n << " branch lengths."
	 << endl;
    delete this;
    exit(EXIT_FAILURE);
  }

  consumeComments(file, line);

  // Read transition rates.
  i = 0;
  while (i < n && !file.eof()) {
    rates[i] = - strtod(line.c_str(), NULL);
    i++;
    getline(file, line);
  }
  if (i < n) {
    cerr << "HerdGen::HerdGen(): Could not read " << n << " transition rates."
	 << endl;
    delete this;
    exit(EXIT_FAILURE);
  }

  consumeComments(file, line);

  // Read initial probabilities.
  double p;
  double ps = 0;
  i = 0;
  while (i < n && !file.eof()) {
    p = strtod(line.c_str(), NULL);
    initialProbabilities[i] = p;
    ps += p;
    getline(file, line);
    i++;
  }

  if (i < n) {
    cerr << "HerdGen::HerdGen(): Could not read " << n << " initial probabilities."
	 << endl;
    delete this;
    exit(EXIT_FAILURE);
  }

  if (ps <= 0) {
    cerr << "HerdGen::HerdGen(): Sum of probabilities is <= 0." << endl;
    delete this;
    exit(EXIT_FAILURE);
  }

  /* Normalise initial probabilty vector, just in case. */
  for (i=0; i < n; i++) {
    initialProbabilities[i] /= ps;
  }



#ifdef DEBUG
   cerr << "Initial probabilities: " << endl;
   for (int i = 0; i < n; i++)
     cerr << initialProbabilities[i] << endl;
   cerr << "Lengths: " << endl;
   for (int i = 0; i < n; i++)
     cerr << lengths[i] << endl;
   cerr << "Rates: " << endl;
   for (int i = 0; i < n; i++)
     cerr << rates[i] << endl;
#endif

}

HerdGen::~HerdGen() {
  delete[] initialProbabilities;
  delete[] rates;
  delete[] lengths;
}

void HerdGen::setUniformRandomSource(RandomSourceWrapper * randomSource) {

  rsw = randomSource;

}

/* Draw an Erlang sample.
 * The routine uses the fact that \sum \log u = \log \product u to reduce the
 * number of logarithms required (see e.g. Haverkort91). Furthermore, in order
 * to avoid numerical problems that can arise as \product u -> 0 for very long
 * branches, the computation is split up into blocks such that only a maximum
 * number of N uniforms u are multiplied in each block (see the remark in my
 * thesis). N is currently 100; this might be too conservative and could be
 * raised.
 */
inline double HerdGen::drawErlang(int length, double rate) {

  double x = 0;
  double u = 1;
  int i;

  for (i=1; i<=length; i++) {
    u *= rsw->uni();
    
    if (i % 100 == 0) {
      x += log(u);
      u = 1;
    }

  }

  x += log(u);

  return x / rate;
  //  return(log(u) / rate);  /* Note that we stored the negative rates. */

}


/* Draw an initial branch with probabilities as given in initialProbabilities[].
 * This comes from Levente's phnepa.cc.
 */
inline int HerdGen::drawBranch() {

  double
    r = rsw->uni(),
    x = initialProbabilities[0];

  int i = 0;

  while (x <= r && i < n) {
    i++;
    x = x + initialProbabilities[i];
  }

  return i;
}

/* Generate and return a random variate from the distribution stored in
 * this HerdGen object. This method uses the SimpleCount algorithm for HErD
 * in branch form. */
double HerdGen::getVariate() {

  int i = drawBranch();
  double r = 0;

  double rate = rates[i];
  int length = lengths[i];

  return drawErlang(length, rate);

}
