#!/usr/bin/env ruby

# Author:: Ben J. Woodcroft
# Copyright:: 2011
# License:: The Ruby License

require 'rubygems'
require 'bio'
require 'bio-tm_hmm'
require 'optparse'

# If being run directly instead of being require'd,
# output one transmembrane per line, and
# indicate that a particular protein has no transmembrane domain
options = {
  :filter_in => false,
  :filter_out => false,
}
o = OptionParser.new do |opts|
    opts.banner = "
    Usage: #{File.basename __FILE__} [-f NUM] [-g NUM] [fasta_filename]

      Fasta file can be piped in on STDIN, or specified as the first unqualified argument.

      Without arguments, a description of the transmembrane domains is printed out for each input sequence\n\n"

  opts.on('-f','--filter-in MIN_TRANSMEMBRANE_DOMAINS','Print those sequences that have at _least_ MIN_TRANSMEMBRANE_DOMAINS transmembrane domain(s). Prints out the sequences in FASTA format.') do |m|
    options[:filter_in] = m.to_i #gets set to 0 when optional MIN_TRANSMEMBRANE_DOMAINS is omitted
  end
  opts.on('-g','--filter-out MAX_TRANSMEMBRANE_DOMAINS','Print those sequences that have at _most_ MAX_TRANSMEMBRANE_DOMAINS transmembrane domain(s). Prints out the sequences in FASTA format.') do |m|
    if m.nil?
      options[:filter_out] = 1
    else
      options[:filter_out] = m.to_i
    end
  end
end
o.parse!

runner = Bio::TMHMM::TmHmmWrapper.new

Bio::FlatFile.auto(ARGF).each do |seq|
  result = runner.calculate(seq.seq)
  name = seq.definition

  # Default output - a description of the TMDs for each input aaseq
  if options[:filter_in] == false and options[:filter_out] == false
    if result.has_domain?
      # At least one TMD found. Output each on a separate line
      result.transmembrane_domains.each do |tmd|
        puts [
        name,
        result.transmembrane_type,
        tmd.start,
        tmd.stop,
        tmd.orientation
      ].join("\t")
      end
    else
      puts [
      name,
      'No Transmembrane Domain Found'
    ].join("\t")
    end

  elsif options[:filter_in] != false
  	if result.transmembrane_domains.length >= options[:filter_in]
  		puts seq
  	end
  elsif options[:filter_out] != false
  	unless result.transmembrane_domains.length > options[:filter_out]
  		puts seq
  	end
  end
end