bionode-ncbi

var fs = require('fs')
var mkdirp = require('mkdirp')
var async = require('async')
var request = require('request')
var through = require('through2')
var JSONStream = require('JSONStream')
var xml2js = require('xml2js').parseString
var dld = require('dld')
var tool = require('tool-stream')


module.exports = exports = ncbi = new NCBI()

function NCBI() {
  this.APIROOT = 'http://eutils.ncbi.nlm.nih.gov/entrez/eutils/'
  this.DEFAULTS = 'retmode=json&version=2.0'
  this.RETURNMAX = 1000
  this.XMLPROPERTIES = {
    'sra': ['expxml', 'runs'],
    'biosample': ['sampledata'],
    'assembly': ['meta' ]
  }
  return this
}

NCBI.prototype.search = function(db, term) {
  var stream = through.obj(transform)
  if (term) { stream.write(term); stream.end() }
  return stream

  function transform(obj, enc, next) {
    var self = this
    var query = [
      ncbi.APIROOT + 'esearch.fcgi?',
      ncbi.DEFAULTS,
      'db=' + db,
      'term=' + encodeURI(obj),
      'usehistory=y'
    ].join('&')
    var getUIDs = _requestData(db)
    request({ uri: query, json: true })
    .pipe(JSONStream.parse())
    .pipe(_requestsSplit(db))
    .pipe(getUIDs)
    _attachStandardEvents(getUIDs, self, next)
  }
}

function _requestsSplit(db) {
  return through.obj(transform)
  function transform(obj, enc, next) {
    var self = this
    var webenv = obj.esearchresult.webenv
    var count = obj.esearchresult.count
    var numRequests = Math.floor(count / ncbi.RETURNMAX)
    for (var i = 0; i <= numRequests; i++) {
      var start = { db: db, webenv: webenv, retstart: i * ncbi.RETURNMAX }
      self.push(start)
    }
    next()
  }
}

function _requestData(db) {
  return through.obj(transform)
  function transform(obj, enc, next) {
    var self = this

    var query = [
      ncbi.APIROOT + 'esummary.fcgi?',
      ncbi.DEFAULTS,
      'db=' + obj.db,
      'query_key=1',
      'WebEnv=' + obj.webenv,
      'retmax=' + ncbi.RETURNMAX,
      'retstart=' + obj.retstart
    ].join('&')
    var xmlProperties = ncbi.XMLPROPERTIES[db] || []

    var finish = through.obj()
    if (db === 'sra') {
      var sraFinish = tool.filterObjectsArray('total_bases', '', 'runs.Run')
      _attachStandardEvents(sraFinish, self, next)
      finish
      .pipe(tool.ensureIsArray('runs.Run'))
      .pipe(sraFinish)
    }
    else {
      _attachStandardEvents(finish, self, next)
    }

    request({uri: query, json: true})
    .pipe(JSONStream.parse())
    .pipe(tool.extractProperty('result'))
    .pipe(tool.deleteProperty('uids'))
    .pipe(tool.arraySplit())
    .pipe(tool.XMLToJSProperties(xmlProperties))
    .pipe(finish)
  }
}

NCBI.prototype.link = function(srcDB, destDB, srcUID) {
  var stream = through.obj(getDestUID)
  if (srcUID) { stream.write(srcUID); stream.end() }
  return stream

  function getDestUID(srcUID, enc, next) {
    var self = this

    var query = [
      ncbi.APIROOT + 'elink.fcgi?',
      'dbfrom=' + srcDB,
      'db=' + destDB,
      'id=' + srcUID
    ].join('&')

    var link = {
      srcDB: srcDB,
      destDB: destDB,
      srcUID: srcUID
    }

    var linkName = srcDB+'_'+destDB

    var getLink = tool.attachToObject(link, 'destUID')

    request({ uri: query, json: true })
    .pipe(_wait(500))
    .pipe(tool.XMLToJS(true))
    .pipe(tool.extractProperty('LinkSet.0.LinkSetDb'))
    .pipe(tool.arraySplit())
    .pipe(tool.collectMatch('LinkName.0', linkName))
    .pipe(tool.extractProperty('Link'))
    .pipe(tool.arraySplit())
    .pipe(tool.extractProperty('Id.0'))
    .pipe(getLink)
    _attachStandardEvents(getLink, self, next)
  }
}

NCBI.prototype.download = function(db, term) {
  var stream = through.obj(transform)
  if (term) { stream.write(term); stream.end() }
  return stream
  function transform(obj, enc, next) {
    var self = this
    var download = _download()
    ncbi.search(db, obj)
    .pipe(_getURLs(db))
    .pipe(download)
    _attachStandardEvents(download, self, next)
  }
}

NCBI.prototype.urls = function(db, term) {
  var stream = through.obj(transform)
  if (term) { stream.write(term); stream.end() }
  return stream
  function transform(obj, enc, next) {
    var self = this
    var getURLs = _getURLs(db)
    ncbi.search(db, obj)
    .pipe(getURLs)
    _attachStandardEvents(getURLs, self, next)
  }
}

function _getURLs(db) {
  return through.obj(transform)
  function transform(obj, enc, next) {
    var self = this
    var parseURL = {
      sra: sraURL,
      assembly: assemblyURL
    }

    parseURL[db]()

    function sraURL() {
      var runs = obj.runs.Run
      async.eachSeries(runs, printSRAURL, next)
      function printSRAURL(run, cb) {
        var acc = run.acc
        var runURL = [
          'http://ftp.ncbi.nlm.nih.gov/sra/sra-instant/reads/ByRun/sra/',
          acc.slice(0,3) + '/',
          acc.slice(0,6) + '/',
          acc + '/',
          acc + '.sra',
        ].join('')
        self.push({url: runURL, uid: obj.uid})
        cb()
      }
    }

    function assemblyURL() {
      var rootURL = obj.meta.FtpSites.FtpPath._.replace('ftp://', 'http://')
      var assemblyURL = rootURL + 'Primary_Assembly/unplaced_scaffolds/FASTA/unplaced.scaf.fa.gz'
      self.push({url: assemblyURL, uid: obj.uid})
      next()
    }
  }
}

function _download(db, term) {
  return through.obj(transform)
  function transform(obj, enc, next) {
    var self = this
    var prevTime = Date.now()
    var currTime
    var chunkSizeMB = 1
    var chunkSize = chunkSizeMB * 1024 * 1024 //bytes
    var folder = obj.uid + '/'
    var path = folder + obj.url.replace(/.*\//, '')

    mkdirp.sync(obj.uid)
    if (!fs.existsSync(path)) {
      dld(obj.url, folder, chunkSize)
      .on('data', log)
      .on('end', function() {
        self.push(path)
        next()
      })
      .on('error', function(err) { self.emit('error', err) })
    }
    else {
      self.push(path)
      next()
    }

    function log(position, size) {
      var progress = (position * 100 / size).toFixed(2) + ' %'
      var sizeMB = Math.round(size / 1024 / 1024) + ' MB'
      currTime = Date.now()
      var diffTimeSec = (currTime - prevTime) / 1000
      prevTime = currTime
      var speed =  (chunkSizeMB / diffTimeSec).toFixed(2) + ' MB/s'
      var log = 'Downloading ' + path+' '+ progress + ' of ' + sizeMB + ' at ' + speed
      self.push(log)
    }
  }
}


function _attachStandardEvents(stream, self, next) {
  stream
  .on('data', function(data) { self.push(data) })
  .on('end', function() { next() })
  .on('error', function(err) { self.emit('error', err) })
}

function _wait(ms) {
  return through.obj(transform)
  function transform(obj, enc, next) {
    var self = this
    setTimeout(pushObj, ms)
    function pushObj() {
      self.push(obj)
      next()
    }
  }
}