/**
 *    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: short.cpp 157 2012-07-25 05:58:09Z yinan $ */

#include <string.h>
#include <stdlib.h>
#include <assert.h>
#include <math.h>
#include "lib.h"
#include "error.h"
#include "short.h"

ShortRead::ShortRead() {
  memset(this, 0, sizeof(ShortRead));
}

ShortRead::ShortRead(ShortRead * read) {
  memcpy(this, read, sizeof(ShortRead));
}

ShortRead::~ShortRead() {
}

/*
 *	ShortRead::readLine
 *	read a line from the specified file.
 */
int ShortRead::readLine(FILE * file, char * str, int maxLength) {
  char c;
  int num = 0;

  if (file == NULL
    )
    return 0;

  while (1) {
    c = fgetc(file);
    if (c == EOF || c == '\n' || num == maxLength - 1)
      break;

    str[num++] = c;
  }
  str[num] = '\0';

  return num;
}

/*
 *	ShortRead::init
 *	store the names and the number of files that contain 
 *	the short reads.
 */
int ShortRead::init(char ** files, int numFile, bool fw, bool bw,
    char * alignFileName, char * unalignFileName) {
  int i;

  if (numFile <= 0)
    return ERR_PARA;

  forward = fw;
  backward = bw;
  nReadPerAlign = 0;
  if (forward) {
    sidx[FORWARD] = nReadPerAlign;
    nReadPerAlign++;
  }
  if (backward) {
    sidx[BACKWARD] = nReadPerAlign;
    nReadPerAlign++;
  }

  if (nReadPerAlign <= 0)
    return ERR_PARA;

  length = 0;
  nFile = numFile;
  fnames = new char *[numFile];
  for (i = 0; i < numFile; i++) {
    fnames[i] = new char[256];

    if (strlen(files[i]) >= 256)
      return ERR_PARA;

    strcpy(fnames[i], files[i]);
  }

  nRead = 0;

  storeName = true;
  storeQual = true;
  lenName = 0;
  lenQual = 0;

  if (alignFileName != NULL)
  {
    strcpy(alFileName, alignFileName);

    alfile = fopen(alignFileName, "w");
    if (alfile == NULL)
    {
      return ERR_FILE;
    }
  }

  if (unalignFileName != NULL)
  {
    strcpy(unFileName, unalignFileName);

    unfile = fopen(unalignFileName, "w");
    if (unfile == NULL)
    {
      return ERR_FILE;
    }
  }

  code2Base[BASE_A] = 'A';
  code2Base[BASE_C] = 'C';
  code2Base[BASE_G] = 'G';
  code2Base[BASE_T] = 'T';
  code2Base[BASE_N] = 'N';

  return SUCCESS;
}

/*
 *	ShortRead::allocate()
 *	allocate space for query sequence, name, and qual.
 */
int ShortRead::allocate() {
  /* allocate the space for the short reads */
  reads = (int64 *) malloc(
      (int64) nRead * nReadPerAlign * WORDS_PER_READ * sizeof(int64));
  if (reads == NULL)
  {
    elog(
        ERROR,
        "failed to allocate space for compressed reads. [%lldMB]\n",
        (int64) nRead * nReadPerAlign * WORDS_PER_READ * sizeof(int64) / 1024
            / 1024);
    return ERR_MEM;
  }

  /* allocate the qname space */
  if (storeName) {
    names = (char *) malloc((int64) nRead * lenName * sizeof(char));
    if (names == NULL)
    {
      elog(ERROR, "failed to allocate space for read names.[%lldMB]\n",
          (int64) nRead * lenName * sizeof(char) / 1024 / 1024);
      return ERR_MEM;
    }
  }

  /* allocate the qual space */
  if (storeQual) {
    quals = (char *) malloc((int64) nRead * lenQual * sizeof(char));
    if (quals == NULL)
    {
      elog(ERROR, "failed to allocate space for qual scores. [%lldMB]\n",
          (int64) nRead * lenQual * sizeof(char) / 1024 / 1024);
      return ERR_MEM;
    }
  }

  return SUCCESS;
}

/*
 *	ShortRead::preProcess
 *	count the number of short reads in the specified files.
 */
int ShortRead::preProcess(int maxLen) {
  int i;
  int len;
  FILE * file;
  char str[256];

  if (fnames == 0) {
    elog(ERROR, "no specified file name of short reads.\n");
    return ERR_PARA;
  }

  if (nReadPerAlign <= 0)
    return ERR_PARA;

  nRead = 0;
  for (i = 0; i < nFile; i++) {
    file = fopen(fnames[i], "r");
    if (file == NULL)
    {
      elog(ERROR, "failed to open file %s.\n", fnames[i]);
      return ERR_FILE;
    }

    int lineid = 0;
    while (len = readLine(file, str, 256)) {
      lineid++;

      if (str[0] == '\0')
        continue;
      switch (lineid % 4) {
      case 1:
        if (str[0] == '@') {
          /* HERE CAN BE IMPROVED */
          /* remove first '@', add '\0' in the end */
          len = getReadNameLength(str);
          if (len > lenName)
            lenName = len;
        } else {
          elog(ERROR, "illegal character at line %d of file %s.\n", lineid,
              fnames[i]);
          return ERR_READ_FORMAT;
        }
        break;
      case 2:
        /* initialize the read length */
        if (length == 0) {
          length = len;
          lenQual = len + 1;
        }

        if (length != len) {
          elog(ERROR, "Read at line %d of file %s has illegal read length.\n",
              lineid, fnames[i]);
          return ERR_READ_FORMAT;
        }

        if (maxLen > len) {
          elog(
              ERROR,
              "Read at line %d of file %s is shorter than the specified length.\n",
              lineid, fnames[i]);
          return ERR_PARA;
        }
        break;
      case 3:
        if (str[0] == '+') {
          //				readLine(file, str, 256);
        }
        break;
      case 0:
        //query sequence
        nRead++;

        /* initialize the read length */
        if (length == 0) {
          length = len;
          lenQual = len + 1;
        }

        /* the read file contain shorts with varied length */
        if (length != len) {
          elog(ERROR, "Read at line %d of file %s has illegal read length.\n",
              lineid, fnames[i]);
          return ERR_READ_FORMAT;
        }

        if (maxLen > len) {
          elog(
              ERROR,
              "Read at line %d of file %s is shorter than the specified length.\n",
              lineid, fnames[i]);
          return ERR_PARA;
        }
        break;
      }
    }

    if (fclose(file) != 0) {
      elog(ERROR, "failed to close file %s.\n", fnames[i]);
      return ERR_FILE;
    }
  }

  if (maxLen > 0) {
    length = maxLen;
    lenQual = maxLen + 1;
  }

  return SUCCESS;
}

/*
 *	ShortRead::load
 *	This function is used to load the short reads in the specified 
 *	files. The short reads are loaded into a 64-bit integer array, 
 *	using 4 64-bit integer to represent a short read. The short read 
 *	can be loaded in forward or backward order.
 */
int ShortRead::load(int maxLen) {
  int i, j, k, curRead;
  int ret;
  int forward_offset, backward_offset, offset;
  int64 forward_word, backward_word, code;
  int64 reverse[8] = { 3, 2, 1, 0, 4, 5, 6, 7 };
  char str[256];
  FILE * file;

  /* 	statistics */
  ret = preProcess(maxLen);
  if (ret != SUCCESS
    )
    return ret;

  /*	allocate the space for reads. */
  ret = allocate();
  if (ret != SUCCESS)
  {
    elog(ERROR, "failed to allocate space for reads.\n");
    return ret;
  }

  i = 0;
  curRead = 0;
  offset = (WORDS_PER_READ * BITS_PER_LONGWORD - length * BITS_PER_BASE)
      / BITS_PER_LONGWORD;

  for (k = 0; k < nFile; k++) {
    file = fopen(fnames[k], "r");
    if (file == NULL
      )
      return ERR_FILE;

    int lineid = 0;
    while (readLine(file, str, 256)) {
      lineid++;
      if (str[0] == '\0')
        continue;
      if (str[0] == '+') {
        /* read the score line */
        assert(curRead < nRead);

        readLine(file, str, 256);
        lineid++;

        if (storeQual) {
          strncpy(&quals[curRead * lenQual], str, length);
          quals[curRead * lenQual + length] = '\0';
        }

        curRead++;
      } else if (str[0] == '@') {
        /* copy the query sequence name */
        assert(curRead < nRead);
        if (storeName)
          extractReadName(&names[curRead * lenName], &str[1]);
        //	strcpy(&names[curRead * lenName], &str[1]);

        /* store the query sequence using 3 bits to represent each base */
        readLine(file, str, 256);
        lineid++;

        /* cut the sequence if necessary */
        str[length] = '\0';

        /* initialize the values */
        reads[i] = 0;
        reads[i + 1] = 0;
        reads[i + 2] = 0;
        reads[i + 3] = 0;
        if (nReadPerAlign > 1) {
          reads[i + 6] = 0;
          reads[i + 7] = 0;
          reads[i + 8] = 0;
          reads[i + 9] = 0;
        }

        /* initialize the current word in forward/backward format */
        forward_word = 0;
        backward_word = 0;

        /* initialize the begining offset in forward/backward format */
        backward_offset = 0;
        forward_offset = (WORDS_PER_READ * BITS_PER_LONGWORD
            - length * BITS_PER_BASE) % BITS_PER_LONGWORD;

        /*
         *	scan the sequence and generate the compact representation
         *	in forward or/and backward format.
         */
        for (j = 0; j < length; j++) {
          if (str[j] == 'A')
            code = BASE_A;
          else if (str[j] == 'C')
            code = BASE_C;
          else if (str[j] == 'G')
            code = BASE_G;
          else if (str[j] == 'T')
            code = BASE_T;
          else if (str[j] == 'N')
            code = BASE_N;
          else
            elog(ERROR, "ERROR: unknown character in short read files.\n");

          if (forward) {
            /*	forward format */
            if (forward_offset + BITS_PER_BASE >= BITS_PER_LONGWORD)
            {
              /* on the boundary of 64-bit word */
              reads[i + offset + j * BITS_PER_BASE / BITS_PER_LONGWORD] =
                  (forward_word << (BITS_PER_LONGWORD - forward_offset))
                      | (code
                          >> (forward_offset + BITS_PER_BASE - BITS_PER_LONGWORD));
              forward_offset = forward_offset + BITS_PER_BASE
                  - BITS_PER_LONGWORD;
              forward_word = ~((~code) | (-1LL << forward_offset));
            } else {
              forward_word = (forward_word << BITS_PER_BASE) | code;
              forward_offset += BITS_PER_BASE;
            }
          }

          if (backward) {
            /* backward format */
            if (backward_offset + BITS_PER_BASE >= BITS_PER_LONGWORD)
            {
              /* on the boundary of 64-bit word */
              reads[i + nReadPerAlign * WORDS_PER_READ - 1
                  - j * BITS_PER_BASE / BITS_PER_LONGWORD] = backward_word
                  | (reverse[code] << backward_offset);
              backward_offset = backward_offset + BITS_PER_BASE
                  - BITS_PER_LONGWORD;
              backward_word = reverse[code]
                  >> (BITS_PER_BASE - backward_offset);
            } else {
              backward_word = backward_word
                  | (reverse[code] << (backward_offset % BITS_PER_LONGWORD));
              backward_offset += BITS_PER_BASE;
            }
          }
        }
        if (backward)
          reads[i + nReadPerAlign * WORDS_PER_READ - 1
              - j * BITS_PER_BASE / BITS_PER_LONGWORD] = backward_word;

        //		outputSegment(reads + i, length);
        //		outputSegment(reads + i + WORDS_PER_READ, length);

        /* jump to the begining 64-bit integer of next short read */
        i += WORDS_PER_READ * nReadPerAlign;
      } else {
        elog(ERROR, "illegal character at line %d of file %s\n", lineid,
            fnames[k]);
        return ERR_READ_FORMAT;
      }
    }

    if (fclose(file) != 0)
      return ERR_FILE;
  }

  return SUCCESS;
}

ShortRead ** ShortRead::split(int nPartition) {
  int i;
  int szPartition, curPartition;
  unsigned int curRead = 0;
  ShortRead ** partitions;

  szPartition = (int) ceil((double) nRead / (double) nPartition);
  partitions = new ShortRead *[nPartition];

  for (i = 0; i < nPartition; i++) {
    partitions[i] = new ShortRead(this);

    if (szPartition <= nRead - curRead)
      curPartition = szPartition;
    else
      curPartition = nRead - curRead;
    partitions[i]->nRead = curPartition;

    partitions[i]->reads = &reads[curRead * WORDS_PER_READ * nReadPerAlign];

    if (storeName && names != NULL)
    {
      partitions[i]->names = &names[curRead * lenName];
    }

    if (storeQual && quals != NULL)
    {
      partitions[i]->quals = &quals[curRead * lenQual];
    }

    if (alfile != NULL)
    {
      sprintf(partitions[i]->alFileName, "%s.t%d", alFileName, i);
      partitions[i]->alfile = fopen(partitions[i]->alFileName, "w");
      assert(partitions[i]->alfile != NULL);
    }

    if (unfile != NULL)
    {
      sprintf(partitions[i]->unFileName, "%s.t%d", unFileName, i);
      partitions[i]->unfile = fopen(partitions[i]->unFileName, "w");
      assert(partitions[i]->unfile != NULL);
    }

    curRead += curPartition;
  }

  return partitions;
}

void ShortRead::printRead(int id, FILE * file) {
  char readstr[256];
  assert(file != NULL);

  CompactSequence::decompose(readstr, length, getRead(id));

  fprintf(file, "@%s\n", getReadName(id));
  fprintf(file, "%s\n", readstr);
  fprintf(file, "+%s\n", getReadName(id));
  fprintf(file, "%s\n", getQual(id));
}