#!/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

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"

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

Aai.check_files opts[:infiles]

FileUtils.mkdir_p opts[:outdir]

seq_lengths, clean_fnames = Aai.process_input_seqs! opts[:infiles]

blast_db_basenames = Aai.make_blastdbs! clean_fnames,
                                        opts[:cpus]

btabs = Aai.blast_permutations! clean_fnames,
                                blast_db_basenames,
                                opts[:cpus]

best_hits = Aai.get_best_hits btabs, seq_lengths

one_way = Aai.one_way_aai best_hits

two_way = Aai.two_way_aai best_hits

aai_strings = Aai.aai_strings one_way, two_way

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|
  Aai.aai_strings(one_way, two_way).each do |str|
    f.puts str
  end
end

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]

move_files clean_fnames, clean_fasta_dir
move_files btabs, btab_dir

blast_db_basenames.each do |basename|
  move_files Dir.glob("#{basename}*"), blast_db_dir
end