main.cc

#include "AbstractStatCollector.h"
#include "BasicStatsCollector.h"
#include "HistogramStatsCollector.h"
#include "CoverageMapStatsCollector.h"

#include "FpsModulator.h"

#include <iostream>
#include <fstream>
#include <streambuf>
#include <string>

static unsigned int totalReads;
static unsigned int updateRate;
static unsigned int firstUpdateRate;
static unsigned int wallReadCount = 400000;
static unsigned int coverageSkipFactor;
static std::string regionJson;
static std::string regionJsonFile;
static bool hasRegionSpec = false;
static bool isBatch = false;

static size_t fps;

using namespace std;
using namespace BamstatsAlive;

void printStatsJansson(AbstractStatCollector& rootStatCollector);

int main(int argc, char* argv[]) {
	
	string filename;
	updateRate = 100;
	firstUpdateRate = 0;
	coverageSkipFactor = 10;

	/* process the parameters */
	
	/* In order to use the -s/-l for coverage map statistics, 
	 * One need to make sure that the reads passed in all came
	 * from the same chromosome. Otherwise, the result is undefined.
	 */

	int ch;
	while((ch = getopt(argc, argv, "u:f:k:r:t:b")) != -1) {
		switch(ch) {
			case 'u':
				fps = atoi(optarg);
				break;
			case 'f':
				firstUpdateRate = atoi(optarg);
				break;
			case 'r':
				regionJson = std::string(optarg);
				hasRegionSpec = true;
				break;
            case 't':
                {
                    regionJsonFile = std::string(optarg);
                    std::ifstream fs(regionJsonFile);
                    regionJson = std::string((std::istreambuf_iterator<char>(fs)), std::istreambuf_iterator<char>());
                    hasRegionSpec = true;
                    break;
                }
            case 'k':
				coverageSkipFactor = atoi(optarg);
				break;
            case 'b':
                isBatch = true;
                coverageSkipFactor = 1;
                break;
		}
	}

	argc -= optind;
	argv += optind;

	if (argc == 0) 
		filename = "-";
	else 
		filename = *argv;


	/* open BAM file */

	BamTools::BamReader reader;
	reader.Open(filename);
	if(!reader.IsOpen()) {
		cout<<"{\"status\":\"error\", \"message\":\"Cannot open the specified file\"}"<<endl;
		exit(1);
	}

	const BamTools::RefVector refVector = reader.GetReferenceData();

	map<int32_t, string> chromIDNameMap;
	for(size_t i=0; i<refVector.size(); i++) {
		chromIDNameMap[reader.GetReferenceID(refVector[i].RefName)] = refVector[i].RefName;
	}

	/* Construct the statistics collectors */

	// NOTICE: The following codes utilize the new c++11 unique_ptr data type
	//         to automatically manage the life-time of heap based objects.
	//         See: http://www.cplusplus.com/reference/memory/unique_ptr/

	// Basic Scalar Statistics
	BasicStatsCollector bsc;
	HistogramStatsCollector* hsc = NULL;
	GenomicRegionStore *regionStore = NULL;

	// Histogram Statistics
	if(hasRegionSpec) {
		LOGS<<"Has Region Spec"<<std::endl;
		try {
			regionStore = new GenomicRegionStore(regionJson);
			LOGS<<regionStore->regions().size()<<" Regions specified"<<endl;
			hsc = new HistogramStatsCollector(chromIDNameMap, coverageSkipFactor, regionStore);
		}
		catch(...) {
			cout<<"{\"status\":\"error\", \"message\":\"Cannot parse region json string\"}"<<endl;
			exit(1);
		}
	}
	else {
		LOGS<<"Does not have region spec"<<std::endl;
		hsc = new HistogramStatsCollector(chromIDNameMap, coverageSkipFactor);
	}
	bsc.addChild(hsc);

	/* Process read alignments */
	BamTools::BamAlignment alignment;

    if(isBatch) {
        while(reader.GetNextAlignment(alignment)) {
            totalReads++;
            bsc.processAlignment(alignment, refVector);
        }
        printStatsJansson(bsc);
    }
    else {
        YiCppLib::FpsModulator<decltype(updateRate)> fpsModulator(updateRate, fps, 250);
        while(reader.GetNextAlignment(alignment) && totalReads <= wallReadCount) {
            totalReads++;
            bsc.processAlignment(alignment, refVector);

            if((totalReads > 0 && totalReads % updateRate == 0)) {
                printStatsJansson(bsc);
                fpsModulator.redraw();
            }
        }
        // count for all regions from which no read came
        printStatsJansson(bsc);
    }

	if(hsc) delete hsc;
	if(regionStore) delete regionStore;
}

void printStatsJansson(AbstractStatCollector& rootStatCollector) {

	// Create the root object that contains everything
	json_t * j_root = json_object();

	// Let the root object of the collector tree create Json
	rootStatCollector.appendJson(j_root);

	// Dump the json
	cout<<json_dumps(j_root, JSON_COMPACT | JSON_ENSURE_ASCII | JSON_PRESERVE_ORDER)<<";"<<endl;
}