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

#include "writer.h"

extern char pgversion[];
extern char * pgcommand;

/*
 * A = 0, C = 1, G = 2, T = 2, N = 7.
 * If the value of A, C, G, T is changed, modify this array.
 */
const char code2Gene[8] = { 'A', 'C', 'G', 'T', 'N', 'N', 'N', 'N' };

void SimpleWriter::writeValidAlignment(int readId, HitSet * set) {
  int h;
  uint32 sid = 0, soffset = 0;
  char qseq[128], rseq[128];
  char field[128];
  char * refName;
  int64 * reference, *query;
  uint32 offset;
  char s;

  for (h = 0; h < set->getNumHits(); h++) {
//		query = set->getQuerySeq(h);
    query = reader1->getRead(readId, set->getStrand(h));
    reference = set->getReferenceSeq(h);
    offset = set->getOffset(h);

    sequence->itree->lookup(offset, sid, soffset);

    //	sprintf(str, "%16u ", offset);
    //	shortread = str + 17;
    //	change = str + 17 + length;

    refName = sequence->getSeqName(sid);
    //	sprintf(str, "chr%u %10u  ", sid, soffset);

    if (set->getStrand(h) == FORWARD
      )
      s = '+';
    else
      s = '-';

    CompactSequence::decompose(qseq, length, query);
    CompactSequence::decompose(rseq, set->getErrorVector(h).len, reference);
//		uncompressSequence(query, length, shortread);

    fprintf(file, "%c\t%s\t%d\t%s\t", s, refName, soffset, qseq);

    field[0] = '\0';
    if (set->getNumMismatch(h) > 0)
      writeField(field, qseq, rseq, set->getStrand(h), set->getErrorVector(h));
//			writeField(query, reference, set->getStrand(h));

    fprintf(file, "%s\n", field);
  }

  return;
}

void SimpleWriter::writeField(char * str, char * query, char * reference,
    strand s, ErrorVector error) {
  int i, j, e;
  int type;
  int len;

  len = strlen(query);
  if (s == FORWARD)
  {
    e = 0;
    for (i = 0, j = 0; i < len; i++, j++) {
      if (query[i] == reference[j])
        continue;
      type = GET_ERROR(error.vec, e);
      switch (type) {
      case ERROR_VECTOR_MIS:
        sprintf(str, "%2d:%c>%c,", i, reference[j], query[i]);
        str += 7;
        break;
      case ERROR_VECTOR_INS:
        sprintf(str, "%2d:_>%c,", i, query[i]);
        str += 7;
        j--;
        break;
      case ERROR_VECTOR_DEL:
        sprintf(str, "%2d:%c>_,", i, reference[j]);
        str += 7;
        i--;
        break;
      default:
        break;
      }
      e++;
    }
  } else {
    e = error.num - 1;
    for (i = len - 1, j = error.len - 1; i >= 0; i--, j--) {
      if (query[i] == reference[j])
        continue;
      type = GET_ERROR(error.vec, e);
      switch (type) {
      case ERROR_VECTOR_MIS:
        sprintf(str, "%2d:%c>%c,", len - 1 - i, reference[j], query[i]);
        str += 7;
        break;
      case ERROR_VECTOR_INS:
        sprintf(str, "%2d:_>%c,", len - 1 - i, query[i]);
        str += 7;
        j++;
        break;
      case ERROR_VECTOR_DEL:
        sprintf(str, "%2d:%c>_,", len - 1 - i, reference[j]);
        str += 7;
        i++;
        break;
      default:
        break;
      }
      e--;
    }
  }
  //remove the last comma
  str[strlen(str) - 1] = '\0';
}

void SimpleWriter::writeField(int64 * query, int64 * reference, strand s) {
  int i, j = 0, k;
  int64 code1, code2;
  uint32 sid = 0, soffset = 0;
  bool firstMismatch = true;

  if (s == FORWARD)
  {
    k = (4 * BITS_PER_LONGWORD - length * BITS_PER_BASE) / BITS_PER_LONGWORD;
    j = BITS_PER_LONGWORD - (length * BITS_PER_BASE) & BITS_LONGWORD_MASK;
    for (i = 0; i < length; i++) {
      if (j + BITS_PER_BASE > BITS_PER_LONGWORD)
      {
        code1 = ((query[k] << (j + BITS_PER_BASE - BITS_PER_LONGWORD))
            | (query[k + 1]
                >> (BITS_PER_LONGWORD + BITS_PER_LONGWORD - j - BITS_PER_BASE)))
            & 0x7;
        code2 = ((reference[k] << (j + BITS_PER_BASE - BITS_PER_LONGWORD))
            | (reference[k + 1]
                >> (BITS_PER_LONGWORD + BITS_PER_LONGWORD - j - BITS_PER_BASE)))
            & 0x7;

        j = j + BITS_PER_BASE - BITS_PER_LONGWORD;
        k++;
      } else {
        code1 = (query[k] >> (BITS_PER_LONGWORD - BITS_PER_BASE - j)) & 0x7;
        code2 = (reference[k] >> (BITS_PER_LONGWORD - BITS_PER_BASE - j)) & 0x7;

        j += BITS_PER_BASE;
      }

      //			shortread[i] = code2Gene[code1];
      if (code1 != code2) {
        if (!firstMismatch)
          fprintf(file, ",");
        firstMismatch = false;
        fprintf(file, "%2d:%c>%c", i, code2Gene[code2], code2Gene[code1]);
      }
    }
  } else {
    k = 3;
    j = 0;
    for (i = length - 1; i >= 0; i--) {
      if (j + BITS_PER_BASE > BITS_PER_LONGWORD)
      {
        if (j < BITS_PER_LONGWORD)
        {
          code1 = ((query[k] >> j) | (query[k - 1] << (BITS_PER_LONGWORD - j)))
              & 0x7;
          code2 = ((reference[k] >> j)
              | (reference[k - 1] << (BITS_PER_LONGWORD - j))) & 0x7;
        } else {
          code1 = query[k - 1] & 0x7;
          code2 = reference[k - 1] & 0x7;
        }

        j = j + BITS_PER_BASE - BITS_PER_LONGWORD;
        k--;
      } else {
        code1 = (query[k] >> j) & 0x7;
        code2 = (reference[k] >> j) & 0x7;

        j += BITS_PER_BASE;
      }

      //			shortread[i] = code2Gene[code1];
      if (code1 != code2) {
        if (!firstMismatch)
          fprintf(file, ",");
        firstMismatch = false;
        fprintf(file, "%2d:%c>%c", length - 1 - i, code2Gene[code2],
            code2Gene[code1]);
      }
    }
  }
}

void SimplePairWriter::writeValidAlignment(int readId, HitSet * set) {
  int h;
  uint32 sid = 0, soffset = 0;
  char shortread[128];
  char * refName;
  int64 * reference, *query;
  uint32 offset;
  char s;

  if (!set->isProperMatch())
    return;

  for (h = 0; h < set->getNumHits(); h++) {
    /* currently skip all partial match */
    if (set->getNumMismatch(h) < 0 || set->getNumMismatch(h + 1) < 0)
      continue;

    if (h % 2 == 0)
      query = reader1->getRead(readId, set->getStrand(h));
    else
      query = reader2->getRead(readId, set->getStrand(h));
    reference = set->getReferenceSeq(h);
    offset = set->getOffset(h);

    sequence->itree->lookup(offset, sid, soffset);

    refName = sequence->getSeqName(sid);

    if (set->getStrand(h) == FORWARD
      )
      s = '+';
    else
      s = '-';

    CompactSequence::decompose(shortread, length, query);

    fprintf(file, "%c\t%s\t%d\t%s\t", s, refName, soffset, shortread);

    if (set->getNumMismatch(h) > 0)
      writeField(query, reference, set->getStrand(h));

    fprintf(file, "\n");
  }

  return;
}

void SamWriter::writeValidAlignment(int readId, HitSet * set) {

  int i, j, k, h, lastDiff;
  int64 code1, code2;
  uint32 sid = 0, soffset = 0;
  int64 * reference, *query;
  uint32 offset;
  char * refName, *qual;
  int flag, mapq = 0;
  double p, sump;
  char qseq[128], rseq[128], tag[128];
  char cigar[128];
  char shortread[256];

  //QNAME
  char * queryName = reader1->getReadName(readId);

  for (h = 0; h < set->getNumHits(); h++) {
    //	query = set->getQuerySeq(h);
    query = reader1->getRead(readId, set->getStrand(h));
    reference = set->getReferenceSeq(h);
    offset = set->getOffset(h);

    sequence->itree->lookup(offset, sid, soffset);

    //FLAG
    flag = 0;
    if (set->getStrand(h) == BACKWARD
      )
      flag |= SAM_FLAG_QUERY_STRAND;
    if (h > 0)
      flag |= SAM_FLAG_NOT_PRIMARY;

    //RNAME
    refName = sequence->getSeqName(sid);

    //POS(soffset)
    soffset++;

    //MAPQ
    mapq = set->getQual(h);

    //SEQ
    CompactSequence::decompose(qseq, length, query);
    CompactSequence::decompose(rseq, set->getErrorVector(h).len, reference);

    //CIGAR(length-M)
    getCIGAR(cigar, qseq, rseq, set->getErrorVector(h));

    //MRNM(*)
    //MPOS(0)
    //ISIZE(0)

    //QUAL
    qual = reader1->getQual(readId);

    fprintf(file, "%s\t%d\t%s\t%d\t%d\t%s\t*\t0\t0\t%s\t%s", queryName, flag,
        refName, soffset, mapq, cigar, qseq, qual);

    writeOptionalField(qseq, rseq, set->getErrorVector(h));
  }
}

void SamWriter::writeInvalidAlignment(int readId, HitSet * set) {
  char queryStr[128];
  char * qualStr;
  int flag;
  int64 * query;

  //QNAME
  char * name = reader1->getReadName(readId);

  //FLAG
  flag = SAM_FLAG_UNMAPPED;

  //QUERY
  query = reader1->getRead(readId, FORWARD);
  CompactSequence::decompose(queryStr, length, query);

  //QUAL
  qualStr = reader1->getQual(readId);

  fprintf(file, "%s\t%d\t*\t0\t0\t*\t*\t0\t0\t%s\t%s\n", name, flag, queryStr,
      qualStr);
}

void SamWriter::writeOptionalField(char * query, char * reference,
    ErrorVector error) {
  char mdstr[256];

  //MD field
  getMDfield(mdstr, query, reference, error);
  fprintf(file, "\tMD:Z:%s", mdstr);

  //NM field
  fprintf(file, "\tNM:i:%d\n", error.num);
}

void SamWriter::getCIGAR(char * str, char * query, char * reference,
    ErrorVector error) {
  int i, j, e;
  int lastErrPos;
  int type;

  //quick path for non-indel match
  if (error.gap == 0) {
    sprintf(str, "%dM", (int) strlen(query));
    return;
  }

  str[0] = '\0';
  lastErrPos = 0;
  e = 0;
  for (i = 0, j = 0; i < strlen(query); i++, j++) {
    if (query[i] == reference[j])
      continue;
    type = GET_ERROR(error.vec, e);
    switch (type) {
    case ERROR_VECTOR_MIS:
      //nothing to do
      break;
    case ERROR_VECTOR_INS:
      if (j - lastErrPos > 0)
        sprintf(str, "%dM", j - lastErrPos);
      strcat(str, "1I");
      str += strlen(str);
      j--;
      lastErrPos = j + 1;
      break;
    case ERROR_VECTOR_DEL:
      if (j - lastErrPos > 0)
        sprintf(str, "%dM", j - lastErrPos);
      strcat(str, "1D");
      str += strlen(str);
      i--;
      lastErrPos = j + 1;
      break;
    default:
      break;
    }
    e++;
  }

  if (j > lastErrPos)
    sprintf(str, "%dM", j - lastErrPos);
}

void SamWriter::getMDfield(char * str, char * query, char * reference,
    ErrorVector error) {
  int i, j, e;
  int lastErrPos;
  int type;

  //quick path for exact match
  if (error.num == 0) {
    sprintf(str, "%d", (int) strlen(query));
    return;
  }

  str[0] = '\0';
  lastErrPos = 0;
  e = 0;
  for (i = 0, j = 0; i < strlen(query); i++, j++) {
    if (query[i] == reference[j])
      continue;
    type = GET_ERROR(error.vec, e);
    switch (type) {
    case ERROR_VECTOR_MIS:
      sprintf(str, "%d", j - lastErrPos);
      str += strlen(str);
      str[0] = reference[j];
      str[1] = '\0';
      str++;
      lastErrPos = j + 1;
      break;
    case ERROR_VECTOR_INS:
      //nothing to do
      j--;
      break;
    case ERROR_VECTOR_DEL:
      sprintf(str, "%d", j - lastErrPos);
      str += strlen(str);
      str[0] = '^';
      str[1] = reference[j];
      str[2] = '\0';
      str += 2;
      i--;
      lastErrPos = j + 1;
      break;
    default:
      break;
    }
    e++;
  }

  sprintf(str, "%d", j - lastErrPos);
}

void SamWriter::writeHead() {
  int i;

  /* for pair-end reads, we don't guarantee the coordinate order */
  fprintf(file, "@HD\tVN:1.3\tSO:queryname\n");

  for (i = 0; i < sequence->getNumSeq(); i++) {
    fprintf(file, "@SQ\tSN:%s\tLN:%d\n", sequence->getSeqName(i),
        sequence->getSeqLen(i));
  }

  fprintf(file, "@PG\tID:WHAM\tVN:%s\tCL:\"%s\"\n", pgversion, pgcommand);
}

void RawWriter::writeValidAlignment(int readId, HitSet * set) {
  int nHit, ret;

  /* In RAW formate, we have to output all hits even the number of
   * hits exceeds -m <int>. Otherwise (discard all hits), the merged
   * results may be wrong.
   * we use getNumAllHits() instead of getNumHIts().
   */
  nHit = set->getNumAllHits();
  ret = fwrite(set->getHits(), sizeof(Hit), nHit, file);
  if (ret != nHit) {
    elog(ERROR, "failed to write RAW align data.\n");
  }
}

/**
 * SamPairWriter::writeValidAlignment()
 * write a mapped alignment in SAM format for a pair-end read.
 */
void SamPairWriter::writeValidAlignment(int readId, HitSet * set) {
  int h;
  uint32 sid1 = 0, soffset1 = 0;
  uint32 sid2 = 0, soffset2 = 0;
  int64 * reference1, *reference2;
  int64 * query1, *query2;
  char * qualStr1, *qualStr2;
  uint32 offset1, offset2;
  char * refName1, *refName2;
  char * mateRefName1, *mateRefName2;
  char equalStr[] = "=";
  int flag1, flag2, mapq1 = 0, mapq2 = 0;
  int tlen1 = 0, tlen2 = 0;
  char queryStr1[128], queryStr2[128], rcQualStr1[128], rcQualStr2[128];
  char refStr1[128], refStr2[128];
  char cigar1[16], cigar2[16];
  bool mate1Unmatch, mate2Unmatch, properMatch;

  //QNAME
  char * queryName1 = reader1->getReadName(readId);
  char * queryName2 = reader2->getReadName(readId);

  for (h = 0; h < set->getNumHits(); h += 2) {
    mate1Unmatch = (set->getNumMismatch(h) < 0);
    mate2Unmatch = (set->getNumMismatch(h + 1) < 0);
    properMatch = set->isProperMatch();

//		query1 = set->getQuerySeq(h);
    query1 = reader1->getRead(readId, set->getStrand(h));
    reference1 = set->getReferenceSeq(h);
    offset1 = set->getOffset(h);

//		query2 = set->getQuerySeq(h + 1);
    query2 = reader2->getRead(readId, set->getStrand(h + 1));
    reference2 = set->getReferenceSeq(h + 1);
    offset2 = set->getOffset(h + 1);

    //FLAG
    flag1 = SAM_FLAG_PAIRED;
    if (set->getStrand(h) == BACKWARD
      )
      flag1 |= SAM_FLAG_QUERY_STRAND;
    if (set->getStrand(h + 1) == BACKWARD
      )
      flag1 |= SAM_FLAG_MATE_STRAND;
    flag1 |= SAM_FLAG_FIRST_IN_PAIR;
    if (properMatch)
      flag1 |= SAM_FLAG_MAPPED_PAIRED;
    if (mate1Unmatch)
      flag1 |= SAM_FLAG_UNMAPPED;
    else if (mate2Unmatch)
      flag1 |= SAM_FLAG_MATE_UNMAPPED;
    if (h > 0)
      flag1 |= SAM_FLAG_NOT_PRIMARY;

    flag2 = SAM_FLAG_PAIRED;
    if (set->getStrand(h + 1) == BACKWARD
      )
      flag2 |= SAM_FLAG_QUERY_STRAND;
    if (set->getStrand(h) == BACKWARD
      )
      flag2 |= SAM_FLAG_MATE_STRAND;
    flag2 |= SAM_FLAG_SECOND_IN_PAIR;
    if (properMatch)
      flag2 |= SAM_FLAG_MAPPED_PAIRED;
    if (mate1Unmatch)
      flag2 |= SAM_FLAG_MATE_UNMAPPED;
    else if (mate2Unmatch)
      flag2 |= SAM_FLAG_UNMAPPED;
    if (h > 0)
      flag2 |= SAM_FLAG_NOT_PRIMARY;

    //POS(soffset)
    sequence->itree->lookup(offset1, sid1, soffset1);
    soffset1++; /* 1-based coordination */
    sequence->itree->lookup(offset2, sid2, soffset2);
    soffset2++; /* 1-based coordination */

    //RNAME
    refName1 = sequence->getSeqName(sid1);
    if (sid1 == sid2) {
      refName2 = refName1;
      mateRefName1 = equalStr;
      mateRefName2 = equalStr;
    } else {
      refName2 = sequence->getSeqName(sid2);
      mateRefName1 = refName2;
      mateRefName2 = refName1;
    }

    //MAPQ
    mapq1 = set->getQual(h);
    mapq2 = set->getQual(h + 1);

    //SEQ
    CompactSequence::decompose(queryStr1, length, query1);
    CompactSequence::decompose(queryStr2, length, query2);
    CompactSequence::decompose(refStr1, set->getErrorVector(h).len, reference1);
    CompactSequence::decompose(refStr2, set->getErrorVector(h + 1).len,
        reference2);

//		uncompressSequence(query1, length, queryStr1);
//		uncompressSequence(query2, length, queryStr2);

    //CIGAR(length-M)
    if (mate1Unmatch)
      strcpy(cigar1, "*");
    else
      getCIGAR(cigar1, queryStr1, refStr1, set->getErrorVector(h));
//			sprintf(cigar1, "%dM", length);
    if (mate2Unmatch)
      strcpy(cigar2, "*");
    else
      getCIGAR(cigar2, queryStr2, refStr2, set->getErrorVector(h + 1));
//			sprintf(cigar2, "%dM", length);

    //MRNM(*)
    //MPOS(0)
    //ISIZE(0)
    if (!mate1Unmatch && !mate2Unmatch & sid1 == sid2) {
      if (set->getStrand(h) == set->getStrand(h + 1)) {
        tlen1 = soffset2 - soffset1;
        tlen2 = soffset1 - soffset2;
      } else if (set->getStrand(h) == FORWARD)
      {
        tlen1 = soffset2 - soffset1 + length;
        tlen2 = soffset1 - soffset2 - length;
      } else {
        tlen1 = soffset2 - soffset1 - length;
        tlen2 = soffset1 - soffset2 + length;
      }
    }

    //QUAL quals1 quals2
    qualStr1 = reader1->getQual(readId);
    if (set->getStrand(h) == BACKWARD
      )
      qualStr1 = reverseSequence(rcQualStr1, qualStr1);
    qualStr2 = reader2->getQual(readId);
    if (set->getStrand(h + 1) == BACKWARD
      )
      qualStr2 = reverseSequence(rcQualStr2, qualStr2);

    /* check whether mate1 is matched */
    fprintf(file, "%s\t%d\t%s\t%d\t%d\t%s\t%s\t%d\t%d\t%s\t%s", queryName1,
        flag1, refName1, soffset1, mapq1, cigar1, mateRefName1, soffset2, tlen1,
        queryStr1, qualStr1);

    if (!mate1Unmatch)
      writeOptionalField(queryStr1, refStr1, set->getErrorVector(h));
    else
      fprintf(file, "\n");

    /* check whether mate2 is matched */
    fprintf(file, "%s\t%d\t%s\t%d\t%d\t%s\t%s\t%d\t%d\t%s\t%s", queryName2,
        flag2, refName2, soffset2, mapq2, cigar2, mateRefName2, soffset1, tlen2,
        queryStr2, qualStr2);

    if (!mate2Unmatch)
      writeOptionalField(queryStr2, refStr2, set->getErrorVector(h + 1));
    else
      fprintf(file, "\n");
  }
}

/**
 * SamPairWriter::writeValidAlignment()
 * write a unmapped alignment in SAM format for a pair-end read.
 */
void SamPairWriter::writeInvalidAlignment(int readId, HitSet * set) {
  int flag1, flag2;
  char queryStr1[128], queryStr2[128];
  char * qualStr1, *qualStr2;

  //QNAME
  char * queryName1 = reader1->getReadName(readId);
  char * queryName2 = reader2->getReadName(readId);

  flag1 = SAM_FLAG_PAIRED | SAM_FLAG_UNMAPPED | SAM_FLAG_MATE_UNMAPPED
      | SAM_FLAG_FIRST_IN_PAIR;
  flag2 = SAM_FLAG_PAIRED | SAM_FLAG_UNMAPPED | SAM_FLAG_MATE_UNMAPPED
      | SAM_FLAG_SECOND_IN_PAIR;

  //SEQ query1 query2
  CompactSequence::decompose(queryStr1, length, reader1->getRead(readId));
  CompactSequence::decompose(queryStr2, length, reader2->getRead(readId));

  //QUAL quals1 quals2
  qualStr1 = reader1->getQual(readId);
  qualStr2 = reader2->getQual(readId);

  fprintf(file, "%s\t%d\t*\t0\t0\t*\t*\t0\t0\t%s\t%s\n", queryName1, flag1,
      queryStr1, qualStr1);
  fprintf(file, "%s\t%d\t*\t0\t0\t*\t*\t0\t0\t%s\t%s\n", queryName2, flag2,
      queryStr2, qualStr2);

  return;
}

char * Writer::reverseSequence(char * str2, char * str1) {
  int i;
  int len = strlen(str1);

  for (i = 0; i < len; i++)
    str2[len - i - 1] = str1[i];
  str2[len] = '\0';

  return str2;
}
