#ifndef _HITSET_H_
#define _HITSET_H_

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

#include <assert.h>
#include <stdlib.h>
#include "lib.h"
#include "sequence.h"
#include "short.h"
#include "bitread.h"
#include "error.h"
#include "pair.h"

#define MAX_NUM_HITS 100
#define MAX_NUM_COUNT 4
#define MAX_INT 2147483647
#define UNASSIGNED_QUAL 256

#define MODE_NORMAL	0
#define MODE_RAW	1
#define MODE_SAM	2
#define MODE_CAM	3

/*
 *	Valid alignment struct
 */
typedef struct Hit {
  int id;
  unsigned int pos;
  int strand;
  ErrorVector error;
//	int nMismatch;
  int qual;
  int64 reference[8];
  int64 query[8]; /* this becomes an optional field */
} Hit;

/* 
 *	the set that containing all valid alignemtns of a particular read. 
 */
class HitSet {
public:
  int mode;
  Hit * hits;
  int maxWrite; //number of hits outputed
  int nHit; //number of hits stored
  int maxHit;
  int maxMatch;
  int maxQual;

  int length;
  CompactSequence * sequence;

  char str[256];

  bool sorted;
  bool strata;

  bool unique; /* only one valid alignment */
  bool properMatch;
  int counts[MAX_NUM_COUNT]; //counters for different error models
  double p[MAX_NUM_COUNT];
  double psum;

public:
  HitSet(int maxhit, int maxmatch, int maxqual, bool isSorted, bool isStrata) {
    int size;

    if (maxhit == 0)
      maxhit = MAX_INT;

    if (maxmatch == 0)
      maxmatch = MAX_INT;

    size = maxhit < maxmatch ? maxhit : maxmatch;
    hits = (Hit *) malloc(sizeof(Hit) * MAX_NUM_HITS);

    maxHit = maxWrite = maxhit;
    maxMatch = maxmatch;
    maxQual = maxqual;
    sorted = isSorted;
    strata = isSorted && isStrata;
    if (maxQual < 255 && maxhit < 10) {
      //keep up to 10 hits
      maxHit = 10;
    }

    unique = true;
    properMatch = true;
  }

  ~HitSet() {
    delete[] hits;
  }

  inline void init(CompactSequence * seq, int len) {
    sequence = seq;
    length = len;
    nHit = 0;

    for (int i = 0; i < MAX_NUM_COUNT; i++)
      counts[i] = 0;

    double pbase = 0.99;
    int i, j;
    for (i = 0; i < MAX_NUM_COUNT; i++) {
      //compute the binomial distribution
      p[i] = 1;
      for (j = 0; j < i; j++)
        p[i] *= 1 - pbase;
      for (j = 0; j < length - i; j++)
        p[i] *= pbase;
      for (j = 0; j < i; j++) {
        p[i] *= length - j;
        p[i] /= (j + 1);
      }
    }
    psum = 0;
  }

  inline int add(int64 * query, int64 * reference, unsigned int pos, strand s,
      ErrorVector * error, int qual, int rid) {
    int i;

    /* if there has been one element in the set,
     * there are multiple alignments for this read.
     */
//		if (nHit == 1 && hits[0].pos != pos)
//			unique = false;
    if (nHit >= MAX_NUM_HITS)
    {
//			elog(WARNING, "WARNING: result set is full. Some reads may be discarded.\n");
      return MSG_HITSETFULL;
    }

    if ((maxMatch == MAX_INT && nHit >= maxHit && !sorted) || nHit > maxMatch)
      return MSG_HITSETFULL;

    for (i = 0; i < nHit; i++) {
      /* linear search. can be replaced by binary search. */
      if (hits[i].pos == pos)
        return SUCCESS;
    }

    if (error->num < MAX_NUM_COUNT)
    {
      counts[error->num]++;
      psum += p[error->num];
    }

    i = nHit;
    if (sorted) {
      for (; i > 0; i--) {
        if (hits[i - 1].error.num <= error->num)
//				if (hits[i-1].error.num <= error.num
//				|| (hits[i-1].error.num == error.num && hits[i-1].qual <= qual))
          break;
        hits[i] = hits[i - 1];
      }
    }
    hits[i].pos = pos;
    hits[i].error = *error;
    hits[i].strand = s;
    hits[i].qual = UNASSIGNED_QUAL;
    hits[i].id = rid;
    //	BITMAP_COPY(query, hits[i].query);
    BitRead::copy(reference, hits[i].reference);

    if (nHit < maxHit || (maxMatch != MAX_INT && nHit <= maxMatch))
      nHit++;

    if ((maxMatch == MAX_INT && nHit >= maxHit && !sorted) || nHit > maxMatch)
      return MSG_HITSETFULL;

    //sorted

    return SUCCESS;
  }

  inline bool isFull() {
    if ((maxMatch == MAX_INT && nHit >= maxHit && !sorted) || nHit > maxMatch)
      return true;
    return false;
  }

  inline void reset() {
    nHit = 0;
    unique = true;
    properMatch = true;

    psum = 0.0001;
    for (int i = 0; i < MAX_NUM_COUNT; i++)
      counts[i] = 0;
  }

  /*
   * verify the quality score and shrink set under the limitation of maxWrite
   */
  inline void verifyQual() {
    int i, j;

    /*
     * suppress all alignments if more than maxMatch valid
     * alignmetns exist for it.
     */
    if (nHit > maxMatch) {
      nHit = 0;
      return;
    }

    for (i = 0, j = 0; i < nHit; i++) {
      if (getQual(i) < maxQual)
        continue;
      if (i != j)
        hits[j] = hits[i];
      j++;
      if (j >= maxWrite)
        break;
    }
    nHit = j;
  }

  inline int64 * getQuerySeq(int i) {
    return hits[i].query;
  }

  inline int64 * getReferenceSeq(int i) {
    return hits[i].reference;
  }

  inline uint32 getOffset(int i) {
    return hits[i].pos;
  }

  inline ErrorVector getErrorVector(int i) {
    return hits[i].error;
  }

  inline int getQual(int i) {
    if (hits[i].qual == UNASSIGNED_QUAL)
    {
      int num = hits[i].error.num;
      if (maxQual >= 255)
        hits[i].qual = 255;
      else {
        //		if (num < MAX_NUM_COUNT && counts[num] >= 1)
        //			return 0;
        if (num < MAX_NUM_COUNT) {
          hits[i].qual = (int) (-10 * log10(1 - p[num] / psum) + 0.5);
          if (hits[i].qual > 255)
            hits[i].qual = 255;
        } else {
          hits[i].qual = 3;
        }
      }
    }
    return hits[i].qual;
  }

  inline int getNumMismatch(int i) {
    return hits[i].error.num;
  }

  inline int getStrand(int i) {
    return hits[i].strand;
  }

  inline Hit * getHits() {
    return hits;
  }

  inline bool isProperMatch() {
    return properMatch;
  }

  inline int getNumHits() {
    return nHit;
  }

  inline int getNumAllHits() {
    return nHit;
  }
};

/*
 * the set containing the valid aligments of paired-end reads.
 */
class HitPairSet: public HitSet {
public:
  HitPairSet(int maxhit, int maxmatch, int maxqual, bool isSorted,
      bool isStrata) :
      HitSet(maxhit * 2, maxmatch * 2, maxqual, isSorted, isStrata) {
  }

  int build(HitSet * set1, HitSet * set2, uint32 maxins, uint32 minins,
      bool pairStrand[][2], bool mateMatch);

  inline int getQual(int i) {
    i -= i % 2;
    if (hits[i].qual == UNASSIGNED_QUAL)
    {
      int num = hits[i].error.num + hits[i + 1].error.num;
      if (maxQual >= 255)
        hits[i].qual = 255;
      else {
        //		if (num < MAX_NUM_COUNT && counts[num] >= 1)
        //			return 0;

        hits[i].qual = (int) (-10 * log10(1 - p[num] / psum) + 0.5);
        if (hits[i].qual > 255)
          hits[i].qual = 255;
      }
      hits[i + 1].qual = hits[i].qual;
    }
    return hits[i].qual;
  }

  /*
   * verify the quality score and shrink set under the limitation of maxWrite
   */
  inline void verifyQual() {
    int i, j;

    /*
     * suppress all alignments if more than maxMatch valid
     * alignmetns exist for it.
     */
    if (nHit > maxMatch) {
      nHit = 0;
      return;
    }

    for (i = 0, j = 0; i < nHit; i += 2) {
      if (getQual(i) < maxQual)
        continue;
      if (i != j)
        hits[j] = hits[i];
      j += 2;
      if (j >= maxWrite)
        break;
    }
    nHit = j;
  }

private:
  int add(Hit * hit1, Hit * hit2);
};

#endif

