/*    Copyright 2010, 2011 Ernest Turro
**
**    This program is free software; you can redistribute it and/or modify
**    it under the terms of the GNU General Public License as published by
**    the Free Software Foundation; either version 2 of the License, or
**    (at your option) any later version.
**
**    This program is distributed in the hope that it will be useful,
**    but WITHOUT ANY WARRANTY; without even the implied warranty of
**    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
**    GNU General Public License for more details.
**
**    You should have received a copy of the GNU General Public License
**    along with this program; if not, write to the Free Software
**    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/

#include <iostream>
#include <iomanip>
#include <map>
#include <sstream>
#include <math.h>
#include <string>
#include <string.h>
#include <fstream>
#include <assert.h>
#include <algorithm>
#include <time.h>
#include <boost/numeric/ublas/matrix_sparse.hpp>
#include <boost/numeric/ublas/matrix_proxy.hpp>
#include <boost/numeric/ublas/operation.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_rng.h>
#ifdef _OPENMP
  #include <omp.h>
  #define OMP_GET_MAX_THREADS omp_get_max_threads()
  #define OMP_GET_THREAD_NUM omp_get_thread_num()
#else
  #define OMP_GET_MAX_THREADS 1
  #define OMP_GET_THREAD_NUM 0
#endif // _OPENMP

#include "sokal.hh"

#define QUOTE_(x) #x
#define QUOTE(x) QUOTE_(x)

#ifndef VERSION
  #define VERSION ?
#endif

#define HITCHAR '>'
#define COMMENTCHAR '@'

namespace ublas = boost::numeric::ublas;
using namespace std;
using ublas::prod;
using ublas::inner_prod;
using ublas::matrix_row;
using ublas::matrix_column;

typedef ublas::compressed_matrix<bool> boolMat;
typedef boolMat::iterator1 boolMatIt1;
typedef boolMat::iterator2 boolMatIt2;
typedef ublas::compressed_matrix<int> intMat;
typedef intMat::iterator1 intMatIt1;
typedef intMat::iterator2 intMatIt2;
typedef ublas::compressed_matrix<double> doubleMat;
typedef intMat::iterator1 doubleMatIt1;
typedef intMat::iterator2 doubleMatIt2;

bool secondRowGreater(matrix_row<boolMat> a, matrix_row<boolMat> b) {
  for(matrix_row<boolMat>::iterator ita = a.begin(), itb = b.begin(); ; ita++, itb++) {
    bool enda = ita==a.end();
    bool endb = itb==b.end();
    if(enda && endb) return false;
    if(enda && !endb) return true;
    if(endb && !enda) return false;
    
    if(ita.index() < itb.index()) return false;
    if(ita.index() > itb.index()) return true;
  }
  cerr << "Error in secondRowGreater()\n";
  exit(1);
}

boolMat M(0,0); // GLOBAL M matrix
bool compareRowsByInt(int i, int j) {
  return(secondRowGreater(matrix_row<boolMat>(M,j), matrix_row<boolMat>(M,i)));
}

int powerof2 (unsigned int x) {
  while (((x & 1) == 0) && x > 1) x >>= 1;
  return (x == 1);
}


bool equalRows(matrix_row<boolMat> a, matrix_row<boolMat> b) {
  for(matrix_row<boolMat>::iterator ita = a.begin(), itb = b.begin(); ; ita++, itb++) {
    bool enda = ita==a.end();
    bool endb = itb==b.end();
    if(enda && endb) return true;
    if(enda && !endb) return false;
    if(endb && !enda) return false;
    if(ita.index() != itb.index()) return false;
  }
  cerr << "Error in equalRows()\n";
  exit(1);
}

void printUsage(char *bin, ostream& out) {
  out << "Usage: mmseq [OPTIONS...] hits_file output_base" << endl
      << endl
      << "Mandatory arguments:"
      << endl
      << "  hits_file          hits file generated with `bam2hits`\n" 
      << "  output_base        base name for output files"
      << endl << endl
      << "Optional arguments:\n"
      << "  -alpha FLOAT       value of alpha in Gamma prior for mu (default: 0.1)" << endl
      << "  -beta FLOAT        value of beta in Gamma prior for mu (default: 0.1)" << endl
      << "  -max_em_iter INT   maximum number of EM iterations (default: 1000)" << endl
      << "  -epsilon FLOAT     minimum loglik ratio between successive EM iterations (default: 0.1)" << endl
      << "  -gibbs_iter INT    number of Gibbs iterations (default: 16384)" << endl
      << "  -gibbs_ss INT      subsampling interval for Gibbs output (default: gibbs_iter/1024)" << endl
      << "  -seed INT          seed for the PRNG in thread 0 (default: 1234)" << endl
      << "  -traces            save Gibbs posterior trace files" << endl
      << "  -debug             output additional diagnostic files" << endl
      << "  -help              print this help message" << endl
      << "  -version           print the version" << endl
      << endl;
}

int main(int argc, char **argv) {
  // Set max_threads here because on Mac OS X omp_get_max_threads() erroneously returns 1 in parallel regions
  int max_threads=OMP_GET_MAX_THREADS;

  // DEFAULT PARAMETER VALUES
  double alpha=0.1; 
  double beta=0.1;

  int max_em_iter=1000;
  double epsilon=0.1;

  int gibbs_iter=16384;
  int trace_length=1024;
  int gibbs_ss=gibbs_iter/trace_length;

  int seed=1234;

  bool debug=false;
  bool traces=false;


  // parse arguments
  vector<string> arguments;
  for(int i=1; i < argc; i++) arguments.push_back(string(argv[i]));

  while(true) {
    if(arguments.size() >0 && arguments[0]=="-alpha") {
      arguments.erase(arguments.begin());
      alpha=strtod(arguments[0].c_str(), NULL);
      arguments.erase(arguments.begin());
    } else if(arguments.size() >0 && arguments[0]=="-beta") {
      arguments.erase(arguments.begin());
      beta=strtod(arguments[0].c_str(), NULL);
      arguments.erase(arguments.begin());
    } else if(arguments.size() > 0 && arguments[0]=="-max_em_iter") {
      arguments.erase(arguments.begin());
      max_em_iter=atoi(arguments[0].c_str());
      arguments.erase(arguments.begin());
    } else if(arguments.size() > 0 && arguments[0]=="-epsilon") {
      arguments.erase(arguments.begin());
      epsilon=strtod(arguments[0].c_str(), NULL);
      arguments.erase(arguments.begin());
    } else if(arguments.size() > 0 && arguments[0]=="-gibbs_iter") {
      arguments.erase(arguments.begin());
      gibbs_iter=atoi(arguments[0].c_str());
      arguments.erase(arguments.begin());
    } else if(arguments.size() > 0 && arguments[0]=="-gibbs_ss") {
      arguments.erase(arguments.begin());
      gibbs_ss=atoi(arguments[0].c_str());
      arguments.erase(arguments.begin());
    } else if(arguments.size() > 0 && arguments[0]=="-seed") {
      arguments.erase(arguments.begin());
      seed=atoi(arguments[0].c_str());
      arguments.erase(arguments.begin());
    } else if(arguments.size() > 0 && arguments[0]=="-traces") {
      traces=true;
      arguments.erase(arguments.begin());
    } else if(arguments.size() > 0 && arguments[0]=="-debug") {
      debug=true;
      arguments.erase(arguments.begin());
    } else if(arguments.size() > 0 && ( arguments[0]=="-h" || arguments[0]=="--help" || arguments[0]=="-help")) {
      cerr << "Calculate mmseq expression estimates.\n";
      printUsage(argv[0], cerr);
      exit(1);
    } else if(arguments.size() > 0 && ( arguments[0]=="-v" || arguments[0]=="--version" || arguments[0]=="-version")) {
      cerr << "mmseq-" << QUOTE(VERSION) << endl;
      exit(1);
    } else {
      if(arguments.size()==2) {
        break;
      } else {
        if(arguments.size()>0 && arguments[0][0]=='-') {
          cerr << "Error: unrecognised option " << arguments[0] << ".\n";
        } else {
          cerr << "Error: mandatory arguments missing.\n";
        }
        printUsage(argv[0], cerr);
        exit(1);
      }
    }
  }

  if( gibbs_iter % gibbs_ss != 0 ) {
    cerr << "Error: gibbs_iter must be divisible by gibbs_ss.\n";
    printUsage(argv[0],cerr);
    exit(1);
  }

  gibbs_ss=gibbs_iter/trace_length;

  if( gibbs_iter <= 0 || trace_length <= 0 ) {
    cerr << "Error: no. of iteratons or trace length <= 0. Possible integer overflow - is gibbs_iter too high?\n";
    printUsage(argv[0],cerr);
    exit(1);
  }

  if(!powerof2(trace_length)) {
    cerr << "Error: gibbs_iter/gibbs_ss must be a power of 2.\n";
    printUsage(argv[0],cerr);
    exit(1);
  }

  string hits_file = arguments[0];
  string output_base(arguments[1]);

  // START helpers
  ifstream ifs;
  ofstream ofs;
  string str, str2;
  vector<string>::iterator strit;
  char buf16384[16384];
  // END helpers

  // Open hits file for reading
  ifs.open(hits_file.c_str());
  if(!ifs.good()) {
    cerr << "Error reading hits file \"" << hits_file << "\".\n";
    exit(1);
  }

  cout << "Running mmseq with parameters:\n"
       << "  alpha:         " << alpha << endl
       << "  beta:          " << beta << endl
       << "  max_em_iter:   " << max_em_iter << endl
       << "  epsilon:       " << epsilon << endl
       << "  gibbs_iter:    " << gibbs_iter << endl
       << "  gibbs_ss:      " << gibbs_ss << endl
       << "  seed[0]:       " << seed << endl
       << "  traces:        " << traces << endl
       << "  debug:         " << debug << endl
       << "  threads:       " << max_threads << endl;


  // Read meta-data at top of hits file 
  // Expecting several lines of @TranscriptMetaData ID length 
  // followed by several lines of @GeneIsoforms GID TID1 TID2 ...
  // followed optionally by several lines of @IdenticalTranscripts TID1 TID2 ...
  // first line after header lines should contain >read_id and is saved in str 
  
  map<string, int> sidLen; 
  vector<string> transcriptList;
//  map<string, bool> sidRemcand; 

  map<string, vector<string> > gene2transcripts;
  vector< vector<string> > identical_transcripts;
 
  {
    while(true) {
      ifs.getline(buf16384,16384);
      if(ifs.eof()) {
        cerr << "Error. Malformed hits file.\n";
        exit(1);
      }
      str=buf16384;
      istringstream iss(str);
      vector<string> tokens;
      copy(istream_iterator<string>(iss),
        istream_iterator<string>(),
        back_inserter<vector<string> >(tokens));
      if(tokens[0].compare("@TranscriptMetaData")==0) {
        sidLen.insert(make_pair(tokens[1], atoi(tokens[2].c_str())));
        transcriptList.push_back(tokens[1]);
//        sidRemcand.insert(make_pair(tokens[1], tokens[3].compare("true")==0));
      } else if(tokens[0].compare("@GeneIsoforms")==0) {
        vector<string> isoforms;
        for(int i=2; i < tokens.size(); i++) isoforms.push_back(tokens[i]);
        gene2transcripts.insert(make_pair(tokens[1], isoforms));
      } else if(tokens[0].compare("@IdenticalTranscripts")==0) {
        vector<string> transcripts;
        for(int i=1; i < tokens.size(); i++) transcripts.push_back(tokens[i]);
        identical_transcripts.push_back(transcripts);
      } else if(tokens[0][0] == HITCHAR) {
        break;
      } else {
        cerr << "Error. Malformed hits file.\n";
        exit(1);
      }
    }
  }

  // START Initialise variables M,  m, n, k, mu, l
  map<string, int> sidIndex;
  map<int, string> indexSid;
  map<vector<int>, int > indexComb;
  vector<int> k; 

  int m=0; // counts distinct combinations encountered
  int n=0; // counts distinct transcripts encountered

  vector<int> doublehits;
  bool doublehit;

  int numbermappedreads=0;
  
  while(!ifs.eof()) {
    doublehit=false;
    if(str[0]==HITCHAR) {
      numbermappedreads++;
      vector<int> comb;
      while(!(ifs >> str).eof() && str[0] != HITCHAR) {
        if(sidIndex.count(str)==0) { sidIndex[str] = n; indexSid[n]=str;n++; doublehits.push_back(0); }
        if(find(comb.begin(), comb.end(), sidIndex[str]) == comb.end()) comb.push_back(sidIndex[str]); 
	      else {
  //        cout << "Double hit: " << str << " " << indexSid[sidIndex[str]] << endl;
          doublehits[sidIndex[str]]++;
          doublehit=true;
        }
        // TODO: think about what to do when a read hits the same transcript in more than one place.
      }
      sort(comb.begin(), comb.end());
      if(indexComb.count(comb) == 0) {
	      cout << "Found " << n << " transcripts in " << m << " transcript combinations.\r";
        //	for(vector<int>::iterator intit=comb.begin(); intit != comb.end(); intit++) cout << *intit << " ";
        //	cout << endl;
	      indexComb[comb] = m;
        m++;
        k.resize(m);
        /* Here I'd rather do M.resize(c,t,true), but this is not implemented in Boost yet */
        // Dynamic matrix
        int factor1 = 1; int factor2 = 1;
        if(M.size1() < m ) factor1=2;
        if(M.size2() < n ) factor2=2;
        if(factor1==2 || factor2==2 ) {
          boolMat temp(M);
          int s1=factor1==2?m*factor1: M.size1();
          int s2=factor2==2?n*factor2: M.size2();
          M.resize(s1,s2,false);
          //cout <<"Now: M.size1()=="<<M.size1()<<", M.size2() == " << M.size2() << endl << flush;
          for(boolMatIt1 i1 = temp.begin1(); i1 != temp.end1(); ++i1) {
            for (boolMatIt2 i2 = i1.begin(); i2 != i1.end(); ++i2) {
              M(i2.index1(),i2.index2())= *i2;
            }
          }
        }
        for(vector<int>::iterator intit=comb.begin(); intit != comb.end(); intit++)
          M(m-1, *intit) = 1;
      }
      k[indexComb[comb]]++;
    }
  }
  ifs.close(); ifs.clear();
  cout << endl;

/*  cout << "Checking if sufficient memory (" << sizeof(double)*(n + identical_transcripts.size() + gene2transcripts.size())* trace_length *1/1048576<< "MB) is available...";
  try {
    double *temp=new double[(n + identical_transcripts.size() + gene2transcripts.size())* trace_length];
    delete [] temp;
  } catch(...) {
    cerr << "Insufficient memory available. Aborting." << endl;
    exit(1);
  }
  cout << "done.\n";
*/

  // Resize M to get rid of excessive rows/columns from dynamic bit above
  boolMat temp(M);
  M.resize(m,n,false);
  for(boolMatIt1 i1 = temp.begin1(); i1 != temp.end1(); ++i1) {
    for (boolMatIt2 i2 = i1.begin(); i2 != i1.end(); ++i2) {
      M(i2.index1(),i2.index2())= *i2;
    }
  }

  time_t before,after;

/*
  // Absent transcript removal step, based on M and removal candidate status

  cout << "Transposing M (this may take a while)..." << flush;
  time(&before);
  M = trans(M);
  time(&after); 
  cout << "done (" << difftime(after,before) << " seconds).\n";

  cout << "Removing absent transcripts..." << flush;

  // Sort rows
  {
    vector<int> Mord(n,0);
    for(int i=1; i < n; i++) Mord[i] = Mord[i-1]+1;
    sort(Mord.begin(), Mord.end(), compareRowsByInt);
    map<int, string> indexSid_ = indexSid;
    for(int t=0; t < n; t++) indexSid[Mord[t]] = indexSid_[t];
    boolMat newMat(n,m);
    for(int t=0; t < n; t++) {
      matrix_row<boolMat> mr(M, t);
      for(matrix_row<boolMat>::iterator it = mr.begin(); it != mr.end(); ++it) {
        newMat(Mord[t],it.index()) = 1;
      }
    }
    M = newMat;
  }

  // Remove removal candidate iff it has duplicates and not all its duplicates are also removal candidates
  vector<int> removal_indices;
  for(int t=0; t < n; t++) {
    if(sidRemcand[indexSid[t]]) {
      vector<int> dup_indices(t,1);
      int start=t,end=t;
      while(start > 0 && equalRows(matrix_row<boolMat>(M,t), matrix_row<boolMat>(M, start-1))) start--;
      while(end < n-2 && equalRows(matrix_row<boolMat>(M,t), matrix_row<boolMat>(M, end+1))) end++; 
      if(end-start > 1) {
        bool allremcands=true;
        vector<int> potential_indices;
        for(int v=start; v <= end; v++) {
          if(!sidRemcand[indexSid[v]]) {
            allremcands=false;
          } else {
            potential_indices.push_back(v);
          }
        }
        if(!allremcands) removal_indices.insert(removal_indices.end(), potential_indices.begin(), potential_indices.end());
      }
    }
  }
  sort(removal_indices.begin(), removal_indices.end());
  removal_indices.resize(unique(removal_indices.begin(), removal_indices.end()) - removal_indices.begin());
#ifdef DEBUG
  ofs.open("removedTranscripts.txt");
  for(vector<int>::iterator it=removal_indices.begin(); it != removal_indices.end(); it++) {
    ofs << indexSid[*it] << " " << *it << endl;
  }
  ofs.close(); ofs.clear();
#endif

//  ofs.open("Mcols_before.txt");
//  for(int t=0; t < n; t++) {
//    ofs << indexSid[t] << "(" << sidRemcand[indexSid[t]] << ","
//    << !(find(removal_indices.begin(),removal_indices.end(),t)==removal_indices.end())
//    << ") ";
//    matrix_row<boolMat> mr(M, t);
//    for (matrix_row<boolMat>::iterator it = mr.begin(); it != mr.end(); ++it) {
//      ofs << it.index() << " ";
//    }
//    ofs << endl;
//  }
//  ofs.close();ofs.clear();
//
  if(removal_indices.size()>0) {
    // Shift indexSid
    int offset=0;
    for(int i=0; i < n; i++) {
      if(i==removal_indices[offset]) offset++;
      else {
        indexSid[i-offset]=indexSid[i];
        doublehits[i-offset]=doublehits[i];
      }
    }
    for(int i=n-removal_indices.size(); i <n; i++) indexSid.erase(i);

    // Update sidIndex just in case we need it
    sidIndex.clear();
    for(int i=0; i < indexSid.size(); i++) sidIndex[indexSid[i]] = i;
     
    // Reconstruct Mt without removed transcripts
    boolMat newMat(n-removal_indices.size(),m);
    int remind=0;
    int currow=0;
    for(int t=0; t < n; t++) {
      if(remind == removal_indices.size() || t != removal_indices[remind]) {
        matrix_row<boolMat> mr(M, t);
        for (matrix_row<boolMat>::iterator it = mr.begin(); it != mr.end(); ++it) {
          newMat(currow,it.index()) = 1;
        }
        currow++;
      } else {
        remind++;
      }
    }
    M = newMat;
    n=n-removal_indices.size();
  }

  cout << "done\n";
*/  
  cout << "Transposing M (this may take a while)..." << flush;
  time(&before);
  boolMat Mt = trans(M);
  time(&after);
  cout << "done (" << difftime(after,before) << " seconds).\n";

  vector<int> Mrowsum(m,0);
  vector<int> Mcolsum(n,0);
  for(boolMatIt1 i1 = M.begin1(); i1 != M.end1(); ++i1) {
    for (boolMatIt2 i2 = i1.begin(); i2 != i1.end(); ++i2) {
      Mrowsum[i2.index1()]++;
      Mcolsum[i2.index2()]++;
    }
  }

  ///////////////
  
  vector<double> l(n); // l
  for(unsigned t=0; t < n; t++)  {
    l[t] = (double)sidLen[indexSid[t]] * (double)numbermappedreads/1000000000.0; // apply this constant to get RPKM
    if(l[t]==0) { // sidLen[indexSid[t]] doesn't exist
      cerr << "Error: transcript '" << indexSid[t] << "' has hits but not found in @TranscriptMetaData header section.\n";
      exit(1);
    }
  }

  vector<vector<int> > counts_shared; // counts shared with different sized sets of transcripts

  for(int t=0; t < n; t++){
    vector<int> counts(100,0); // bin together sets >= 100
    counts_shared.push_back(counts);
  }
  // Set starting mu's
  ublas::vector<double> mu(n);
  for(int t=0; t < n; t++) mu[t] = 0.0;
  #pragma omp parallel
  {
    boolMatIt2 i1 = M.begin2();
    bool firstloop=true;
    int row;
    #pragma omp for schedule(static) 
    for(int t=0; t < n; t++){
      if(firstloop) {
        advance(i1,t);
        firstloop=false;
      }
      for (boolMatIt1 i2 = i1.begin(); i2 != i1.end(); ++i2) {
        row=i2.index1();
        mu[t] += (double)k[row]/Mrowsum[row];
        counts_shared[t][min(Mrowsum[row], 100)-1] += k[row];
      }
      mu[t] /= l[t];
      i1++;
    }
  }

  // END Initialise variables M, m, n, k, mu, l

  if(debug) {
    // Output shared counts histograms
    ofs.open((output_base + ".sharedcounts").c_str());
    for(strit = transcriptList.begin(); strit != transcriptList.end(); strit++) {
      ofs << *strit << "\t";
      for(int i=0; i < 100; i++) {
        if(sidIndex.count(*strit) > 0 ) {
          ofs << counts_shared[sidIndex[*strit]][i] << "\t";
        } else {
          ofs << "0\t";
        }
      }
      ofs << endl;
    }
    ofs.close();ofs.clear();

    ofs.open((output_base + ".doublehits").c_str());
    for(int i=0;i < n; i++) ofs << doublehits[i] <<  endl;
    ofs.close(); ofs.clear();

    ofs.open((output_base + ".k").c_str());
    for(int i=0;i < m; i++) ofs << k[i] <<  endl;
    ofs.close(); ofs.clear();

    ofs.open((output_base+".M").c_str());
    for (boolMatIt1 i1 = M.begin1(); i1 != M.end1(); ++i1) {
      for (boolMatIt2 i2 = i1.begin(); i2 != i1.end(); ++i2) {
        ofs << i2.index1() << "\t" << i2.index2() << endl;
      }
    }
    ofs.close(); ofs.clear();
   
    ofs.open((output_base+".Mt-nodups").c_str());
    vector<int> dupTs;
    ofstream ofs2((output_base + ".dupIDs").c_str());
    for (boolMatIt1 i1 = Mt.begin1(); i1 != Mt.end1(); ++i1) {
      if(i1.index1()>0 && equalRows(matrix_row<boolMat>(Mt, i1.index1()), matrix_row<boolMat>(Mt,i1.index1()-1))) {
        ofs2 << indexSid[i1.index1()] << endl;
      } else {
        for (boolMatIt2 i2 = i1.begin(); i2 != i1.end(); ++i2) {
          ofs << i2.index1() << "\t" << i2.index2() << endl;
        }
      }
    }
    ofs.close(); ofs.clear();
    ofs2.close();ofs2.clear();
  } // end debug

  if(debug) { 
    ofs.open((output_base + ".trace_em").c_str());
  }

  ublas::vector<double> mu_temp(n); // build up next iteration's mu before overwriting mu

  double loglik=0; // log likelihood

  #pragma omp parallel
  {
    #pragma omp for schedule(static) reduction(+:loglik) 
    for(int i=0; i < m; i++) {
      matrix_row<boolMat > mr (M, i);
      loglik += k[i]*log(inner_prod(mr, mu));
    }
    #pragma omp for schedule(static) reduction(+:loglik) 
    for(int t=0; t < n; t++) loglik -= mu[t]*l[t];
  }
  double loglik_temp=0; // next iteration's log likelihood
  double llr=epsilon+1; // log likelihood ratio between successive iterations

  int iter=0;
  cout.precision(5);
  cout.setf(ios::fixed,ios::floatfield);
  while(iter < max_em_iter && llr > epsilon) {
    cout << "EM iteration " << iter << flush;

    if(debug) {
      for(unsigned t=0; t < n; t++)
        ofs << mu[t] << " "; 
      ofs << endl;
    }

    loglik_temp=0;

    #pragma omp parallel 
    {
      int threadid=OMP_GET_THREAD_NUM;

      boolMatIt1 i1 = Mt.begin1();
      bool firstloop=true;
      int row;

      #pragma omp for schedule(static) 
      for(int t=0; t < n; t++) {
       if(firstloop) {
          advance(i1,t);
          firstloop=false;
        }
        double sum=0;
        for (boolMatIt2 i2 = i1.begin(); i2 != i1.end(); ++i2) {
          row=i2.index2();
          matrix_row<boolMat > mr (M, row);
          sum += k[row]/inner_prod(mr, mu);
        }
        mu_temp[t] = mu[t]*sum/l[t];
        i1++;
      }
    
      #pragma omp for schedule(static) reduction(+:loglik_temp) 
      for(int i=0; i < m; i++) {
        matrix_row<boolMat > mr (M, i);
        loglik_temp += k[i]*log(inner_prod(mr, mu_temp));
      }
      #pragma omp for schedule(static) reduction(-:loglik_temp) 
      for(int t=0; t < n; t++) loglik_temp -= mu_temp[t]*l[t];
    }
    mu = mu_temp;
    llr = loglik_temp - loglik;
    loglik = loglik_temp;
        
    cout << ", log likelihood ratio: " << llr << "            \r";
   
    iter++;
  }

  if(debug) {
    ofs.close(); ofs.clear();
  }

  cout <<  endl;

  ublas::vector<double> mu_em = mu;

  // gibbs trace file
  char *tg_file;
  if(traces) {
    tg_file = new char[(output_base + ".trace_gibbs").length()+1];
    strcpy(tg_file, (output_base + ".trace_gibbs").c_str());
  } else {
    tg_file = tmpnam (NULL);
    cerr << "Saving temporary traces in " << tg_file << endl;
  }
  ofs.open(tg_file, ios_base::trunc);

  for(unsigned t=0; t < n; t++)
    ofs << indexSid[t] << " "; 
  ofs << endl;

//  gsl_rng * rg = gsl_rng_alloc (gsl_rng_mt19937); //random number generator
  gsl_rng * rg[max_threads];
  for(int i=0; i < max_threads; i++) {
    rg[i] =  gsl_rng_alloc (gsl_rng_mt19937); 
    gsl_rng_set(rg[i], seed+i);
  }
  

  // set up X matrix with non-zero's in appropriate cells
  intMat X(m,n);
  for(boolMatIt1 i1 = M.begin1(); i1 != M.end1(); ++i1) {
    for(boolMatIt2 i2 = i1.begin(); i2 != i1.end(); ++i2) {
      X(i2.index1(), i2.index2()) = 1;
    }
  }

  int *Xcolsum = new int[n];
  int * Xcolsums = new int[n*max_threads]; // each thread calculates partial Xcolsums
  for(iter=0; iter < gibbs_iter; iter++) {
    cout << "Gibbs iteration " << iter << "       \r";
    
    memset(Xcolsum,0,n*sizeof(int));
    memset(Xcolsums,0, n*max_threads*sizeof(int));

    #pragma omp parallel
    {
      int threadid=OMP_GET_THREAD_NUM;

      // Update X's
      intMatIt1 i1 = X.begin1();
      bool firstloop=true;
      #pragma omp for schedule(static) 
      for(int i=0; i < m; i++) {
        if(firstloop) {
          advance(i1,i);
          firstloop=false;
        }

        unsigned int x[Mrowsum[i1.index1()]];
        double p[Mrowsum[i1.index1()]];
        int count=0;

        for(intMatIt2 i2 = i1.begin(); i2 != i1.end(); ++i2) {
          p[count] = mu[i2.index2()];
          count++;
        }
        
        gsl_ran_multinomial(rg[threadid], Mrowsum[i1.index1()], k[i1.index1()], p, x);

        count=0;
        for(intMatIt2 i2 = i1.begin(); i2 != i1.end(); ++i2) {
          X(i2.index1(), i2.index2()) = x[count];
        //  #pragma omp atomic
        //  myXcolsum[i2.index2()] += x[count];
          Xcolsums[i2.index2()+n*threadid] += x[count];
          count++;
        }
        i1++;
      }

      #pragma omp barrier
      
      // sum up threads' partial Xcolsums
      #pragma omp for schedule(static)
      for(int t=0; t < n ; t++) {
        for(int i=0; i < max_threads; i++) Xcolsum[t] += Xcolsums[t+i*n];
      }

      #pragma omp barrier
      // Now Xcolsum is ready

      // Update mu's
      #pragma omp for schedule(static)
      for(int t=0; t < n ; t++) {
        mu[t] = gsl_ran_gamma(rg[threadid], alpha + Xcolsum[t] , 1.0/(beta + l[t]));
      }
    }

    if(iter % gibbs_ss == 0) {
      for(unsigned t=0; t < n; t++)
        ofs << mu[t] << " "; 
      ofs << endl;
    }
  }

  delete [] Xcolsum;
  delete [] Xcolsums;

  cout << endl;
  for(int i=0; i < max_threads; i++) gsl_rng_free (rg[i]);
  ofs.close(); ofs.clear();

  cout << "Amalgamating transcripts and calculating Monte Carlo standard errors..." << flush;

  // read in gibbs traces and add up identical transcripts and isoforms within genes

  double * mu_trace = new double[n * trace_length]; // IMPORTANT to divide in the brackets else potential overflow

  ifs.open(tg_file);
  assert(ifs.good());
  getline(ifs,str); // skip first line which contains transcript IDs
  for(int i=0; i < trace_length; i++) {
    for(int t=0; t < n; t++) ifs >> mu_trace[t*trace_length + i];
  }
  ifs.close();ifs.clear();

  if(traces) {
    delete tg_file;
  } else {
    // Try to remove temporary gibbs trace file
    remove(tg_file);
  }

  double * mu_trace_identical = new double[identical_transcripts.size() * trace_length];

  vector<double> mu_em_identical(identical_transcripts.size(), 0);
  for(int i=0; i < identical_transcripts.size() * trace_length; i++) mu_trace_identical[i]=0;
  {
    int t=0;
    for(vector< vector<string> >::iterator idit = identical_transcripts.begin();
      idit != identical_transcripts.end(); idit++) {
      for(strit = idit->begin(); strit != idit->end(); strit++) {
        if(sidIndex.find(*strit) != sidIndex.end()) {
          mu_em_identical[t] += mu_em[sidIndex[*strit]];
          for(int i=0; i < trace_length; i++) {
            mu_trace_identical[t*trace_length + i] += mu_trace[sidIndex[*strit]*trace_length + i]; 
          }
        }
      }
      t++;
    }
  }

  double * mu_trace_gene = new double[gene2transcripts.size() * trace_length];
  vector<double> mu_em_gene(gene2transcripts.size(), 0);
  for(int i=0; i < gene2transcripts.size() * trace_length; i++) mu_trace_gene[i]=0;
  {
    int g=0;
    for(map<string, vector<string> >::iterator git = gene2transcripts.begin();
      git != gene2transcripts.end(); git++) {
      for(strit = (git->second).begin(); strit != (git->second).end(); strit++) {
        if(sidIndex.find(*strit) != sidIndex.end()) {
          mu_em_gene[g] += mu_em[sidIndex[*strit]];
          for(int i=0; i < trace_length; i++) {
            mu_trace_gene[g*trace_length + i] += mu_trace[sidIndex[*strit]*trace_length + i]; 
          }
        }
      }
      g++;
    }
  }

  // output the traces for sum of identical transcripts and transcripts belonging to the same gene
  if(traces) {
    char *tgi_file = new char[(output_base + ".identical.trace_gibbs").length()+1];
    strcpy(tgi_file, (output_base + ".identical.trace_gibbs").c_str());
    char *tgg_file = new char[(output_base + ".gene.trace_gibbs").length()+1];
    strcpy(tgg_file, (output_base + ".gene.trace_gibbs").c_str());

    ofs.open(tgi_file);
    {
      int t=0;
      for(vector< vector<string> >::iterator idit = identical_transcripts.begin(); idit != identical_transcripts.end(); idit++) {
        if(isfinite(log(mu_trace_identical[t*trace_length])) != 0) {
          for(strit = idit->begin(); strit != idit->end(); strit++) {
            ofs << *strit;
            if(strit->compare(idit->back())!=0) ofs << "+";
          }
          ofs << " ";
        }
        t++;
      }
    }
    ofs << endl;
    for(int i=0; i < trace_length; i++) {
      int t=0;
      for(vector< vector<string> >::iterator idit = identical_transcripts.begin(); idit != identical_transcripts.end(); idit++) {
        if(isfinite(log(mu_trace_identical[t*trace_length])) != 0) {
          ofs << mu_trace_identical[t*trace_length + i] << " ";
        }
        t++;
      }
      ofs << endl;
    }
    ofs.close(); ofs.clear();

    ofs.open(tgg_file);
    {
      int g=0;
      for(map<string, vector<string> >::iterator git = gene2transcripts.begin(); git != gene2transcripts.end(); git++) {
        if(isfinite(log(mu_trace_gene[g*trace_length])) != 0) {
          ofs << git->first  << " ";
        }
        g++;
      }
    }
    ofs << endl;
    for(int i=0; i < trace_length; i++) {
      int g=0;
      for(map<string, vector<string> >::iterator git = gene2transcripts.begin(); git != gene2transcripts.end(); git++) {
        if(isfinite(log(mu_trace_gene[g*trace_length])) != 0) {
          ofs << mu_trace_gene[g*trace_length + i] << " ";
        }
        g++;
      }
      ofs << endl;
    }
    ofs.close(); ofs.clear();
  }

  // get means of logged posterior samples
  vector<double> meanmu(n,0);
  vector<double> meanmu_identical(identical_transcripts.size(), 0);
  vector<double> meanmu_gene(gene2transcripts.size(), 0);

  // log mu traces to make them more normal before calculating standard errors
  for(int i=0; i < trace_length; i++) {
    for(int t=0; t < n; t++) {
      mu_trace[t*trace_length + i] = log(mu_trace[t*trace_length + i] );
      meanmu[t] +=  mu_trace[t*trace_length + i] ;
    }
    for(int t=0; t < identical_transcripts.size() ; t++) {
      mu_trace_identical[t*trace_length + i] = log(mu_trace_identical[t*trace_length + i] );
      meanmu_identical[t] += mu_trace_identical[t*trace_length + i];
    }
    for(int t=0; t < gene2transcripts.size(); t++){
      mu_trace_gene[t*trace_length + i] = log(mu_trace_gene[t*trace_length + i] );
      meanmu_gene[t] += mu_trace_gene[t*trace_length + i];
    }
  }

  for(int t=0; t < n; t++) meanmu[t] /= trace_length;
  for(int i=0; i < identical_transcripts.size() ; i++) meanmu_identical[i] /= trace_length; 
  for(int i=0; i < gene2transcripts.size(); i++) meanmu_gene[i]/= trace_length;

  // now gibbs means and traces are on log scale.

  // calculate mumcse and iact

  double mumcse[n];
  double iact[n];
  for(int t=0; t < n; t++) {
    double var=0;
    double tau=0;
    int m=0;
    int len=trace_length;
    if(sokal(&len,&mu_trace[t*trace_length],&var,&tau,&m)!=0) {
      mumcse[t] = trace_length;
      iact[t] = NAN;
    } else {
      mumcse[t] = sqrt(tau*var/trace_length);
      iact[t] = tau;
    }
  }
  delete [] mu_trace;

  double mumcse_identical[identical_transcripts.size()];
  double iact_identical[identical_transcripts.size()];
  for(int t=0; t < identical_transcripts.size(); t++) {
    double var=0;
    double tau=0;
    int m=0;
    int len=trace_length;
    if(sokal(&len,&mu_trace_identical[t*trace_length],&var,&tau,&m)!=0) {
      mumcse_identical[t] = trace_length;
      iact_identical[t] = NAN;
    } else {
      mumcse_identical[t] = sqrt(tau*var/trace_length);
      iact_identical[t] = tau;
    }
  }
  delete [] mu_trace_identical;

  double mumcse_gene[gene2transcripts.size()];
  double iact_gene[gene2transcripts.size()];
  for(int g=0; g < gene2transcripts.size(); g++) {
    double var=0;
    double tau=0;
    int m=0;
    int len=trace_length;
    if(sokal(&len,&mu_trace_gene[g*trace_length],&var,&tau,&m)!=0) {
      mumcse_gene[g] = trace_length;
      iact_gene[g] = NAN;
    } else {
      mumcse_gene[g] = sqrt(tau*var/(double)trace_length);
      iact_gene[g] = tau;
    }
  }
  delete [] mu_trace_gene;

  // Calc gene-level weighted sum of transcript lengths
  vector<double> gene_lengths(gene2transcripts.size(), 0.0);
  {
    int g=0;
    for(map<string, vector<string> >::iterator git = gene2transcripts.begin(); git != gene2transcripts.end(); git++) {
      if(isfinite(meanmu_gene[g]) != 0) {
        double sum=0;
        for(strit = (git->second).begin(); strit != (git->second).end(); strit++) {
          if(sidIndex.find(*strit) != sidIndex.end()) {
            gene_lengths[g] += sidLen[*strit] * exp(meanmu[sidIndex[*strit]]);
            sum += exp(meanmu[sidIndex[*strit]]);
          }
        }
        gene_lengths[g] /= sum;
      }
      g++;
    }
  }
  
  ofs.open((output_base + ".mmseq").c_str());
  ofs << "# Mapped reads (or read pairs): " << numbermappedreads << endl;
  ofs << "transcript_id\tlog_mu_em\tlog_mu_gibbs\tmcse\tlength\n";
  for(strit = transcriptList.begin(); strit != transcriptList.end(); strit++) {
    if(sidIndex.count(*strit) > 0 ) {
      ofs << *strit << "\t" << log(mu_em[sidIndex[*strit]]) << "\t" << meanmu[sidIndex[*strit]] << "\t" 
          << mumcse[sidIndex[*strit]] << "\t" << sidLen[*strit] << endl;
    } else {
      ofs << *strit << "\t" << "NA" << "\t" << "NA" << "\t"
          << "NA" << "\t" << sidLen[*strit] << endl;
    }
  }
  ofs.close();ofs.clear();
  
  ofs.open((output_base + ".identical.mmseq").c_str());
  ofs << "# Mapped reads (or read pairs): " << numbermappedreads << endl;
  ofs << "identical_transcript_ids\tlog_mu_em\tlog_mu_gibbs\tmcse\n";
  {
    int t=0;
    for(vector< vector<string> >::iterator idit = identical_transcripts.begin();
      idit != identical_transcripts.end(); idit++) {
      if(isfinite(meanmu_identical[t])) {
        for(strit = idit->begin(); strit != idit->end(); strit++) {
          ofs << *strit;
          if(strit->compare(idit->back())!=0) ofs << "+";
          else ofs << "\t" << log(mu_em_identical[t]) << "\t" << meanmu_identical[t] << "\t" << mumcse_identical[t] << endl;
        }
      }
      t++;
    }
  }
  ofs.close();ofs.clear();

  ofs.open((output_base + ".gene.mmseq").c_str());
  ofs << "# Mapped reads (or read pairs): " << numbermappedreads << endl;
  ofs << "gene_id\tlog_mu_em\tlog_mu_gibbs\tmcse\tweighted_length\n";
  {
    int g=0;
    for(map<string, vector<string> >::iterator git = gene2transcripts.begin();
      git != gene2transcripts.end(); git++) {
      if(isfinite(meanmu_gene[g]) != 0) {
        ofs << git->first <<  "\t" <<  log(mu_em_gene[g]) << "\t" << meanmu_gene[g] << "\t" 
            << mumcse_gene[g] << "\t" <<  gene_lengths[g] << endl;
      } else {
        ofs << git->first <<  "\t" <<  "NA" << "\t" << "NA" << "\t" 
            << "NA" << "\t" << "NA" << endl;
      }
      g++;
    }
  }
  ofs.close();ofs.clear();

  cout << "done." << endl;

  cout << "Output files: " << endl
       << "  " << output_base << ".mmseq" << endl
       << "  " << output_base << ".identical.mmseq" << endl
       << "  " << output_base << ".gene.mmseq" << endl;
  if(traces) {
  cout << "  " << output_base << ".trace_gibbs" << endl
       << "  " << output_base << ".identical.trace_gibbs" << endl
       << "  " << output_base << ".gene.trace_gibbs" << endl;
  }
  if(debug) {
    cout << "  " << output_base << ".trace_em" << endl
       << "  " << output_base << ".k" << endl
       << "  " << output_base << ".M" << endl
       << "  " << output_base << ".Mt-nodups" << endl
       << "  " << output_base << ".doublehits" << endl
       << "  " << output_base << ".dupIDs" << endl;
  }

/*
cout << "Outputting information matrix.. ";

  ofs.open("infomat.txt");

  ublas::vector<double> gamma(m);
  for(int i=0; i < m; i++) {
    matrix_row<boolMat > mr (M, i);
    gamma[i] = inner_prod(mr, mu);
  }

  doubleMat J(n,n);

  for(int t1=0; t1 < n; t1++) {
    for(int t2=0; t2 < n; t2++) {
      double sum=0;
      for(int i=0; i < m; i++) {
        if(M(i,t1)==1 && M(i,t2)==1) {
          sum += k[i]/(gamma[i]*gamma[i]);
        }
      }
      J(t1,t2) = sum;
      ofs << J(t1,t2) << " ";
    }
    ofs << endl;
  }
         
  ofs.close();
  cout << "done.\n";
*/
  return 0;
}
