#!/usr/bin/env ruby

Signal.trap("PIPE", "EXIT")

require "aai"
require "abort_if"
require "fileutils"
require "trollop"

include AbortIf
include AbortIf::Assert

Aai.extend Aai
Aai.extend Aai::Utils
Time.extend Aai::CoreExtensions::Time

def move_files fnames, dest
  fnames.each do |fname|
    FileUtils.mv fname, dest
  end
end

opts = Trollop.options do
  version Aai::VERSION_BANNER

  banner <<-EOS

#{Aai::VERSION_BANNER}

  Seanie's AAI calculator.

  Each input file is treated as a file of ORFs for a single genome.

  Options:
  EOS

  opt(:cpus, "Number of CPUs to use", type: :integer, default: 1)
  opt(:infiles, "Input files", type: :strings)
  opt(:outdir, "Output directory", type: :string, default: ".")
  opt(:basename, "Base name for output file", type: :string,
      default: "aai_scores")
end

abort_if opts[:infiles].nil? || opts[:infiles].empty?,
         "No infiles given. Try #{__FILE__} --help for help."

Aai.check_command "blastp"
Aai.check_command "makeblastdb"
Aai.check_command "parallel"

Aai.check_files opts[:infiles]

abort_if File.exists?(opts[:outdir]),
         "#{opts[:outdir]} already exists. Choose a new outdir."

FileUtils.mkdir_p opts[:outdir]

seq_lengths = nil
clean_fnames = nil
blast_db_basenames = nil
btabs = nil
best_hits = nil
one_way = nil
two_way = nil
aai_strings = nil
score_dir = nil
outf = nil

Time.time_it "Processing input seqs", AbortIf.logger do
  seq_lengths, clean_fnames = Aai.process_input_seqs! opts[:infiles]
end

Time.time_it "Making blast databases", AbortIf.logger do
  blast_db_basenames = Aai.make_blastdbs! clean_fnames,
                                          opts[:cpus]
end

Time.time_it "Running blast jobs", AbortIf.logger do
  btabs = Aai.blast_permutations! clean_fnames,
                                  blast_db_basenames,
                                  opts[:cpus]
end

Time.time_it "Getting best hits", AbortIf.logger do
  best_hits = Aai.get_best_hits btabs, seq_lengths
end

Time.time_it "Getting one way hits", AbortIf.logger do
  one_way = Aai.one_way_aai best_hits
end

Time.time_it "Getting two way hits", AbortIf.logger do
  two_way = Aai.two_way_aai best_hits
end

Time.time_it "Making aai strings", AbortIf.logger do
  aai_strings = Aai.aai_strings one_way, two_way
end

Time.time_it "Writing score file", AbortIf.logger do
  score_dir = File.join opts[:outdir], "aai_scores"
  FileUtils.mkdir_p score_dir
  outf = File.join score_dir, "#{opts[:basename]}.aai.txt"
  File.open(outf, "w") do |f|
    f.puts %w[genome_A genome_B A_to_B_aai B_to_A_aai two_way_aai].
            join "\t"

    Aai.aai_strings(one_way, two_way).each do |str|
      f.puts str
    end
  end
end

Time.time_it "Cleaning up", AbortIf.logger do
  blast_db_dir = File.join opts[:outdir], "blastdbs"
  btab_dir = File.join opts[:outdir], "btabs"
  clean_fasta_dir = File.join opts[:outdir], "clean_fastas"

  FileUtils.mkdir_p [blast_db_dir, btab_dir, clean_fasta_dir]

  AbortIf.logger.debug { "Moving clean fastas" }
  move_files clean_fnames, clean_fasta_dir

  AbortIf.logger.debug { "Moving btabs" }
  move_files btabs, btab_dir

  AbortIf.logger.debug { "Moving blast dbs" }
  blast_db_basenames.each do |basename|
    move_files Dir.glob("#{basename}*"), blast_db_dir
  end
end