json output

[bamtools.git] / bamtools_count.cpp

diff --git a/bamtools_count.cpp b/bamtools_count.cpp
index 1560e1c5d4bdc5c79b973b9f5fef6c72e145cfab..67d1e9b1149494db1d1d96d9f0fb277e6f2c83c5 100644 (file)
--- a/bamtools_count.cpp
+++ b/bamtools_count.cpp
@@ -12,6 +12,7 @@
 // ***************************************************************************
 
 #include <iostream>
+#include <sstream>
 #include <string>
 #include <vector>
 
@@ -29,16 +30,19 @@ using namespace BamTools;
 struct CountTool::CountSettings {
 
     // flags
+    bool HasBamIndexFilename;
     bool HasInputBamFilename;
     bool HasRegion;
 
     // filenames
-    std::string InputBamFilename;
-    std::string Region;
+    string BamIndexFilename;
+    string InputBamFilename;
+    string Region;
     
     // constructor
     CountSettings(void)
-        : HasInputBamFilename(false)
+        : HasBamIndexFilename(false)
+        , HasInputBamFilename(false)
         , HasRegion(false)
         , InputBamFilename(Options::StandardIn())
     { }  
@@ -52,14 +56,15 @@ CountTool::CountTool(void)
     , m_settings(new CountSettings)
 { 
     // set program details
-    Options::SetProgramInfo("bamtools count", "prints alignment counts for a BAM file", "-in <filename> ");
+    Options::SetProgramInfo("bamtools count", "prints alignment counts for a BAM file", "-in <filename> [-region <REGION> [-index <filename>]]");
     
     // set up options 
     OptionGroup* IO_Opts = Options::CreateOptionGroup("Input & Output");
-    Options::AddValueOption("-in", "BAM filename", "the input BAM file", "", m_settings->HasInputBamFilename, m_settings->InputBamFilename, IO_Opts, Options::StandardIn());
+    Options::AddValueOption("-in",    "BAM filename",        "the input BAM file", "", m_settings->HasInputBamFilename, m_settings->InputBamFilename, IO_Opts, Options::StandardIn());
+    Options::AddValueOption("-index", "BAM index filename", "the BAM index file",  "", m_settings->HasBamIndexFilename, m_settings->BamIndexFilename, IO_Opts);
     
     OptionGroup* FilterOpts = Options::CreateOptionGroup("Filters");
-    Options::AddValueOption("-region", "REGION", "genomic region. See README for more details", "", m_settings->HasRegion, m_settings->Region, FilterOpts);
+    Options::AddValueOption("-region", "REGION", "genomic region. Index file is recommended for better performance. See \'bamtools help index\' for more details on creating one", "", m_settings->HasRegion, m_settings->Region, FilterOpts);
 }
 
 CountTool::~CountTool(void) { 
@@ -78,31 +83,86 @@ int CountTool::Run(int argc, char* argv[]) {
     Options::Parse(argc, argv, 1);
 
     //open our BAM reader
-//     BamReader reader;
-//     reader.Open(m_settings.InputBamFilename);
+    BamReader reader;
+    reader.Open(m_settings->InputBamFilename);
 
-    // count alignments
-    string startChrom;
-    string stopChrom;
-    int startPos;
-    int stopPos;
+    // alignment counter
+    int alignmentCount(0);
     
+    // set up error handling
+    ostringstream errorStream("");
+    bool foundError(false);
+    
+    // if no region specified, count entire file 
     if ( !m_settings->HasRegion ) {
-        cerr << "Counting all alignments " << endl;
-    } else {
-        if ( Utilities::ParseRegionString(m_settings->Region, startChrom, startPos, stopChrom, stopPos) ) {
-            cerr << "Counting only alignments in region " << m_settings->Region << endl;
-            cerr << "StartChrom: " << startChrom << endl;
-            cerr << "StartPos:   " << startPos << endl;
-            cerr << "StopChrom:  " << stopChrom << endl;
-            cerr << "StopPos:    " << stopPos << endl;
+        BamAlignment al;
+        while ( reader.GetNextAlignment(al) ) 
+            ++alignmentCount;
+    } 
+    
+    // more complicated - region specified
+    else {
+        
+        Region region;
+        if ( Utilities::ParseRegionString(m_settings->Region, reader, region) ) {
+
+            // has index, we can jump directly to 
+            if ( m_settings->HasBamIndexFilename ) {
+              
+                // this is kind of a hack...?
+                // re-open BamReader, this time with the index file
+                reader.Close();
+                reader.Open(m_settings->InputBamFilename, m_settings->BamIndexFilename);
+              
+                // attempt Jump(), catch error
+                if ( !reader.Jump(region.StartChromID, region.StartPosition) ) {
+                    foundError = true;
+                    errorStream << "Could not jump to desired REGION: " << m_settings->Region << endl;
+                }
+            }
+            
+            else {
+              
+                // read through sequentially, until first overlapping read is found
+                BamAlignment al;
+                bool alignmentFound(false);
+                while( reader.GetNextAlignment(al) ) {
+                    if ( (al.RefID == region.StartChromID ) && ( (al.Position + al.Length) >= region.StartPosition) ) {
+                        alignmentFound = true;
+                        break;
+                    }
+                }
+                
+                // if overlapping alignment found (not EOF), increment counter
+                if ( alignmentFound) ++ alignmentCount; 
+            }
+            
+            // -----------------------------
+            // count alignments until stop hit
+            
+            if ( !foundError ) {
+                // while valid alignment AND
+                // ( either current ref is before stopRefID OR
+                //   current ref stopRefID but we're before stopPos )
+                BamAlignment al;
+                while ( reader.GetNextAlignment(al) && ((al.RefID < region.StopChromID ) || (al.RefID == region.StopChromID && al.Position <= region.StopPosition)) )
+                    ++alignmentCount;
+            }
+
+        } else {
+            foundError = true;
+            errorStream << "Could not parse REGION: " << m_settings->Region << endl;
+            errorStream << "Be sure REGION is in valid format (see README) and that coordinates are valid for selected references" << endl;
         }
     }
      
-    cerr << " from " << m_settings->InputBamFilename << endl;
-    cerr << "FEATURE NOT YET IMPLEMENTED!" << endl;
-
+    // print errors OR results 
+    if ( foundError )
+        cerr << errorStream.str() << endl;
+    else
+        cout << alignmentCount << endl;
+    
     // clean & exit
-//     reader.Close();
-    return 0;
+    reader.Close();
+    return (int)foundError;
 }
\ No newline at end of file