/**
 *    WHAM - high-throughput sequence aligner
 *    Copyright (C) 2011  WHAM Group, University of Wisconsin
 *
 *    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 3 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, see <http://www.gnu.org/licenses/>.
 */

/*	$Id: main.cpp 165 2012-11-26 10:23:16Z yinan $ */

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <assert.h>
#include <time.h>
#include <pthread.h>
#ifndef WIN32
#include <sys/time.h>
#endif
#include <unistd.h>
#include <fstream>

#include "error.h"
#include "aligner.h"
#include "sequence.h"
#include "pair.h"
#include "short.h"
#include "edit_distance.h"
#include "perfcounters.h"
#include "rdtsc.h"
#include "util.h"

using namespace std;

extern int ELOG_LEVEL;

char pgversion[] = "0.1.5";
char * pgcommand;

Aligner * aligner;
char outputpath[MAX_LENGTH_PATH] = "";
int maxerr = 0;
int maxgap = 0;
int nThread;
int k = 1, m = 0;
bool sorted = false, strata = false;
PerfCounters perfcounter;
unsigned long long perfctr[4];
int minins = 0, maxins = 250;
char * alignFileName = NULL, *unalignFileName = NULL;
bool concatenate = true;
AlignRes * threadres;

extern int64 statHashLookup;
extern int64 statHashLookupEntry;
extern int64 statEmbedHashLookup;
extern int64 statEmbedHashLookupEntry;

typedef struct ThreadInfo {
  int id;
  AlignInfo info;
  char outputpath[MAX_LENGTH_PATH];
} ThreadInfo;

void printStatInfos() {
  elog(DEBUG1, "Hash Lookups: %llu\n", statHashLookup);
  elog(DEBUG1, "Hash Lookup Entries: %llu\n", statHashLookupEntry);
  elog(DEBUG1, "Embed Hash Lookups: %llu\n", statEmbedHashLookup);
  elog(DEBUG1, "Embed Hash Lookup Entries: %llu\n", statEmbedHashLookupEntry);
}

/*
 * skip a line in a text file
 */
int skipLine(FILE * file) {
  char c;
  int i;

  i = 0;
  while (1) {
    c = fgetc(file);
    i++;
    if (c == 10 || c == 13 || c == EOF
    )
      break;
  }
  return i;
}

/*
 * split a comma list
 */
char ** commaList(char * str, int & num) {
  int i;
  char * pch;
  char ** tok;

  num = 1;
  pch = strchr(str, ',');
  while (pch != NULL) {
    num++;
    pch = strchr(pch + 1, ',');
  }

  tok = new char *[num];
  i = 0;
  pch = strtok(str, ",");
  while (pch != NULL) {
    tok[i] = new char[strlen(pch) + 1];strcpy(tok[i], pch);
    i++;
    pch = strtok (NULL, ",");
  }

  return tok;
}

/*
 * convert string parameter to integer
 */
int getArguVal(char * str) {
  char * c;
  int base;

  c = strchr(str, 'k');
  if (c == NULL
  )
    c = strchr(str, 'K');
  if (c == NULL
  )
    c = strchr(str, 'm');
  if (c == NULL
  )
    c = strchr(str, 'M');
  if (c == NULL
  )
    c = strchr(str, 'g');
  if (c == NULL
  )
    c = strchr(str, 'G');

  if (c == NULL
  )
    return atoi(str);

  if (*c == 'k' || *c == 'K')
    base = 1024;
  else if (*c == 'm' || *c == 'M')
    base = 1024 * 1024;
  else if (*c = 'g' || *c == 'G')
    base = 1024 * 1024 * 1024;
  else
    base = 1;

  *c = '\0';
  return base * atoi(str);
}

void getVersion(char * buf) {
  ifstream versionfile("VERSION", ifstream::in);

  versionfile.getline(buf, 16);

  versionfile.close();
}

int printversion() {
  printf("WHAM Version %s\n", pgversion);
  return SUCCESS;
}

int printhelp() {
  printf("Usage:\n");
  printf(
      "  wham [options]* {<readfile> | -1 <m1file> -2 <m2file>} <basepath> <output>\n");
  printf(
      "  <readfile>         comma-separated list of files containing unpaired reads\n");
  printf(
      "  <m1file>           comma-separated list of files containing upstream mates\n");
  printf(
      "  <m2file>           comma-separated list of files containing downstream mates\n");
  printf(
      "  <basepath>         write wham data to files with this dir/basename\n");
  printf("  <output>           file to write alignments to\n");
  printf("Input options:\n");
  printf("  -l <int>           use first <int> bases in each read\n");
  printf("Alignment options:\n");
  printf(
      "  -v <int>           specify the max number of errors in a reported alignment.\n");
  printf(
      "  -g/--gap <int>     specify the max number of gaps in a reported alignment.\n");
  printf("  -e/--maqerr <int>  max sum of mismatch quals across alignment\n");
  printf(
      "  --nofw/--norc      do not align to forward/reverse-complement ref strand\n");
  printf(
      "  --nofr/--norf      do not align to mate1/mate2 strand: fw/rev, rev/fw.\n");
//	printf("  --noff/--norr      do not align to mate1/mate2 strand: fw/fw, rev/rev.\n");
  printf(
      "  -I/--minins <int>  minimum insert size for paired-end alignment (default: 0).\n");
  printf(
      "  -X/--maxins <int>  maximum insert size for paired-end alignment (default: 250).\n");
  printf("Reporting options:\n");
  printf(
      "  -k <int>           report up to <int> valid alignemtns per read (default: 1).\n");
  printf("  -a/--all           report all valid alignments per read.\n");
  printf(
      "  --best             reprot valid alignments in a sorted order of quality.\n");
  printf(
      "  -m <int>           discard reads with more than <int> valid alignmetns.\n");
  printf("Output options:\n");
  printf("  -S/--sam           write alignment in SAM format\n");
  printf("  --al <fname>       write aligned reads/pairs to file(s) <fname>\n");
  printf(
      "  --un <fname>       write unaligned reads/pairs to file(s) <fname>\n");
  printf("Performance options:\n");
  printf("  -t <int>           specify the number of threads\n");
  printf(
      "  --nocat            do not concatenate results from various threads\n");
  printf(
      "  --step <int>       specify the number of indexes that fit into memory.\n");
  printf("Other options:\n");
//	printf("  --pipeline         load and lookup hash indexes one by one\n");
  printf("  --version          print version information\n");
  printf("  -h/--help          print this usage message\n");

  return SUCCESS;
}

/*
 * concatenate all command options
 */
char * getCommand(int argc, char * argv[]) {
  int i, len = 0;
  char * str;

  for (i = 0; i < argc; i++)
    len += strlen(argv[i]) + 1;

  str = new char[len];

  for (i = 0; i < argc; i++) {
    if (i > 0)
      strcat(str, " ");
    strcat(str, argv[i]);
  }

  return str;
}

/*
 * merge output files generated by various threads
 */
void mergeFiles(string fname) {
  int i;
  string command;
  char buf[10];

  command = "cat ";
  for (i = 0; i < nThread; i++) {
    sprintf(buf, ".t%d ", i);
    command += fname + buf;
  }
  command += "> " + fname;
  system(command.c_str());

  for (i = 0; i < nThread; i++) {
    sprintf(buf, ".t%d", i);
    command = "rm -f " + fname + buf;
    system(command.c_str());
  }

}

/*
 * merge alignments, aligned reads, unaligned reads generated by various threads
 */
void merge(char * outputpath, char * alignFileName, char * unalignFileName,
    bool paired) {
  string path;

  // merge output files
  if (outputpath[0] != '\0') {
    path.assign(outputpath);
    mergeFiles(path);
  }

  //merge aligned read files
  if (alignFileName != NULL)
  {
    path.assign(alignFileName);
    if (paired) {
      mergeFiles(path + "_1");
      mergeFiles(path + "_2");
    } else {
      mergeFiles(path);
    }
  }

  //merge unaligned read files
  if (unalignFileName != NULL)
  {
    path.assign(unalignFileName);
    if (paired) {
      mergeFiles(path + "_1");
      mergeFiles(path + "_2");
    } else {
      mergeFiles(path);
    }
  }
}

/*
 * thread procedure to align reads
 */
void * alignThreadProc(void * value) {
  AlignRes res;
  ThreadInfo * info = (ThreadInfo *) value;

//	cpu_set_t mask;
//	CPU_ZERO(&mask);
//	CPU_SET(info->id, &mask);

//	if (sched_setaffinity(0, sizeof(cpu_set_t), &mask) == -1)
//		elog(ERROR, "set affinity error\n");

  res = aligner->align(&(info->info), info->outputpath);

  info->info.reader1->flush();
  if (info->info.reader2)
    info->info.reader2->flush();

  threadres[info->id] = res;

  return NULL;
}

/*
 * load indexes, align reads, and output results. Using multithreading if necessary
 */
int align(char * basepath, AlignInfo * info, char * filename, int nThread) {
  int i, j, ret;
  double t;
  pthread_attr_t attr;
  ShortRead ** partitions1, **partitions2;
  pthread_t * threadpool;
  Timer timer;
  AlignRes res;
  ThreadInfo * infos;

  pthread_attr_init(&attr);
  pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE);
  pthread_attr_setscope(&attr, PTHREAD_SCOPE_SYSTEM);

  /* make sure the index can fit into memory */

  elog(INFO, "loading WHAM indexes...\n");

  /* load index header */
  aligner = new Aligner(basepath);

  if (aligner->getReadLength() != info->reader1->getReadLength()) {
    elog(ERROR, "Unmatched read length (index: %d, reads: %d)\n",
        aligner->getReadLength(), info->reader1->getReadLength());
    return ERR_PARA;
  }

  if (maxgap > 0 && aligner->allowGap() == false) {
    elog(ERROR, "The index does not support indel.\n");
    return ERR_PARA;
  }

  aligner->printInfo();

  //!!!this function should be modified
  aligner->setErrorModel(maxerr, maxgap, info->maxQual);

  /* load hash tables */
  timer.start();
  ret = aligner->loadHashtables(basepath);
  if (ret != SUCCESS) {
    if (ret == ERR_MEM)
      elog(ERROR, "No enough memory\n");
    return ret;
  }
  elog(INFO, "loading time: %.0f sec\n", timer.stop());

//	aligner->setScanThreshold(info->scanThreshold);

  /*
   elog(INFO, "testing WHAM indexes...");
   ret = aligner->check(1000);
   if (ret != SUCCESS)
   {
   elog(ERROR, "failed to pass the tests. \n");
   return ret;
   }
   elog(INFO, "pass.\n");
   */

  elog(INFO, "align reads using indexes...\n");
  if (nThread == 1) {
    /* quick path for single thread */
#ifdef PERFCOUNT
    unsigned long long timer1;
    perfcounter.init();

    startTimer(&timer1);
    perfcounter.threadinit();
    perfcounter.writeCounters(&perfctr[0], &perfctr[1]);
#endif

    info->showBar = true;

    timer.start();

    /* single-thread execution */
    res = aligner->align(info, filename);

    t = timer.stop();

#ifdef PERFCOUNT
    perfcounter.writeCounters(&perfctr[2], &perfctr[3]);
    stopTimer(&timer1);
    printf("Cycles: %lld, PerfCnt 1:  %lld, PerfCnt 2: %lld\n", timer1, perfctr[2]-perfctr[0], perfctr[3]-perfctr[1]);
#endif
  } else {
    threadpool = new pthread_t[nThread];
    infos = new ThreadInfo[nThread];
    threadres = new AlignRes[nThread];

    bool pair = (info->reader2 != NULL);

    /* multithreading execution */
    partitions1 = info->reader1->split(nThread);
    if (pair)
      partitions2 = info->reader2->split(nThread);

    for (int i = 0; i < nThread; i++) {
      infos[i].id = i;
      infos[i].info = *info;
      infos[i].info.reader1 = partitions1[i];
      if (pair)
        infos[i].info.reader2 = partitions2[i];
      else
        infos[i].info.reader2 = NULL;

      //only the first thread shows the progress bar
      if (i == 0)
        infos[i].info.showBar = true;
      else
        infos[i].info.showBar = false;

      if (outputpath[0] == '\0')
        infos[i].outputpath[0] = '\0';
      else
        sprintf(infos[i].outputpath, "%s.t%d", outputpath, i);
      assert(
          !pthread_create(&threadpool[i], &attr, alignThreadProc, (void *)&infos[i]));
    }

    pthread_attr_destroy(&attr);

    timer.start();

    for (int i = 0; i < nThread; i++)
      assert(!pthread_join(threadpool[i], NULL));

    t = timer.stop();

    elog(INFO, "\n");
    res.nRead = res.nValidRead = res.nValidAlignment = 0;
    for (int i = 0; i < nThread; i++) {
      res.nValidAlignment += threadres[i].nValidAlignment;
      res.nValidRead += threadres[i].nValidRead;
      res.nRead += threadres[i].nRead;
      elog(INFO, "Thread %d: Valid Reads: %u/%u, Valid Alignments: %d\n", i,
          threadres[i].nValidRead, threadres[i].nRead,
          threadres[i].nValidAlignment);
    }

    if (concatenate) {
      elog(INFO, "collect results...\n");
      merge(filename, alignFileName, unalignFileName, pair);
    }

    delete[] infos;
    delete[] threadpool;
    delete[] threadres;
  }

  /* remove hash tables */
  ret = aligner->removeHashTables();
  if (ret != SUCCESS
    )
    return ret;

  PairAligner::printTimePairAlign();

  elog(INFO, "\n***************************\n");
  elog(INFO, "Total Align Time: %.6f sec\n", t);

  float ratio = (float) res.nValidRead * 100 / res.nRead;
  elog(INFO, "Valid Reads: %u/%u=%.2f%%\n", res.nValidRead, res.nRead, ratio);
  elog(INFO, "Valid Alignments: %u\n", res.nValidAlignment);
  elog(INFO, "***************************\n\n");

  aligner->printStat();

  return SUCCESS;
}

int main(int argc, char* argv[]) {
  int i, j;
  int numKeys;
  char ** m1_files = NULL, **m2_files = NULL;
  int m1_nfile, m2_nfile;
  int nSeq = 1;
  bool m1_forward = true, m1_backward = true;
  bool m2_forward = true, m2_backward = true;
  bool paired = false;
  int outputMode = MODE_NORMAL;
  int ret;
  CompactSequence * sequence;
  char basepath[256] = "";
  pthread_t * threadpool;
  ShortRead * reader1 = NULL, *reader2 = NULL;
  int * threadid;
  int maxlen = 0;
  int maxqual = 255;
  int maxMate = 100;
  double scanThreshold = 0.001;
  bool mateMatch = true;
  bool pairStrand[2][2];

  srand((unsigned int) time(NULL));
  nThread = 1;

  pairStrand[FORWARD][FORWARD] = false;
  pairStrand[FORWARD][BACKWARD] = true;
  pairStrand[BACKWARD][FORWARD] = true;
  pairStrand[BACKWARD][BACKWARD] = false;

  j = 0;
  for (i = 1; i < argc; i++) {
    if (strcmp(argv[i], "-l") == 0)
      maxlen = atoi(argv[++i]);
    else if (strcmp(argv[i], "-k") == 0)
      k = atoi(argv[++i]);
    else if (strcmp(argv[i], "-m") == 0)
      m = atoi(argv[++i]);
    else if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--all") == 0)
      m = k = 0;
    else if (strcmp(argv[i], "-e") == 0 || strcmp(argv[i], "--maqerr") == 0)
      maxqual = atoi(argv[++i]);
    else if (strcmp(argv[i], "-v") == 0)
      maxerr = atoi(argv[++i]);
    else if (strcmp(argv[i], "-g") == 0 || strcmp(argv[i], "--gap") == 0)
      maxgap = atoi(argv[++i]);
    else if (strcmp(argv[i], "-s") == 0)
      scanThreshold = atof(argv[++i]);
    else if (strcmp(argv[i], "-t") == 0)
      nThread = atoi(argv[++i]);
    else if (strcmp(argv[i], "--best") == 0)
      sorted = true;
    else if (strcmp(argv[i], "--hit") == 0)
      maxMate = atoi(argv[++i]);
    else if (strcmp(argv[i], "--al") == 0) {
      alignFileName = new char[256];
      strcpy(alignFileName, argv[++i]);
    } else if (strcmp(argv[i], "--un") == 0) {
      unalignFileName = new char[256];
      strcpy(unalignFileName, argv[++i]);
    } else if (strcmp(argv[i], "--nofw") == 0)
      m1_forward = false;
    else if (strcmp(argv[i], "--norc") == 0)
      m1_backward = false;
    else if (strcmp(argv[i], "--nocat") == 0)
      concatenate = false;
    else if (strcmp(argv[i], "-S") == 0 || strcmp(argv[i], "--sam") == 0)
      outputMode = MODE_SAM;
    else if (strcmp(argv[i], "--nofr") == 0) {
      pairStrand[FORWARD][BACKWARD] = false;
    } else if (strcmp(argv[i], "--norf") == 0) {
      pairStrand[BACKWARD][FORWARD] = false;
    } else if (strcmp(argv[i], "--noff") == 0) {
      pairStrand[FORWARD][FORWARD] = false;
    } else if (strcmp(argv[i], "--norr") == 0) {
      pairStrand[BACKWARD][BACKWARD] = false;
    } else if (strcmp(argv[i], "--nomate") == 0) {
      mateMatch = false;
    } else if (strcmp(argv[i], "--info") == 0) {
      i++;
      if (strcmp(argv[i], "ERROR") == 0)
        ELOG_LEVEL = ERROR;
      else if (strcmp(argv[i], "WARNING") == 0)
        ELOG_LEVEL = WARNING;
      else if (strcmp(argv[i], "INFO") == 0)
        ELOG_LEVEL = INFO;
      else if (strcmp(argv[i], "DEBUG1") == 0)
        ELOG_LEVEL = DEBUG1;
    } else if (strcmp(argv[i], "-I") == 0 || strcmp(argv[i], "--minins") == 0)
      minins = atoi(argv[++i]);
    else if (strcmp(argv[i], "-X") == 0 || strcmp(argv[i], "--maxins") == 0)
      maxins = atoi(argv[++i]);
    else if (strcmp(argv[i], "--version") == 0) {
      printversion();
      return SUCCESS;
    } else if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) {
      printhelp();
      return SUCCESS;
    } else if (strcmp(argv[i], "-1") == 0) {
      j = 1;
      paired = true;
      m1_files = commaList(argv[++i], m1_nfile);
      if (strcmp(argv[++i], "-2") == 0)
        m2_files = commaList(argv[++i], m2_nfile);
      else {
        elog(ERROR, "Specify files for mate2 of paired-end reads\n");
        printf("See usage message by specifying -h/--help.\n");
        return ERR_PARA;
      }
    } else {
      if (argv[i][0] == '-') {
        printf("Invalid option %s.\n", argv[i]);
        printf("See usage message by specifying -h/--help.\n");
        return ERR_PARA;
      }

      if (j == 0) {
        paired = false;
        m1_files = commaList(argv[i], m1_nfile);
      } else if (j == 1)
        strcpy(basepath, argv[i]);
      else if (j == 2)
        strcpy(outputpath, argv[i]);
      else {
        printf("Invalid option %s.\n", argv[i]);
        printf("See usage message by specifying -h/--help.\n");
        return ERR_PARA;
      }
      j++;
    }
  }

  if (basepath[0] == '\0') {
    elog(ERROR, "specify the base name.\n");
    printf("See usage message by specifying -h/--help.\n");
    exit(1);
  }

//	if (outputpath[0] == '\0') {
//		elog(ERROR, "specify the output file name.\n");
//		printf("See usage message by specifying -h/--help.\n");
//		exit(1);
//	}

  if (m1_files == NULL)
  {
    elog(ERROR, "specify the read files.\n");
    printf("See usage message by specifying -h/--help.\n");
    return ERR_PARA;
  }

  if (maxgap > MAX_GAP)
  {
    elog(ERROR, "WHAM supports up to %d gaps.\n", MAX_GAP);
    return ERR_PARA;
  }

  if (!pairStrand[FORWARD][FORWARD] && !pairStrand[FORWARD][BACKWARD])
    m1_forward = false;
  if (!pairStrand[FORWARD][FORWARD] && !pairStrand[BACKWARD][FORWARD])
    m2_forward = false;
  if (!pairStrand[BACKWARD][FORWARD] && !pairStrand[BACKWARD][BACKWARD])
    m1_backward = false;
  if (!pairStrand[FORWARD][BACKWARD] && !pairStrand[BACKWARD][BACKWARD])
    m2_backward = false;

  /* catenate the command line */
  pgcommand = getCommand(argc, argv);

  elog(INFO, "loading short reads...\n");

  if (!paired) {
    /*	single-end read */
    reader1 = new ShortRead();
    reader1->init(m1_files, m1_nfile, m1_forward, m1_backward, alignFileName,
        unalignFileName);

    ret = reader1->load(maxlen);
    if (ret != SUCCESS)
    {
      elog(ERROR, "failed to load the short reads.\n");
      return ret;
    }
  } else {
    char * alFileName1 = NULL, *alFileName2 = NULL;
    char * unFileName1 = NULL, *unFileName2 = NULL;

    if (alignFileName != NULL)
    {
      alFileName1 = new char[256];
      alFileName2 = new char[256];
      sprintf(alFileName1, "%s_1", alignFileName);
      sprintf(alFileName2, "%s_2", alignFileName);
    }
    if (unalignFileName != NULL)
    {
      unFileName1 = new char[256];
      unFileName2 = new char[256];
      sprintf(unFileName1, "%s_1", unalignFileName);
      sprintf(unFileName2, "%s_2", unalignFileName);
    }

    /*	paired-end read */
    reader1 = new ShortRead();
    reader1->init(m1_files, m1_nfile, m1_forward, m1_backward, alFileName1,
        unFileName1);

    ret = reader1->load(maxlen);
    if (ret != SUCCESS)
    {
      elog(ERROR, "failed to load the short reads.\n");
      return ret;
    }

    reader2 = new ShortRead();
    reader2->init(m2_files, m2_nfile, m2_forward, m2_backward, alFileName2,
        unFileName2);

    ret = reader2->load(maxlen);
    if (ret != SUCCESS)
    {
      elog(ERROR, "failed to load the short reads.\n");
      return ret;
    }

    /* check mate1 and mate2 */
    if (reader1->getNumReads() != reader2->getNumReads()) {
      elog(ERROR, "the numbers of reads in mate files does not match.\n");
      return ret;
    }

    if (reader1->getReadLength() != reader2->getReadLength()) {
      elog(ERROR, "the lengths of reads in mate files does not match\n");
      return ret;
    }
  }

  if (reader1->getNumReads() == 0) {
    elog(INFO, "Empty read files.\n");
    return SUCCESS;
  }

  AlignInfo info;
  info.reader1 = reader1;
  info.reader2 = reader2;
  info.minins = minins;
  info.maxins = maxins;
  info.sorted = sorted;
  info.strata = strata;
  info.maxHit = k;
  info.maxMatch = m;
  info.maxQual = maxqual;
  info.maxGap = maxgap;
  info.scanThreshold = scanThreshold;
  info.mateMatch = mateMatch;
  info.maxMate = maxMate;
  info.outputFormat = outputMode;
  info.concatenate = concatenate;
  memcpy(info.pairStrand, pairStrand, 4 * sizeof(bool));

  ret = align(basepath, &info, outputpath, nThread);
  if (ret != SUCCESS)
  {
    elog(ERROR, "failed to align read files.\n");
    return ret;
  }

  if (outputpath[0] != '\0')
    elog(INFO, "see %s for all valid alignments.\n", outputpath);

  printStatInfos();

  return SUCCESS;
}