# frozen_string_literal: true

require "json"
require "open3"
require "optparse"
require "tmpdir"
require "zlib"
require "digest"
require "fileutils"
require_relative "data_set"

module MoonModel
  class DataPreparer
    MAGIC = "MOONDATA1\n".b
    DEFAULT_DTM = "/scratch0/lunaserv/layer-new/luna/wac_dtm_numeric/WAC_GLD100_V1.0_GLOBAL_with_LOLA_30M_POLE.0032km.cub"
    DEFAULT_REFLECTANCE = "/scratch0/lunaserv/layer-new/luna/wac_normalized_reflectance/*.tif"

    def self.run(argv)
      options = { dtm: DEFAULT_DTM, reflectance: DEFAULT_REFLECTANCE,
                  output: File.expand_path("../../data/moon_surface.bin.gz", __dir__) }
      OptionParser.new do |opts|
        opts.banner = "Usage: prepare_lunar_data [options]"
        opts.on("--dtm FILE") { |v| options[:dtm] = v }
        opts.on("--reflectance GLOB") { |v| options[:reflectance] = v }
        opts.on("-o", "--output FILE") { |v| options[:output] = v }
      end.parse!(argv)
      new(**options).run
      0
    rescue StandardError => e
      warn "prepare_lunar_data: #{e.message}"
      2
    end

    def initialize(dtm:, reflectance:, output:)
      @dtm = dtm
      @reflectance = reflectance
      @output = output
    end

    def run
      sources = Dir.glob(@reflectance).sort
      raise "DTM not found: #{@dtm}" unless File.file?(@dtm)
      raise "no reflectance files matched #{@reflectance}" if sources.empty?

      info = JSON.parse(capture("gdalinfo", "-json", @dtm))
      width, height = info.fetch("size")
      transform = info.fetch("geoTransform")
      xmin = transform[0]; ymax = transform[3]
      xmax = xmin + transform[1] * width
      ymin = ymax + transform[5] * height

      Dir.mktmpdir("moon-data-") do |tmp|
        elevation_path = File.join(tmp, "elevation.bin")
        reflectance_path = File.join(tmp, "reflectance.bin")
        vrt = File.join(tmp, "reflectance.vrt")
        system!("gdal_translate", "-q", "-of", "ENVI", "-ot", "Float32", @dtm, elevation_path)
        wrapped_sources = sources.map.with_index do |source, index|
          source_info = JSON.parse(capture("gdalinfo", "-json", source))
          source_transform = source_info.fetch("geoTransform")
          source_width, source_height = source_info.fetch("size")
          sx0 = source_transform[0]; sy1 = source_transform[3]
          sx1 = sx0 + source_transform[1] * source_width
          sy0 = sy1 + source_transform[5] * source_height
          next source unless sx0.negative?

          # The western COGs are encoded as -180..0 while GLD100 is 0..360.
          shifted = File.join(tmp, "wrapped_#{index}.vrt")
          circumference = 2.0 * (sx1 - sx0)
          system!("gdal_translate", "-q", "-of", "VRT", "-a_ullr",
                  (sx0 + circumference).to_s, sy1.to_s, (sx1 + circumference).to_s, sy0.to_s,
                  source, shifted)
          shifted
        end
        system!("gdalbuildvrt", "-q", vrt, *wrapped_sources)
        system!("gdalwarp", "-q", "-overwrite", "-of", "ENVI", "-ot", "Byte", "-r", "bilinear",
                "-te", xmin.to_s, ymin.to_s, xmax.to_s, ymax.to_s, "-ts", width.to_s, height.to_s,
                vrt, reflectance_path)
        elevations = File.binread(elevation_path).unpack("e*").map do |radius|
          [[(radius - DataSet::DATUM_RADIUS_M).round, -32_768].max, 32_767].min
        end
        reflectance = File.binread(reflectance_path, width * height)
        raise "unexpected reflectance byte count" unless reflectance.bytesize == width * height
        histogram = Array.new(256, 0)
        reflectance.each_byte { |value| histogram[value] += 1 }
        percentiles = [0.01, 0.99].map do |p|
          target = (p * (reflectance.bytesize - 1)).round
          cumulative = 0
          histogram.index { |count| cumulative += count; cumulative > target }
        end
        minmax = elevations.minmax
        header = {
          "format_version" => 1, "title" => "Reduced LROC WAC GLD100/LOLA terrain and normalized reflectance",
          "width" => width, "height" => height, "longitude" => "0 degrees east, wrapping through 360",
          "latitude" => "+90 degrees north to -90 degrees south", "datum_radius_m" => DataSet::DATUM_RADIUS_M,
          "elevation_encoding" => "little-endian signed int16 metres relative to datum",
          "reflectance_encoding" => "uint8 normalized reflectance", "elevation_offset_range_m" => minmax,
          "reflectance_percentiles" => percentiles,
          "sources" => [File.basename(@dtm)] + sources.map { |path| File.basename(path) },
          "source_sha256" => ([@dtm] + sources).to_h { |path| [File.basename(path), Digest::SHA256.file(path).hexdigest] }
        }
        write_asset(header, elevations.pack("s<*"), reflectance)
      end
      puts "Wrote #{@output}"
    end

    private

    def write_asset(header, elevations, reflectance)
      encoded = JSON.generate(header)
      FileUtils.mkdir_p(File.dirname(@output))
      Zlib::GzipWriter.open(@output) do |gzip|
        gzip.mtime = 0
        gzip.write(MAGIC)
        gzip.write([encoded.bytesize].pack("L<"))
        gzip.write(encoded)
        gzip.write(elevations)
        gzip.write(reflectance)
      end
    end

    def capture(*command)
      output, status = Open3.capture2e(*command)
      raise "command failed: #{command.join(" ")}\n#{output}" unless status.success?
      output
    end

    def system!(*command)
      success = system(*command)
      raise "command failed: #{command.join(" ")}" unless success
    end
  end
end