ExtractLongSTRs.cpp

/***********************************************************************
 * Copyright (C) 2025, Po-Ru Loh, Broad Institute of MIT and Harvard
 * 
 * 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/>.
 *
 ***********************************************************************/

#include <iostream>
#include <fstream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <sys/time.h>
#include <htslib/sam.h>

using namespace std;

inline const char *targetStr(bam_hdr_t *hdr, int tid) {
  return 0 <= tid && tid < hdr->n_targets ? hdr->target_name[tid] : "*";
}

void processCram(const char *cramFile, const char *refFasta, const char *outFile,
		 int maxMismatches) {

  struct timeval tv;
  gettimeofday(&tv, NULL); double tPrev = tv.tv_sec + tv.tv_usec*1e-6;

  ofstream fout(outFile);
  if (!fout) {
    cerr << "ERROR opening " << outFile << " for writing" << endl;
    exit(1);
  }
  fout << "bpDelta" << "\t" << "mismatches" << "\t" << "QNAME" << "\t"
       << "FLAG" << "\t" << "RNAME" << "\t" << "POS" << "\t"
       << "QUAL" << "\t" << "RNEXT" << "\t" << "PNEXT" << "\t"
       << "SEQ" << endl;

  // open cram file
  htsFormat fmt = {0};
  char fmtStr[10 + strlen(refFasta) + 1];
  sprintf(fmtStr, "reference=%s", refFasta);
  hts_opt_add((hts_opt **) &fmt.specific, fmtStr);
  samFile *fin = sam_open_format(cramFile, "r", &fmt);
  if (fin == NULL) {
    cerr << "ERROR opening " << cramFile << endl;
    exit(1);
  }

  // only extract required fields
  hts_set_opt(fin, CRAM_OPT_REQUIRED_FIELDS, SAM_QNAME |
	      SAM_FLAG | SAM_RNAME | SAM_POS | SAM_MAPQ | SAM_RNEXT | SAM_PNEXT | SAM_SEQ);

  bam_hdr_t *hdr = sam_hdr_read(fin);
  if (hdr == NULL) {
    cerr << "ERROR reading header from " << cramFile << endl;
    exit(1);
  }

  bam1_t *aln = bam_init1();
  int ctr = 0, ctrRecorded = 0;
  int maxReadLen = 0;
  unsigned char *bases = NULL;
  int mismatchCtr[6+1];
  
  // iterate through cram file
  while (true) {
    int ret = sam_read1(fin, hdr, aln);
    if (ret == -1) break; // end of stream
    if (ret < -1) {
      fprintf(stderr, "ERROR: sam_read1 returned %d; truncated file?\n", ret);
      exit(1);
    }

    ctr++;

    int flag = aln->core.flag;
    int qual = aln->core.qual;

    if (flag&(BAM_FSECONDARY|BAM_FQCFAIL|BAM_FSUPPLEMENTARY))
      continue;

    int l_qseq = aln->core.l_qseq;
    if (l_qseq > maxReadLen) {
      maxReadLen = l_qseq;
      free(bases);
      bases = (unsigned char *) malloc(maxReadLen);
    }

    for (int d = 2; d <= 6; d++)
      mismatchCtr[d] = maxReadLen - l_qseq; // penalize hard-clipped reads
    bool earlyExited = false;
    for (int i = 0; i < l_qseq; i++) {
      bases[i] = bam_seqi(bam_get_seq(aln), i);
      for (int d = 2; d <= 6 && i-d >= 0; d++)
	mismatchCtr[d] += bases[i]==15 || bases[i]!=bases[i-d];
      if ((i&7)==7) {
	char found = 0;
	for (int d = 2; d <= 6; d++)
	  found |= mismatchCtr[d] <= maxMismatches;
	if (!found) {
	  earlyExited = true;
	  break;
	}
      }
    }
    if (earlyExited)
      continue;
    
    int d_best = 2;
    for (int d = d_best+1; d <= 6; d++)
      if (mismatchCtr[d] < mismatchCtr[d_best])
	d_best = d;
    if (mismatchCtr[d_best] <= maxMismatches) {

      char motif[d_best+1]; motif[d_best] = '\0';
      for (int k = 0; k < d_best; k++) {
	char baseCts[16]; memset(baseCts, 0, sizeof(baseCts));
	for (int i = k; i < l_qseq; i += d_best)
	  baseCts[bases[i]]++;
	if (baseCts[1] >= baseCts[2] && baseCts[1] >= baseCts[4] && baseCts[1] >= baseCts[8])
	  motif[k] = 'A';
	else if (baseCts[2] >= baseCts[4] && baseCts[2] >= baseCts[8])
	  motif[k] = 'C';
	else if (baseCts[4] >= baseCts[8])
	  motif[k] = 'G';
	else
	  motif[k] = 'T';
      }
      bool homopolymer = true;
      for (int k = 1; k < d_best; k++)
	if (motif[k] != motif[k-1])
	  homopolymer = false;

      if (!homopolymer) {
	fout << d_best << "\t" << mismatchCtr[d_best] << "\t" << bam_get_qname(aln) << "\t"
	     << flag << "\t" << targetStr(hdr, aln->core.tid) << "\t" << aln->core.pos+1 << "\t"
	     << qual << "\t" << targetStr(hdr, aln->core.mtid) << "\t" << aln->core.mpos+1 << "\t";
	for (int i = 0; i < l_qseq; i++)
	  fout << seq_nt16_str[bases[i]];
	fout << endl;
	ctrRecorded++;
      }
    }
  }

  free(bases);
  bam_destroy1(aln);
  bam_hdr_destroy(hdr);
  sam_close(fin);
  fout.close();

  gettimeofday(&tv, NULL); double tCur = tv.tv_sec + tv.tv_usec*1e-6;
  cout << "Processed " << ctr << " reads; wrote " << ctrRecorded << " to " << outFile
       << " (" << tCur-tPrev << " sec)" << endl;
}

int main(int argc, char *argv[]){

  if (argc != 4) {
    fprintf(stderr, "Usage:\n");
    fprintf(stderr, "- arg1 = bam/cram file or URL\n");
    fprintf(stderr, "- arg2 = reference fasta (unused if arg1 is bam)\n");
    fprintf(stderr, "- arg3 = output file\n");
    exit(1);
  }

  const char *cramFile = argv[1];
  const char *refFasta = argv[2];
  const char *outFile = argv[3];
  int maxMismatches = 10;

  processCram(cramFile, refFasta, outFile, maxMismatches);

  return 0;
}