prop/lib/microwaveprop/weather/grib2/wgrib2.ex

defmodule Microwaveprop.Weather.Grib2.Wgrib2 do
  @moduledoc """
  Fast GRIB2 grid extraction using the wgrib2 binary.

  Uses wgrib2's `-lola` option to interpolate HRRR Lambert Conformal data
  onto a regular lat-lon grid, outputting IEEE 754 binary floats.
  Falls back to the pure-Elixir decoder if wgrib2 is not available.
  """

  require Logger

  @wgrib2_path System.find_executable("wgrib2")
  @undefined_value 9.999e20

  @doc """
  Extract values from a GRIB2 binary at a regular lat-lon grid.

  Takes the raw GRIB2 binary data, a regex pattern to match desired messages
  (e.g. ":(TMP|DPT|PRES):"), and grid specification.

  Returns `{:ok, %{{lat, lon} => %{"VAR:LEVEL" => float}}}` or `{:error, reason}`.
  """
  def extract_grid(grib_binary, match_pattern, grid_spec) do
    if available?() do
      extract_with_wgrib2(grib_binary, match_pattern, grid_spec)
    else
      {:error, :wgrib2_not_available}
    end
  end

  @doc "Check if wgrib2 is available on the system."
  def available? do
    @wgrib2_path != nil
  end

  defp extract_with_wgrib2(grib_binary, match_pattern, grid_spec) do
    %{
      lon_start: lon_start,
      lon_count: lon_count,
      lon_step: lon_step,
      lat_start: lat_start,
      lat_count: lat_count,
      lat_step: lat_step
    } = grid_spec

    # wgrib2 uses 0-360 longitude convention
    wgrib2_lon_start = normalize_lon(lon_start)

    lon_spec = "#{wgrib2_lon_start}:#{lon_count}:#{lon_step}"
    lat_spec = "#{lat_start}:#{lat_count}:#{lat_step}"

    # Write GRIB to temp file
    tmp_grib = Path.join(System.tmp_dir!(), "hrrr_#{System.unique_integer([:positive])}.grib2")
    tmp_bin = tmp_grib <> ".lola.bin"

    try do
      File.write!(tmp_grib, grib_binary)

      # Run wgrib2: match desired messages, extract to regular lat-lon grid as binary
      args = [
        tmp_grib,
        "-match",
        match_pattern,
        "-lola",
        lon_spec,
        lat_spec,
        tmp_bin,
        "bin"
      ]

      case System.cmd(@wgrib2_path, args, stderr_to_stdout: true) do
        {output, 0} ->
          # Parse message inventory from stdout to know which vars were extracted
          messages = parse_wgrib2_inventory(output)

          case File.read(tmp_bin) do
            {:ok, bin_data} ->
              parse_lola_binary(bin_data, messages, grid_spec)

            {:error, :enoent} ->
              # No output file means no matching messages
              {:ok, %{}}

            {:error, reason} ->
              {:error, "Failed to read wgrib2 output: #{inspect(reason)}"}
          end

        {output, exit_code} ->
          {:error, "wgrib2 failed (exit #{exit_code}): #{String.slice(output, 0, 200)}"}
      end
    after
      File.rm(tmp_grib)
      File.rm(tmp_bin)
    end
  end

  defp parse_wgrib2_inventory(output) do
    output
    |> String.split("\n")
    |> Enum.flat_map(fn line ->
      case String.split(line, ":", parts: 8) do
        [_n, _offset, _date, var, level | _] ->
          [%{var: var, level: level}]

        _ ->
          []
      end
    end)
  end

  defp parse_lola_binary(bin_data, messages, grid_spec) do
    %{lon_count: nx, lat_count: ny, lon_start: lon_start, lon_step: lon_step, lat_start: lat_start, lat_step: lat_step} =
      grid_spec

    points_per_message = nx * ny
    bytes_per_message = points_per_message * 4

    result =
      messages
      |> Enum.with_index()
      |> Enum.reduce(%{}, fn {msg, msg_idx}, acc ->
        offset = msg_idx * bytes_per_message
        key = "#{msg.var}:#{msg.level}"

        if offset + bytes_per_message <= byte_size(bin_data) do
          chunk = binary_part(bin_data, offset, bytes_per_message)
          merge_message_values(acc, key, chunk, nx, ny, lon_start, lon_step, lat_start, lat_step)
        else
          acc
        end
      end)

    {:ok, result}
  end

  defp merge_message_values(acc, key, chunk, nx, ny, lon_start, lon_step, lat_start, lat_step) do
    # Binary is row-major: lat varies slowest, lon varies fastest
    # Each value is a 32-bit IEEE 754 little-endian float
    Enum.reduce(0..(ny - 1), acc, fn j, acc_outer ->
      lat = Float.round(lat_start + j * lat_step, 3)

      Enum.reduce(0..(nx - 1), acc_outer, fn i, acc_inner ->
        offset = (j * nx + i) * 4
        <<_::binary-size(offset), value::float-little-32, _::binary>> = chunk

        if value > @undefined_value / 2 do
          # Skip undefined values (ocean/outside-domain points)
          acc_inner
        else
          lon = Float.round(denormalize_lon(lon_start + i * lon_step), 3)
          point = {lat, lon}
          existing = Map.get(acc_inner, point, %{})
          Map.put(acc_inner, point, Map.put(existing, key, value))
        end
      end)
    end)
  end

  # Convert -125.0 to 235.0 for wgrib2
  defp normalize_lon(lon) when lon < 0, do: lon + 360.0
  defp normalize_lon(lon), do: lon

  # Convert back from 0-360 to -180..180
  defp denormalize_lon(lon) when lon > 180.0, do: lon - 360.0
  defp denormalize_lon(lon), do: lon
end