src/mapparser.cpp

//
//  mapparser.cpp
//  express
//
//  Created by Adam Roberts on 8/19/11.
//  Copyright 2011 Adam Roberts. All rights reserved.
//

#include "mapparser.h"
#include "main.h"
#include "fragments.h"
#include "targets.h"
#include "threadsafety.h"
#include "library.h"
#include <boost/algorithm/string/predicate.hpp>

using namespace std;

const size_t BUFF_SIZE = 9999;

size_t cigar_length(const char* cigar_str, vector<Indel>& inserts,
                    vector<Indel>& deletes) {
  inserts.clear();
  deletes.clear();
  const char* p_cig = cigar_str;
  size_t i = 0; // read index
  size_t j = 0; // genomic index
  while (*p_cig) {
    char* t;
    size_t op_len = (size_t)strtol(p_cig, &t, 10);
    char op_char = toupper(*t);
    switch(op_char) {
      case 'I':
      case 'S':
        inserts.push_back(Indel(i, op_len));
        i += op_len;
        break;
      case 'D':
        deletes.push_back(Indel(i, op_len));
        j += op_len;
        break;
      case 'M':
      case 'N':
        i += op_len;
        j += op_len;
        break;
    }
    p_cig = t + 1;
  }
  return j;
}

size_t cigar_length(vector<BamTools::CigarOp>& cigar_vec,
                    vector<Indel>& inserts, vector<Indel>& deletes) {
  inserts.clear();
  deletes.clear();
  size_t i = 0; // read index
  size_t j = 0; // genomic index
  for (size_t k = 0; k < cigar_vec.size(); ++k) {
    char op_char = cigar_vec[k].Type;
    size_t op_len = cigar_vec[k].Length;
    switch(op_char) {
      case 'I':
      case 'S':
        inserts.push_back(Indel(i, op_len));
        i += op_len;
        break;
      case 'D':
        deletes.push_back(Indel(i, op_len));
        j += op_len;
        break;
      case 'M':
      case 'N':
        i += op_len;
        j += op_len;
        break;
    }
  }
  return j;
}

MapParser::MapParser(Library* lib, bool write_active)
    : _lib(lib), _write_active(write_active) {

  string in_file = lib->in_file_name;
  string out_file = lib->out_file_name;
  bool is_sam = false;
  if (in_file.size() == 0) {
    logger.info("No alignment file specified. Expecting streaming input on "
                "stdin...\n");
    _parser.reset(new SAMParser(&cin));
    is_sam = true;
  } else {
    logger.info("Attempting to read '%s' in BAM format...", in_file.c_str());
    BamTools::BamReader* reader = new BamTools::BamReader();
    if (reader->Open(in_file)) {
      logger.info("Parsing BAM header...");
      _parser.reset(new BAMParser(reader));
      if (out_file.size()) {
        out_file += ".bam";
        BamTools::BamWriter* writer = new BamTools::BamWriter();
        if (writer->Open(out_file, reader->GetHeader(),
                         reader->GetReferenceData())) {
          bool sample = out_file.substr(out_file.length()-8,4) == "samp";
          _writer.reset(new BAMWriter(writer, sample));
        } else {
          logger.severe("Unable to open output BAM file '%s'.",
                        out_file.c_str());
        }
      }
    } else {
      delete reader;
      logger.info("Input is not in BAM format. Trying SAM...");
      ifstream* ifs = new ifstream(in_file.c_str());
      if (!ifs->is_open()) {
        logger.severe("Unable to open input SAM file '%s'.", in_file.c_str());
      }
      _parser.reset(new SAMParser(ifs));
      is_sam = true;
    }
  }

  if (is_sam && out_file.size()) {
    out_file += ".sam";
    ofstream* ofs = new ofstream(out_file.c_str());
    if (!ofs->is_open()) {
      logger.severe("Unable to open output SAM file '%s'.", out_file.c_str());
    }
    *ofs << _parser->header();
    bool sample = out_file.substr(out_file.length()-8,4) == "samp";
    _writer.reset(new SAMWriter(ofs, sample));
  }
}

void MapParser::threaded_parse(ParseThreadSafety* thread_safety_p,
                               size_t stop_at,
                               size_t num_neighbors) {
  ParseThreadSafety& pts = *thread_safety_p;
  bool fragments_remain = true;
  size_t n = 0;
  size_t still_out = 0;

  TargetTable& targ_table = *(_lib->targ_table);

  while (!stop_at || n < stop_at) {
    Fragment* frag = NULL;
    while (fragments_remain) {
      frag = new Fragment(_lib);
      fragments_remain = _parser->next_fragment(*frag);
      if (frag->num_hits()) {
        break;
      }
      delete frag;
      frag = NULL;
    }
    for (size_t i = 0; frag && i < frag->hits().size(); ++i) {
      FragHit& m = *(frag->hits()[i]);
      
      if (m.first_read() && m.first_read()->seq.length() > max_read_len) {
        logger.severe("Length of first read for fragment '%s' is longer than "
                      "maximum allowed read length (%d vs. %d). Increase the "
                      "limit using the '--max-read-len,L' option.",
                      m.frag_name().c_str(),  m.first_read()->seq.length(),
                      max_read_len);
      }
      if (m.second_read() && m.second_read()->seq.length() > max_read_len) {
        logger.severe("Length of second read for fragment '%s' is longer than "
                      "maximum allowed read length (%d vs. %d). Increase the "
                      "limit using the '--max-read-len,L' option.",
                      m.frag_name().c_str(),  m.second_read()->seq.length(),
                      max_read_len);
      }
      
      Target* t = targ_table.get_targ(m.target_id());
      if (!t) {
        logger.severe("Target sequence at index '%s' not found. Verify that it "
                      "is in the SAM/BAM header and FASTA file.",
                      m.target_id());
      }
      m.target(t);
      assert(t->id() == m.target_id());
      
      // Add num_neighbors targets on either side to the neighbors list.
      // Used for experimental feature.
      vector<const Target*> neighbors;
      for (TargID j = 1; j <= num_neighbors;  j++) {
        if (j <= m.target_id()) {
          neighbors.push_back(targ_table.get_targ(m.target_id() - j));
        }
        if (j + m.target_id() < targ_table.size()) {
          neighbors.push_back(targ_table.get_targ(m.target_id() + j));
        }
      }
      m.neighbors(neighbors);
    }
    boost::scoped_ptr<Fragment> done_frag(pts.proc_out.pop(false));
    while (done_frag) {
      if (_writer && _write_active) {
        _writer->write_fragment(*done_frag);
      }
      still_out--;
      done_frag.reset(pts.proc_out.pop(false));
    }

    if (!frag) {
      break;
    }

    pts.proc_in.push(frag);
    n++;
    still_out++;
  }

  pts.proc_in.push(NULL);

  while (still_out) {
    boost::scoped_ptr<Fragment> done_frag(pts.proc_out.pop(true));
    if (_writer && _write_active) {
      _writer->write_fragment(*done_frag);
    }
    still_out--;
  }
}

BAMParser::BAMParser(BamTools::BamReader* reader) : _reader(reader) {
  BamTools::BamAlignment a;

  size_t index = 0;
  foreach(const BamTools::RefData& ref, _reader->GetReferenceData()) {
    if (_targ_index.count(ref.RefName)) {
      logger.severe("Target '%s' appears multiple times in BAM header.",
                    ref.RefName.c_str());
    }
    _targ_index[ref.RefName] = index++;
    _targ_lengths[ref.RefName] = ref.RefLength;
  }

  // Get first valid ReadHit
  _read_buff = new ReadHit();
  do {
    if (!_reader->GetNextAlignment(a)) {
      logger.severe("Input BAM file contains no valid alignments.");
    }
  } while(!map_end_from_alignment(a));
}

bool BAMParser::next_fragment(Fragment& nf) {
  nf.add_map_end(_read_buff);

  BamTools::BamAlignment a;
  _read_buff = new ReadHit();

  while(true) {
    if (!_reader->GetNextAlignment(a)) {
      return false;
    } else if (!map_end_from_alignment(a)) {
      continue;
    } else if (!nf.add_map_end(_read_buff)) {
      return true;
    }
    _read_buff = new ReadHit();
  }
}

bool BAMParser::map_end_from_alignment(BamTools::BamAlignment& a) {
  ReadHit& r = *_read_buff;

  if (!a.IsMapped()) {
    return false;
  }

  bool is_paired = a.IsPaired();

  if (is_paired && !a.IsProperPair()) {
    return false;
  }

  if (is_paired && (direction == F || direction == R)) {
    return false;
  }
  
  bool is_reversed = a.IsReverseStrand();
  bool is_mate_reversed = a.IsMateReverseStrand();
  if (is_paired && (!a.IsMateMapped() || a.RefID != a.MateRefID ||
                    is_reversed == is_mate_reversed ||
                    (is_reversed && a.MatePosition > a.Position) ||
                    (is_mate_reversed && a.MatePosition < a.Position))) {
    return false;
  }
  
  bool left_first = (!is_paired && !is_reversed) ||
                    (a.IsFirstMate() && !is_reversed)||
                    (a.IsSecondMate() && is_reversed);

  if (((direction == RF || direction == R) && left_first) ||
      ((direction == FR || direction == F) && !left_first)) {
    return false;
  }
    
  r.name = a.Name;
  if (boost::algorithm::ends_with(r.name, "\1") ||
      boost::algorithm::ends_with(r.name, "\2")) {
    r.name = r.name.substr(r.name.size()-2);
  }

  r.reversed = is_reversed;
  r.first = !is_paired || a.IsFirstMate();
  r.targ_id = a.RefID;
  r.left = a.Position;
  r.mate_l = a.MatePosition;
  r.seq.set(a.QueryBases, is_reversed);
  r.bam = a;
  r.right = r.left + cigar_length(a.CigarData, r.inserts, r.deletes);
  
  foreach (Indel& indel, r.inserts) {
    if (indel.len > max_indel_size) {
      return false;
    }
  }
  foreach (Indel& indel, r.deletes) {
    if (indel.len > max_indel_size) {
      return false;
    }
  }
  
  return true;
}

void BAMParser::reset() {
  _reader->Rewind();

  // Get first valid FragHit
  BamTools::BamAlignment a;
  delete _read_buff;
  _read_buff = new ReadHit();
  do {
    _reader->GetNextAlignment(a);
  } while(!map_end_from_alignment(a));
}

SAMParser::SAMParser(istream* in) {
  _in = in;

  char line_buff[BUFF_SIZE];
  _read_buff = new ReadHit();
  _header = "";

  // Parse header
  size_t index = 0;
  while(_in->good()) {
    _in->getline(line_buff, BUFF_SIZE-1, '\n');
    if (line_buff[0] != '@') {
      break;
    }
    _header += line_buff;
    _header += "\n";

    string str(line_buff);
    size_t idx = str.find("SN:");
    if (idx!=string::npos) {
      string name = str.substr(idx+3);
      name = name.substr(0,name.find_first_of("\n\t "));
      if (_targ_index.count(name)) {
        logger.severe("Target '%s' appears multiple times in SAM header.",
                      str.c_str());
      }
      _targ_index[name] = index++;
      idx = str.find("LN:");
      if (idx != string::npos) {
        string len = str.substr(idx+3);
        len = len.substr(0,len.find_first_of("\n\t "));
        _targ_lengths[name] = atoi(len.c_str());
      }
    }
  }

  // Load first aligned read
  while(!map_end_from_line(line_buff)) {
    if (!_in->good()) {
      logger.severe("Input SAM file contains no valid alignments.");
    }
    _in->getline(line_buff, BUFF_SIZE-1, '\n');
  }
}

bool SAMParser::next_fragment(Fragment& nf) {
  nf.add_map_end(_read_buff);

  _read_buff = new ReadHit();
  char line_buff[BUFF_SIZE];

  while(_in->good()) {
    _in->getline (line_buff, BUFF_SIZE-1, '\n');
    if (!map_end_from_line(line_buff)) {
      continue;
    }
    if (!nf.add_map_end(_read_buff)) {
      break;
    }
    _read_buff = new ReadHit();
  }

  return _in->good();
}

bool SAMParser::map_end_from_line(char* line) {
  ReadHit& r = *_read_buff;
  string sam_line(line);
  char *p = strtok(line, "\t");
  int sam_flag = 0;
  bool paired = 0;
  bool left_first = 0;
  bool other_reversed = 0;

  int i = 0;
  while (p && i <= 9) {
    switch(i++) {
      case 0: {
        r.name = p;
        if (boost::algorithm::ends_with(r.name, "\1") ||
            boost::algorithm::ends_with(r.name, "\2")) {
          r.name = r.name.substr(r.name.size()-2);
        }
        break;
      }
      case 1: {
        sam_flag = atoi(p);
        if (sam_flag & 0x4) {
          goto stop;
        }
        paired = sam_flag & 0x1;
        if (paired && (direction == F || direction == R)) {
          goto stop;
        }
        if (paired && (!(sam_flag & 0x2) || sam_flag & 0x8)) {
          goto stop;
        }
        r.reversed = sam_flag & 0x10;
        other_reversed = sam_flag & 0x20;
        if (paired && r.reversed == other_reversed) {
          goto stop;
        }
        r.first = !paired || sam_flag & 0x40;
        left_first = ((!paired && !r.reversed) || (r.first && !r.reversed) ||
                        (!r.first && r.reversed));
        if (((direction == RF || direction == R) && left_first) ||
            ((direction == FR || direction == F) && !left_first)) {
          goto stop;
        }
        break;
      }
      case 2: {
        if(p[0] == '*') {
          goto stop;
        }
        if (!_targ_index.count(p)) {
          logger.severe("Target sequence '%s' not found. Verify that it is in "
                        "the SAM/BAM header and FASTA file.", p);
        }
        r.targ_id = _targ_index[p];
        break;
      }
      case 3: {
        r.left = (size_t)(atoi(p)-1);
        break;
      }
      case 4: {
        break;
      }
      case 5: {
        r.right = r.left + cigar_length(p, r.inserts, r.deletes);
        foreach (Indel& indel, r.inserts) {
          if (indel.len > max_indel_size) {
            goto stop;
          }
        }
        foreach (Indel& indel, r.deletes) {
          if (indel.len > max_indel_size) {
            goto stop;
          }
        }
        break;
      }
      case 6: {
        // skip if only one end mapped of paired-end read
        if(paired && p[0] != '=') {
          goto stop;
        }
        break;
      }
      case 7: {
        r.mate_l = atoi(p)-1;
        if (paired && ((r.reversed && r.left < (size_t)r.mate_l) ||
                       (other_reversed && r.left > (size_t)r.mate_l))) {
          goto stop;
        }
        break;
      }
      case 8: {
        break;
      }
      case 9: {
        r.seq.set(p, r.reversed);
        r.sam = sam_line;
        goto stop;
      }
    }
    p = strtok(NULL, "\t");
  }
 stop:
  return i == 10;
}

void SAMParser::reset() {
  // Rewind input file
  _in->clear();
  _in->seekg(0, ios::beg);

  // Load first alignment
  char line_buff[BUFF_SIZE];
  delete _read_buff;
  _read_buff = new ReadHit();

  while(_in->good()) {
    _in->getline(line_buff, BUFF_SIZE-1, '\n');
    if (line_buff[0] != '@') {
      break;
    }
  }

  while(!map_end_from_line(line_buff)) {
    _in->getline(line_buff, BUFF_SIZE-1, '\n');
  }
}

BAMWriter::BAMWriter(BamTools::BamWriter* writer, bool sample)
   : _writer(writer) {
  _sample = sample;
}

BAMWriter::~BAMWriter() {
  _writer->Close();
}

void BAMWriter::write_fragment(Fragment& f) {
  if (_sample) {
    const FragHit* hit = f.sample_hit();
    PairStatus ps = hit->pair_status();
    if (ps != RIGHT_ONLY) {
      _writer->SaveAlignment(hit->left_read()->bam);
    }
    if (ps != LEFT_ONLY) {
      _writer->SaveAlignment(hit->right_read()->bam);
    }
  } else {
    double total = 0;
    foreach(FragHit* hit, f.hits()) {
      total += sexp(hit->params()->posterior);
      PairStatus ps = hit->pair_status();
      if (ps != RIGHT_ONLY) {
        hit->left_read()->bam.AddTag("XP","f",(float)sexp(hit->params()->posterior));
        _writer->SaveAlignment(hit->left_read()->bam);
      }
      if (ps != LEFT_ONLY) {
        hit->right_read()->bam.AddTag("XP","f",(float)sexp(hit->params()->posterior));
        _writer->SaveAlignment(hit->right_read()->bam);
      }
    }
    assert(approx_eq(total, 1.0));
  }
}

SAMWriter::SAMWriter(ostream* out, bool sample) : _out(out) {
  _sample = sample;
}

SAMWriter::~SAMWriter() {
  _out->flush();
}

void SAMWriter::write_fragment(Fragment& f) {
  if (_sample) {
    const FragHit* hit = f.sample_hit();
    PairStatus ps = hit->pair_status();

    if (ps != RIGHT_ONLY) {
      *_out << hit->left_read()->sam << endl;
    }
    if (ps != LEFT_ONLY) {
      *_out << hit->right_read()->sam << endl;
    }
  } else {
    double total = 0;
    foreach(const FragHit* hit, f.hits()) {
      total += sexp(hit->params()->posterior);
      PairStatus ps = hit->pair_status();
      if (ps != RIGHT_ONLY) {
        *_out << hit->left_read()->sam << " XP:f:"
              << (float)sexp(hit->params()->posterior) << endl;
      }
      if (ps != LEFT_ONLY) {
        *_out << hit->right_read()->sam << " XP:f:"
              << (float)sexp(hit->params()->posterior) << endl;
      }
    }
    assert(approx_eq(total, 1.0));
  }
}