#!/usr/bin/env ruby
# encoding: UTF-8
# Copyright (c) 2016 Stuart Glenn, Oklahoma Medical Research Foundation
# Licensed under 3 Clause BSD
# Wrapper to pull together some fastq and alignment metrics for report

# Things for which we are looking and reporting include
# - Percent undetermined for the demuxed flowcell (bcl2fastq)
# <=10
# - Average %Q30 across all lanes for each sample (bcl2fastq)
# =75
# - Average mean Q score across all lanes for each sample (bcl2fastq)
# =30
#
# - Percent pass filter unique reads (picard)
# PCT_PF_UQ_READS
# =75
# - Percent pass filter unique reads aligned to reference (picard)
# PCT_PF_UQ_READS_ALIGNED
# =90
# - Mean target coverage (picard)
# MEAN_TARGET_COVERAGE
# =100
# - Percent target bases > 10x (picard)
# PCT_TARGET_BASES_10X
# =90
# - Percent target bases > 20x (picard)
# PCT_TARGET_BASES_20X
# =85
# - Percent target bases > 30x (picard)
# PCT_TARGET_BASES_30X
# =80
#
# - VerifyBamID (FREELK1/FREELK0) ratio (VerifyBamID)
# =0.90
# verifyBamID --vcf ${vcf_file.vcf.gz}  --bam ${bam_file.bam} --out ${out_name} --ignoreRG --chip-none --maxDepth 1000 --precise
# read then out_name.selfSM

$:.unshift File.dirname(File.realpath(__FILE__))

require 'yaml'
require 'optparse'
require 'fileutils'
require 'uri'
require 'uri-o3'
require 'net/http'
require 'csv'
require 'yaml'
require 'openssl'

def metric_class(n,c,e,perc=false)
  s = eval("lambda {|x| #{c}}")
  Class.new do
    define_singleton_method :name do
      n
    end
    define_singleton_method :check_to_s do
      c
    end
    define_singleton_method :explanation do
      e
    end
    define_method :check do
      s
    end
    define_method :name do
      n
    end

    attr_reader :val

    define_method :initialize do |val|
      @val = val
    end

    def valid?
      if nil == @val
        return false
      elsif nil == self.check
        ""
      else
        self.check.call(@val)
      end
    end

    if perc
      define_method :formated_val do |fmt|
        return "NDA" if nil == @val
        sprintf(fmt,(@val * 100.0).round(2))
      end
    else
      define_method :formated_val do |fmt|
        return "NDA" if nil == @val
        sprintf(fmt,@val.round(2))
      end
    end
  end
end


METRICS = {
  :number_of_runs => metric_class("BCL Number of Runs", "x < 2 && x > 0",
                                  "The number of runs this sample has undergone."),
  :avg_of_lanes => metric_class("BCL Demux Average % of Lanes", "x > 0",
                                "The percentage of the reads across each lane that were identified as belonging to this sample. The percentage should be approximately equal to the fraction of the pooled samples for each respective sample."),
  :num_lanes => metric_class("BCL Demux Number of Lanes", "x >= 1 && x <= 4",
                            "The number of lanes across which this sample was sequenced."),
  :avg_q30 => metric_class("BCL Demux Average Q30", "x >= 75.0",
                          "The percentages of high quality bases generated by the sequencer relative to total bases sequenced. A high number indicates that a majority of bases were of high quality."),
  :avg_unknown => metric_class("BCL Demux Average % Unknown","x <= 10.0",
                              "The percentage of reads run concurrently for this sample for which the multiplexing tag was unknown or unable to be identified. Low numbers indicate that only sequence libraries that were expected to be present were sequenced and that the index read was performed correctly."),
  :avg_mean_q => metric_class("BCL Demux Average Mean Q","x >= 30.0",
                             "Ensures the average accuracy of the bases generated was above 99.999%."),
  :freelk_ratio => metric_class("VerifyBam FreeLK Ratio","x >= 0.90",
                               "The ratio of the two FreeLK outputs from VerifyBamID (FREELK1/FREELK0). A ratio above 0.9 indicates no sample cross-contamination during preparation."),
  :pct_pf_uq_reads => metric_class("Picard % PF UQ Reads","x >= 0.75",
                                   "The percentage of good quality reads that were only seen once. A high number means fewer PCR artifacts.",
                                   true),
  :pct_pf_uq_reads_aligned => metric_class("Picard % PF UQ Reads Aligned","x >= 0.90",
                                           "The percentage of good quality unique reads that actually aligned to the human reference genome. Numbers higher than 0.9 indicate only DNA from a human was sequenced.",
                                           true),
  :mean_target_coverage => metric_class("Picard Mean Target Coverage","x >= 100",
                                       "The average number of times each targeted base in the exome was sequenced with a unique read. High coverage is needed for accurate variant calling."),
  :pct_target_bases_10x => metric_class("Picard % Target Bases 10x","x >= 0.90",
                                        "The percentages of overall targeted bases that were covered at least 10 times with a unique read. High coverage is needed for accurate variant calling.",
                                        true),
  :pct_target_bases_20x => metric_class("Picard % Target Bases 20x","x >= 0.85",
                                        "The percentages of overall targeted bases that were covered at least 20 times with a unique read. High coverage is needed for accurate variant calling.",
                                        true),
  :pct_target_bases_30x => metric_class("Picard % Target Bases 30x","x >= 0.80",
                                        "The percentages of overall targeted bases that were covered at least 30 times with a unique read. High coverage is needed for accurate variant calling.",
                                        true),
  :wgs_coverage => metric_class("Illumina Coverage Depth","x >= 30",
                                       "Illumina's depth of coverage calculation from picard stats"),
  :mean_coverage => metric_class("Picard Mean Coverage","x >= 100",
                                       "The mean coverage in bases of the genome territory, after all filters are applied.. High coverage is needed for accurate variant calling."),
  :pct_10x => metric_class("Picard % Bases 10x","x >= 0.90",
                                        "The percentage of bases that attained at least 10X sequence coverage in post-filtering bases. High coverage is needed for accurate variant calling.",
                                        true),
  :pct_20x => metric_class("Picard % Bases 20x","x >= 0.85",
                                        "The percentage of bases that attained at least 20X sequence coverage in post-filtering bases. High coverage is needed for accurate variant calling.",
                                        true),
  :pct_30x => metric_class("Picard % Bases 30x","x >= 0.80",
                                        "The percentage of bases that attained at least 30X sequence coverage in post-filtering bases. High coverage is needed for accurate variant calling.",
                                        true),
}

@opts = {
  :verbose => false,
  :debug => true,
  :conf => nil,
  :bam => nil,
  :vcf => nil,
  :mode => :wes,
}

def demux_stats(conf,sample_name)
  c = YAML::load_file(conf)[sample_name]
  runs = c.map{|r| demux_stats_from_o3(r[:inputs].first)}.reject {|r| r.empty?}
  results = {
    :number_of_runs => c.size#runs.size
  }
  return results if runs.empty?
  runs.first.keys.each do |k|
    sum = runs.reduce(0.0) {|a,v| a += v[k]}
    if :num_lanes == k
      results[k] = sum
    else
      results[k] = sum/results[:number_of_runs]
    end
  end
  return results
end

def demux_stats_from_o3(object_path)
  unless object_path =~ /\Ao3:\/\//
    return {}
  end
  token = ENV['OS_AUTH_TOKEN']
  u = URI.parse(object_path)

  http = Net::HTTP.new(u.host, u.port)
  http.use_ssl=true
  http.verify_mode = OpenSSL::SSL::VERIFY_PEER
  req  = Net::HTTP::Head.new(u.request_uri)
  req["X-Auth-Token"] = token
  resp = http.request(req)
  msg = case resp
        when Net::HTTPSuccess
          r = {}
          %w/avg-q30 avg-unknown avg-mean-q avg-of-lanes num-lanes/.each do |q|
            h_key = "x-object-meta-sequence-run-library-demux-#{q}"
            if resp.key?(h_key)
              r[q.gsub(/-/,'_').to_sym] = resp[h_key].to_f.round(2)
            end
          end
          return r
        when Net::HTTPUnauthorized, Net::HTTPForbidden
          "Unauthorized, check OS_AUTH_TOKEN"
        when Net::HTTPNotFound
          "Object #{object_path} not found"
        when Net::HTTPServerError
          "Server Error: #{resp.value}"
        else
          "Unknown error: #{resp.value}"
        end
  $stderr.puts "O3 Error: #{msg}"
  return {}
end

# FREELK1/FREELK0
def verify_bam_ratio(bam_path,vcf_path)
  sm_tsv = verify_bam_sm_path_for(bam_path,vcf_path)
  return nil if nil == sm_tsv
  d = CSV.read(sm_tsv,{:headers=>true,:col_sep=>"\t",:converters=>:numeric})
  return (d[0]['FREELK1']/d[0]['FREELK0']).round(2)
end

def verify_bam_sm_path_for(bam_path,vcf_path)
  b = File.basename(bam_path,".bam")
  v = File.basename(vcf_path,".vcf.gz")
  out = File.join(File.dirname(bam_path),"verify-#{b}-#{v}")
  final = "#{out}.selfSM"
  if File.exist?(final)
    return final
  end
  cmd = %W/verifyBamID --self --precise --ignoreRG --chip-none --maxDepth 1000 --vcf
          #{vcf_path} --bam #{bam_path} --out #{out}/

  begin
    pid = spawn(*cmd,STDOUT=>'/dev/null',STDERR=>'/dev/null')
  rescue Errno::ENOENT
    $stderr.puts "Missing verifyBamID from $PATH"
    return nil
  end
  pid, status = Process.wait2(pid)
  if nil == status
    $stderr.puts "Unable to start verifyBamID"
    return nil
  elsif 0 != status.exitstatus
    $stderr.puts "Failure verifyBamID #{$?}"
    return nil
  end
  %w/depthSM log/.each do |ext|
    f = "#{out}.#{ext}"
    begin
      File.delete(f) if File.exists?(f)
    rescue
    end
  end
  return final
end

def picard_hs_ws(bam_path,o,mode,ref,i_file)
  metric,extra = if :wgs == mode
                   ["CollectWgsMetrics",nil]
                 elsif :hs == mode
                   ["CollectHsMetrics",%W/BAIT_INTERVALS=#{i_file} TARGET_INTERVALS=#{i_file}/]
                 end
  cmd = %W/picard #{metric} INPUT=#{bam_path} OUTPUT=#{o}
           VALIDATION_STRINGENCY=LENIENT REFERENCE_SEQUENCE=#{ref}/
  cmd += extra if extra
  begin
    pid = spawn(*cmd,STDOUT=>'/dev/null',STDERR=>'/dev/null')
  rescue Errno::ENOENT
    $stderr.puts "Missing picard from $PATH"
    return nil
  end
  pid, status = Process.wait2(pid)
  if nil == status
    $stderr.puts "Unable to start picard"
    return nil
  elsif 0 != status.exitstatus
    $stderr.puts "Failure picard (#{cmd.join(" ")}) #{$?}"
    return nil
  end
  return o
end

# - Percent pass filter unique reads (picard)
# PCT_PF_UQ_READS
# - Percent pass filter unique reads aligned to reference (picard)
# PCT_PF_UQ_READS_ALIGNED
# - Mean target coverage (picard)
# MEAN_TARGET_COVERAGE
# - Percent target bases > 10x (picard)
# PCT_TARGET_BASES_10X
# - Percent target bases > 20x (picard)
# PCT_TARGET_BASES_20X
# - Percent target bases > 30x (picard)
# PCT_TARGET_BASES_30X
def picard_stats(bam_path,conf)
  b = File.basename(bam_path,".bam")
  c = YAML::load_file(conf)
  picard_mode = if :wgs == @opts[:mode]
             :wgs
           else
             :hs
           end
  i_file = nil
  if :hs == picard_mode
    i_file = c[b].first[:interval_file] || c['DEFAULT'].first[:interval_file]
    if nil == i_file
      $stderr.puts "Exome requires interval file for picard metrics"
      return {}
    end
  end
  out = File.join(File.dirname(bam_path),"#{b}_#{picard_mode.to_s}_metrics.txt")
  unless File.exist?(out)
    picard_hs_ws(bam_path,out,picard_mode,c['DEFAULT'].first[:gatk_ref],i_file)
    unless File.exist?(out)
      $stderr.puts "No picard hs/ws metrics: #{out}"
      return {}
    end
  end
  d = ""
  IO.foreach(out) do |line|
    if "" != d
      d << line
      break
    end
    case picard_mode
    when :hs
      next unless line =~ /\ABAIT_SET/
    when :wgs
      next unless line =~ /\AGENOME_TERRITORY/
    end
    d << line
  end
  c = CSV.parse(d,{:headers=>true,:col_sep=>"\t",:converters=>:numeric})
  r = {}
  keys = case picard_mode
    when :hs
      %w/PCT_PF_UQ_READS PCT_PF_UQ_READS_ALIGNED MEAN_TARGET_COVERAGE PCT_TARGET_BASES_10X PCT_TARGET_BASES_20X PCT_TARGET_BASES_30X/
    when :wgs
      %w/MEAN_COVERAGE PCT_10X PCT_20X PCT_30X PCT_EXC_DUPE PCT_EXC_OVERLAP PCT_EXC_TOTAL/
    end
  keys.each do |k|
    r[k.downcase.to_sym] = c[0][k].round(2)
  end
  if :wgs == picard_mode
    r[:wgs_coverage] = (r[:mean_coverage] * (1.0 - r[:pct_exc_dupe] - r[:pct_exc_overlap]) ) / ( 1.0 - r[:pct_exc_total])
  end
  return r
end

op = OptionParser.new do |o|
  o.banner = "Usage: #{File.basename(__FILE__)} -c CONF_FILE -b BAM_FILE -v VCF_FILE"
  o.on("-V","--verbose","Enable verbose output") do
    @opts[:verbose] = true
  end
  o.on("-D","--debug","Enable debug mode") do
    @opts[:debug] = true
  end
  o.on("-h","--help","Show this help message") do
    puts o
    exit(0)
  end
  o.on("-m","--mode [MODE]",%w/wes wgs/,"Set reporting mode (wes, wgs), default wes") do |m|
    if nil == m
      $stderr.puts "Unknown mode #{m}"
      $stderr.puts o.help
      exit(1)
    end
    @opts[:mode] = m.to_sym
  end
  o.on("-c","--conf YAML","Load YAML for analysis conf, for fastq source") do |y|
    @opts[:conf] = File.expand_path(y)
  end
  o.on("-b","--bam BAM","Specify BAM file to parse") do |b|
    @opts[:bam] = File.expand_path(b)
  end
  o.on("-v","--vcf VCF","Specify VCF file to parse for verifyBamID") do |v|
    @opts[:vcf] = File.expand_path(v)
  end
  o.parse!
end

%w/vcf bam conf/.each do |k|
  unless @opts[k.to_sym] && File.exist?(@opts[k.to_sym])
    $stderr.puts "Missing valid #{k} parameter"
    $stderr.puts op.help()
    exit(1)
  end
end

removes = case @opts[:mode]
when :wes
  %w/pct_30x pct_20x pct_10x mean_coverage wgs_coverage/
when :wgs
  %w/pct_target_bases_30x pct_target_bases_20x pct_target_bases_10x pct_pf_uq_reads mean_target_coverage pct_pf_uq_reads_aligned/
else
  $stderr.puts "Invalid alignment mode: #{@pts[:mode]}"
  $stderr.puts op.help()
  exit(1)
end.map{|x| x.to_sym}

sample_name = File.basename(@opts[:bam],".bam")

stats = demux_stats(@opts[:conf],sample_name)
stats[:freelk_ratio] = verify_bam_ratio(@opts[:bam],@opts[:vcf])
stats.merge!(picard_stats(@opts[:bam],@opts[:conf]))

puts sample_name
puts "=" * sample_name.length
METRICS.each do |k,mc|
  next if removes.include?(k)
  m = mc.new(stats[k])
  pass = m.valid? ? " PASS ✅" : " FAIL ❌"
  printf("    %28s %7s%s\n",m.name,m.formated_val("%7.2f"),pass)
end
puts <<EOF

Check Definitions
=================
EOF
METRICS.each do |k,mc|
  next if removes.include?(k)
  puts "\n### #{mc.name}"
  puts " * #{mc.check_to_s}"
  puts " * #{mc.explanation}"
end
