#!/usr/bin/env ruby
#
# BioRuby bio-blastxmlparser Plugin
# Author:: Pjotr Prins
# License:: MIT License
#
# Copyright (C) 2010-2014 Pjotr Prins <pjotr.prins@thebird.nl>

rootpath = File.dirname(File.dirname(__FILE__))
$: << File.join(rootpath,'lib')

BLASTXML_VERSION = File.new(File.join(rootpath,'VERSION')).read.chomp

$stderr.print "BLAST XML Parser "+BLASTXML_VERSION+" Copyright (C) 2014 Pjotr Prins <pjotr.prins@thebird.nl>\n\n"

USAGE = <<EOM

  blastxmlparser filename(s)

    Use --help switch for more information

EOM

if ARGV.size == 0
  print USAGE
  exit 1
end

require 'bio-blastxmlparser'
require 'optparse'
require 'ostruct'

require 'bio-logger'
Bio::Log::CLI.logger('stderr')
Bio::Log::CLI.trace('info')

options = OpenStruct.new()

opts = OptionParser.new do |o|

  o.on_tail("-h", "--help", "Show help and examples") {
    print(opts)
    print USAGE
    exit()
  }

  o.banner = "== Usage\n  #{File.basename($0)} [options] file(s)"

  o.separator ""

  o.on("-p name", "--parser name", "Use split|nosplit parser (default split)") do |p|
    options.parser = p.to_sym
  end
  
  o.on("--filter filter",String, "Filtering expression") do |s|
    options.filter = s
  end

  o.on("-t num", "--threads num",String, "Use parallel threads") do |num|
    options.threads = num.to_i
  end

  o.on("-e filter","--exec filter",String, "Evaluate filter (deprecated)") do |s|
    $stderr.print "WARNING: -e,--exec switch is deprecated, use --filter instead!\n"
    options.filter = s
  end

  o.separator ""

  o.on("-n fields","--named fields",String, "Print named fields") do |s|
    options.fields = s.split(/,/)
  end
  o.on("--output-fasta","Output FASTA") do |b|
    options.output_fasta = true
  end
  
  o.on("-t erb","--template erb",String, "Use ERB template for output") do |s|
    require 'erb'
    options.template = s
  end
  
  o.separator ""

  o.on("--logger filename",String,"Log to file (default stderr)") do | name |
    Bio::Log::CLI.logger(name)
  end
 
  o.on("--trace options",String,"Set log level (default INFO, see bio-logger)") do | s |
    Bio::Log::CLI.trace(s)
  end
  
  o.on("-q", "--quiet", "Run quietly") do |q|
    Bio::Log::CLI.trace('error')
  end
  
  o.on("-v", "--verbose", "Run verbosely") do |v|
    Bio::Log::CLI.trace('info')
  end
  
  o.on("--debug", "Show debug messages") do |v|
    Bio::Log::CLI.trace('debug')
  end

end

begin
  opts.parse!(ARGV)

  Bio::Log::CLI.configure('bio-blastxmlparser')
  logger = Bio::Log::LoggerPlus['bio-blastxmlparser']

  if options.threads != 1
    begin
      require 'parallel'
    rescue LoadError
      $stderr.print "Error: Missing 'parallel' module. Install with command 'gem install parallel' if you want multiple threads\n"
      options.threads = 1
    end
  end

  if options.template
    include BioRdf
    fn = options.template
    raise "No template #{fn}!" if not File.exist?(fn)
    template = ERB.new(File.read(fn))
  end

  ARGV.each do | fn |
    logger.info("XML parsing #{fn}")
    parser_type = options.parser
    if !parser_type
      # If a file is smaller than 0.5 Gb the nosplit parser is used by default for performance
      if File.size(fn) > 512_000_000
        parser_type = :split
      else
        parser_type = :nosplit
      end
    end
    n = if parser_type == :nosplit
      Bio::BlastXMLParser::NokogiriBlastXml.new(File.new(fn)).to_enum
    else
      # default
      Bio::BlastXMLParser::BlastXmlSplitter.new(fn)
    end
    chunks = []
    chunks_count = 0
    NUM_CHUNKS=10_000

    process = lambda { |iter2,i|  # Process one BLAST iter block
      if parser_type == :nosplit
        iter = iter2 
      else
        xml = Nokogiri::XML.parse(iter2.join) { | cfg | cfg.noblanks }
        iter = Bio::BlastXMLParser::NokogiriBlastIterator.new(xml,self,:prefix=>nil)
      end
      res = []
      line_count = 0
      hit_count = 0
      iter.each do | hit |
        hit_count += 1
        hit.each do | hsp |
          do_print = if options.filter
                       eval(options.filter)
                     else
                       true
                     end
          if do_print
            line_count += 1
            if template
              res << template.result(binding)
            elsif options.output_fasta
              res << ">"+hit.accession+' '+iter.iter_num.to_s+'|'+iter.query_id+' '+hit.hit_id+' '+hit.hit_def+"\n"
              res << hsp.qseq+"\n"
            else
              # Default output
              if options.fields
                out = [iter.iter_num,hit_count,hsp.hsp_num]
                options.fields.each do | f |
                  out << eval(f)
                end
                res << out.join("\t")+"\n"
              else
                res << [iter.iter_num,iter.query_id,hit_count,hit.hit_id,hsp.hsp_num,hsp.evalue].join("\t")+"\n"
              end
            end
          end
        end
      end
      res
    } # end process 

    output = lambda { |collection| 
      collection.each do | result |
        result.each { |line| print line }
      end
    } # end output

    if options.threads == 1
      n.each do | iter |
        process.call(iter,0).each { | line | print line }
      end
    else
      n.each do | iter |
        chunks << iter
        chunks_count += 1
        if chunks.size > NUM_CHUNKS
          out = Parallel.map_with_index(chunks, :in_processes => options.threads) { | iter,i |
            process.call(iter,i)
          }
          # Output is forked to a separate process too
          fork do
            output.call out
            STDOUT.flush
            STDOUT.close
            exit 0
          end
          chunks = []
        end
      end
      output.call Parallel.map_with_index(chunks, :in_processes => options.threads) { | iter,i |
            process.call(iter,i)
          }
    end
  end
rescue OptionParser::InvalidOption => e
  $stderr.print e.message
end