require "base64"
require "csv"
require "fileutils"
require "json"
require "rummager"
require "zip"
require "analytics/popular_queries"
require "relevancy/load_judgements"

namespace :learn_to_rank do
  desc "Fetch data from BigQuery.  This costs money!"
  task :fetch_bigquery_export, [:output, :viewcount] do |_, args|
    raise 'set $ENABLE_LTR to "true" to use learn_to_rank' unless Search::RelevanceHelpers.ltr_enabled?
    raise 'Base64 encoded environment variable "BIGQUERY_CREDENTIALS" is required' unless ENV.key? "BIGQUERY_CREDENTIALS"

    data = LearnToRank::DataPipeline::Bigquery.fetch(JSON.parse(Base64.decode64(ENV["BIGQUERY_CREDENTIALS"])), viewcount: args.viewcount || 10)

    CSV.open("tmp/#{args[:output]}.csv", "wb") do |csv|
      csv << data.first.keys
      data.each do |row|
        csv << row.values
      end

      while data.next?
        data = data.next
        data.each do |row|
          csv << row.values
        end
      end
    end
  end

  desc "Export a CSV of relevancy judgements generated from CTR on popular queries"
  task :generate_relevancy_judgements, [:queries, :output] do |_, args|
    raise 'set $ENABLE_LTR to "true" to use learn_to_rank' unless Search::RelevanceHelpers.ltr_enabled?

    queries = LearnToRank::DataPipeline::LoadSearchQueries.from_csv(args.queries)
    generator = LearnToRank::DataPipeline::RelevancyJudgements.new(queries: queries)
    judgements = generator.relevancy_judgements
    data = judgements.force

    CSV.open("tmp/#{args.output}.csv", "wb") do |csv|
      csv << data.first.keys
      data.each do |row|
        csv << row.values
      end
    end
  end

  desc "Export a CSV of SVM-formatted relevancy judgements for training a model"
  task :generate_training_dataset, [:judgements, :output_dir] do |_, args|
    raise 'set $ENABLE_LTR to "true" to use learn_to_rank' unless Search::RelevanceHelpers.ltr_enabled?

    csv = args.judgements || "tmp/autogenerated_relevancy_judgements.csv"
    svm_dir = args.output_dir || "tmp/ltr_data"
    FileUtils.mkdir_p svm_dir

    judgements_data = Relevancy::LoadJudgements.from_csv(csv)
    judgements = LearnToRank::DataPipeline::EmbedFeatures.new(judgements_data).augmented_judgements
    svm = LearnToRank::DataPipeline::JudgementsToSvm.new(judgements).svm_format_grouped_by_query

    File.open("#{svm_dir}/train.txt", "wb") do |train|
      File.open("#{svm_dir}/validate.txt", "wb") do |validate|
        File.open("#{svm_dir}/test.txt", "wb") do |test|
          files = [train, train, train, train, train, train, train, test, test, validate].shuffle
          svm.each_with_index do |query_set, index|
            # 70% in train 20% in test, 10% in validate
            file = files[index % 10]
            query_set.each { |row| file.puts(row) }
          end
        end
      end
    end
  end

  namespace :reranker do
    desc "Train a reranker model with relevancy judgements"
    task :train, [:svm_dir, :model_dir] do |_, args|
      raise 'set $ENABLE_LTR to "true" to use learn_to_rank' unless Search::RelevanceHelpers.ltr_enabled?

      model_dir = args.model_dir || "tmp/libsvm"
      svm_dir = args.svm_dir || "tmp/ltr_data"
      sh "env OUTPUT_DIR=#{model_dir} TRAIN=#{svm_dir}/train.txt VALI=#{svm_dir}/validate.txt TEST=#{svm_dir}/test.txt ./ltr/scripts/train.sh"
    end

    desc "Serves a trained model"
    task :serve, [:model_dir] do |_, args|
      raise 'set $ENABLE_LTR to "true" to use learn_to_rank' unless Search::RelevanceHelpers.ltr_enabled?

      model_dir = args.model_dir || "tmp/libsvm"
      sh "env EXPORT_PATH=#{__dir__}/../../#{model_dir} ./ltr/scripts/serve.sh"
    end

    desc "Evaluate search performance using nDCG with and without the model"
    task :evaluate, [:relevancy_judgements] do |_, args|
      raise 'set $ENABLE_LTR to "true" to use learn_to_rank' unless Search::RelevanceHelpers.ltr_enabled?

      ndcg_at = "10"

      csv = args.datafile || begin
        bucket_name = ENV["AWS_S3_RELEVANCY_BUCKET_NAME"]
        raise "Missing required AWS_S3_RELEVANCY_BUCKET_NAME envvar" if bucket_name.nil?

        csv = Tempfile.open(["judgements", ".csv"])
        o = Aws::S3::Object.new(bucket_name: bucket_name, key: "autogenerated_judgements.csv")
        o.get(response_target: csv.path)
        csv.path
      end

      rounds = ["relevance:disable", nil]
      begin
        results, results_with_model = rounds.map do |ab_test_round|
          judgements = Relevancy::LoadJudgements.from_csv(csv)
          evaluator = Evaluate::Ndcg.new(judgements, ab_test_round)
          evaluator.compute_ndcg
        end
      ensure
        if csv.is_a?(Tempfile)
          csv.close
          csv.unlink
        end
      end

      merged = results.keys.each_with_object({}) do |query, hsh|
        hsh[query] = {
          without: results[query],
          with_model: results_with_model[query],
        }
      end

      maxlen = results.keys.map { |query, _| query.length }.max
      score_maxlen = results.values.map { |score, _| score[ndcg_at].to_s.length }.max

      merged.map do |(query, scores)|
        winning = scores[:without][ndcg_at] <= scores[:with_model][ndcg_at] ? "√" : "x"
        justified_query = "#{query}:".ljust(maxlen + 1)
        justified_score = scores[:without][ndcg_at].to_s.ljust(score_maxlen + 1)
        puts "#{winning} #{justified_query} #{justified_score} #{scores[:with_model][ndcg_at]}"
      end

      winning = merged.dig("average_ndcg", :without, ndcg_at) <= merged.dig("average_ndcg", :with_model, ndcg_at)

      puts "---"
      puts "without model score: #{merged['average_ndcg'][:without][ndcg_at]}"
      puts "with model score: #{merged['average_ndcg'][:with_model][ndcg_at]}"
      puts "Without model: #{merged['average_ndcg'][:without]}"
      puts "With model: #{merged['average_ndcg'][:with_model]}"
      puts "The model has a #{winning ? 'good' : 'bad'} score"
    end
  end
end
