source: trunk/yat/omic/Pileup.h @ 3310

#ifndef theplu_yat_omic_pileup
#define theplu_yat_omic_pileup

// $Id: Pileup.h 3310 2014-08-25 06:52:19Z peter $

/*
  Copyright (C) 2012, 2013, 2014 Peter Johansson

  This file is part of the yat library, http://dev.thep.lu.se/yat

  The yat library 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.

  The yat library 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/>.
*/

#include "BamRead.h"

#include <algorithm>
#include <cassert>
#include <deque>

namespace theplu {
namespace yat {
namespace omic {

  template<typename Iterator>
  class Pileup
  {
  public:

    /**
       Essentially a BamRead but with additional information regarding
       base we currently pointing to
     */
    class Entry
    {
    public:
      /**
         Create an Entry with first read \a b
       */
      Entry(const yat::omic::BamRead& b);

      /**
         \return const ref to BamRead
       */
      const yat::omic::BamRead& bam(void) const { return bam_; }

      uint32_t cigar(void) const
      {
        assert(c_index_ < bam_.core().n_cigar);
        return bam_.cigar(c_index_);
      }

      /**
         \return base quality of current position
       */
      unsigned short qual(void) const
      {
        assert(qpos_ < bam_.sequence_length());
        return bam_.qual(qpos_);
      }

      /**
         \return base of current position
       */
      uint8_t sequence(void) const;

      /**
         \brief iterate to next position with coverage

         If any read has an insertion here, next position will be a
         skip_ref locus; otherwise we iterate to next locus on
         reference that has coverage.
       */
      void increment(void);

      /**
         \return true if beyond last (aligned) base, i.e., we don't
         count clipped bases
       */
      bool end_of_sequence(void) const
      {
        return qpos_ >= bam_.sequence_length();
      }

      /**
         \return true if this position was deleted in this query
       */
      bool is_del(void) const
      {
        assert(c_index_ < bam_.core().n_cigar);
        return bam_cigar_type(cigar())==2;
      }

      size_t qpos(void) const { return qpos_; }
    private:
      yat::omic::BamRead bam_;
      // index of base pointed to
      size_t qpos_;
      size_t c_index_;
      size_t c_index2_;
      size_t skip_;
      void step_cigar(void);
      uint32_t cigar_op(void) const
      {
        return bam_cigar_op(cigar());
      }
      friend class Pileup;
    };

    // FIXME change to a filter iterator that ignores end_of_sequence entries
    typedef typename std::deque<Entry>::const_iterator const_iterator;

    /**
       Create an empty Pileup
     */
    Pileup(void) {}

    /**
       Create a Pileup that will use reads in range [first, last)
     */
    Pileup(Iterator first, Iterator last);

    /**
       first Entry in pileup
     */
    const_iterator begin(void) const { return data_.begin(); }

    /**
       last Entry in Pileup
     */
    const_iterator end(void) const;

    /**
       \return true if there are more positions to increment into
     */
    bool good(void) const { return !(bam_ == bam_end_ && data_.empty()); }

    /**
     */
    void increment(void);

    /**
       \return template id as defined bam file header
     */
    int32_t tid(void) const { return tid_; }

    /**
       \return position
     */
    int32_t pos(void) const { return pos_; }

    /**
       \return true if position does not exist in reference, i.e.,
       inserted region.
     */
    bool skip_ref(void) const;
  private:
    void step_cigar(void);

    Iterator bam_;
    Iterator bam_end_;
    int32_t pos_;
    int32_t tid_;
    std::deque<Entry> data_;
    uint32_t skip_ref_;
  };


  template<typename Iterator>
  Pileup<Iterator> make_pileup(Iterator first, Iterator last)
  {
    return Pileup<Iterator>(first, last);
  }

  // implementations

  template<typename Iterator>
  Pileup<Iterator>::Pileup(Iterator first, Iterator last)
    : bam_(first), bam_end_(last), skip_ref_(0)
  {
    tid_ = 0;
    pos_ = 0;
    increment();
  }


  template<typename Iterator>
  typename Pileup<Iterator>::const_iterator Pileup<Iterator>::end(void) const
  {
    if (data_.empty())
      return data_.end();
    // if last entry is passed pos, then don't include it in [begin, end) range
    if ((data_.back().bam().tid() > tid() ||
         data_.back().bam().pos() > pos())) {
      return data_.begin() + (data_.size()-1);
    }

    return data_.end();
  }


  template<typename Iterator>
  void Pileup<Iterator>::increment(void)
  {
    if (data_.empty()) {
      if (bam_==bam_end_)
        return;
      tid_ = bam_->core().tid;
      pos_ = bam_->core().pos;
    }
    else {

      // in insertion
      if (skip_ref()) {
        for (size_t i=0; i<data_.size(); ++i) {
          // ignore reads left of current position
          if (data_[i].end_of_sequence())
            continue;
          if (bam_cigar_type(data_[i].cigar())==1) {
            data_[i].increment();
          }
          --data_[i].skip_;
        }
        --skip_ref_;
      }
      // walk one base forward on reference
      else {
        ++pos_;
        for (size_t i=0; i<data_.size(); ++i) {
          if (data_[i].end_of_sequence())
            continue;
          data_[i].increment();
          // check if we're stepping into insertion
          if (!data_[i].end_of_sequence() && (data_[i].cigar_op() == BAM_CINS))
            skip_ref_ = std::max(skip_ref_, bam_cigar_oplen(data_[i].cigar()));
        }
        // stepped into an insertion
        if (skip_ref_) {
          for (size_t i=0; i<data_.size(); ++i)
            data_[i].skip_ = skip_ref_;
        }
      }
    }

    // read data
    while (bam_!=bam_end_ && bam_->core().tid==tid_ && bam_->pos()<=pos_) {
      // skip unmapped reads
      if (!(bam_->core().flag & BAM_FUNMAP)) {
        data_.push_back(Entry(*bam_));
      }
      ++bam_;
    }

    // pop reads that are end of sequence
    while (data_.size() && data_[0].end_of_sequence()) {
      data_.pop_front();
    }
  }


  template<typename Iterator>
  bool Pileup<Iterator>::skip_ref(void) const
  {
    return skip_ref_;
  }


  template<typename Iterator>
  Pileup<Iterator>::Entry::Entry(const yat::omic::BamRead& b)
    : bam_(b), qpos_(0), c_index_(0), c_index2_(0), skip_(0)
  {
    if (bam_.core().n_cigar==0)
      return;
    // iterate cigar to first match
    while (bam_cigar_type(bam_.cigar(c_index_)) != 3) {
      assert(bam_cigar_type(bam_.cigar(c_index_))!=2);
      if (bam_cigar_type(bam_.cigar(c_index_)) == 1) {
        qpos_ += b.cigar_oplen(c_index_);
      }
      ++c_index_;
      if (c_index_ >= bam_.core().n_cigar) {
        qpos_ = bam_.sequence_length();
        return;
      }
      assert(c_index_ < bam_.core().n_cigar);
    }
  }


  template<typename Iterator>
  void Pileup<Iterator>::Entry::increment(void)
  {
    if (bam_cigar_type(cigar()) & 1)
      ++qpos_;
    step_cigar();
  }


  template<typename Iterator>
  uint8_t Pileup<Iterator>::Entry::sequence(void) const
  {
    if (skip_==0) {
      if (bam_cigar_type(cigar()) & 1)
        return bam_.sequence(qpos_);
      else
        return 0;
    }
    if (cigar_op()!=BAM_CINS)
      return 0;
    return bam_.sequence(qpos_);
  }


  template<typename Iterator>
  void Pileup<Iterator>::Entry::step_cigar(void)
  {
    ++c_index2_;
    assert(c_index_<bam_.core().n_cigar);
    if (c_index2_ >= bam_.cigar_oplen(c_index_)) {
      c_index2_ = 0;
      ++c_index_;
    }
  }

}}}
#endif