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

  @typep grid_spec :: %{
           lon_start: float(),
           lon_count: pos_integer(),
           lon_step: float(),
           lat_start: float(),
           lat_count: pos_integer(),
           lat_step: float()
         }

  @typep point_grid :: %{{float(), float()} => %{String.t() => float()}}

  @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}`.
  """
  @spec extract_grid(binary(), String.t(), grid_spec()) :: {:ok, point_grid()} | {:error, term()}
  def extract_grid(grib_binary, match_pattern, grid_spec) do
    Microwaveprop.Instrument.span(
      [:wgrib2, :extract_grid],
      %{bytes: byte_size(grib_binary)},
      fn ->
        if available?() do
          extract_with_wgrib2(grib_binary, match_pattern, grid_spec)
        else
          {:error, :wgrib2_not_available}
        end
      end
    )
  end

  @doc """
  Like `extract_grid/3`, but takes a path to a GRIB2 file on disk instead
  of a binary. Avoids loading the entire file into memory — useful for
  large native-level HRRR files (~530 MB).
  """
  @spec extract_grid_from_file(Path.t(), String.t(), grid_spec()) ::
          {:ok, point_grid()} | {:error, term()}
  def extract_grid_from_file(grib_path, match_pattern, grid_spec) do
    Microwaveprop.Instrument.span([:wgrib2, :extract_grid_from_file], %{}, fn ->
      if available?() do
        extract_file_with_wgrib2(grib_path, match_pattern, grid_spec)
      else
        {:error, :wgrib2_not_available}
      end
    end)
  end

  @doc """
  Like `extract_grid_from_file/3`, but instead of building the full
  `%{{lat, lon} => %{"VAR:LEVEL" => float}}` map (which can be 1+ GB for
  200+ messages × 95k cells), processes each grid cell through a reducer
  function and keeps only the reduced output.

  The reducer receives `%{"VAR:LEVEL" => float}` for one cell and returns
  an arbitrary term. The result is `%{{lat, lon} => reducer_output}`.

  Peak memory: binary output (~76 MB for 200 msgs × 95k cells) + one
  cell's worth of data at a time + output map (~10 MB of scalars).
  """
  @spec extract_grid_from_file_mapped(
          Path.t(),
          String.t(),
          grid_spec(),
          (%{String.t() => float()} -> term())
        ) :: {:ok, %{{float(), float()} => term()}} | {:error, term()}
  def extract_grid_from_file_mapped(grib_path, match_pattern, grid_spec, cell_reducer) do
    Microwaveprop.Instrument.span([:wgrib2, :extract_grid_from_file_mapped], %{}, fn ->
      if available?() do
        extract_file_mapped_with_wgrib2(grib_path, match_pattern, grid_spec, cell_reducer)
      else
        {:error, :wgrib2_not_available}
      end
    end)
  end

  @doc "Check if wgrib2 is available on the system."
  @spec available?() :: boolean()
  def available?, do: wgrib2_path() != nil

  defp wgrib2_path, do: System.find_executable("wgrib2")

  defp extract_file_with_wgrib2(grib_path, 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_lon_start = normalize_lon(lon_start)
    lon_spec = "#{wgrib2_lon_start}:#{lon_count}:#{lon_step}"
    lat_spec = "#{lat_start}:#{lat_count}:#{lat_step}"

    tmp_bin = grib_path <> ".lola.bin"

    try do
      args = [grib_path, "-match", match_pattern, "-lola", lon_spec, lat_spec, tmp_bin, "bin"]

      case System.cmd(wgrib2_path(), args, stderr_to_stdout: true) do
        {output, 0} ->
          messages = parse_wgrib2_inventory(output)

          case File.read(tmp_bin) do
            {:ok, bin_data} -> parse_lola_binary(bin_data, messages, grid_spec)
            {:error, :enoent} -> {: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_bin)
    end
  end

  defp extract_file_mapped_with_wgrib2(grib_path, match_pattern, grid_spec, cell_reducer) 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_lon_start = normalize_lon(lon_start)
    lon_spec = "#{wgrib2_lon_start}:#{lon_count}:#{lon_step}"
    lat_spec = "#{lat_start}:#{lat_count}:#{lat_step}"

    tmp_bin = grib_path <> ".lola.mapped.bin"

    try do
      args = [grib_path, "-match", match_pattern, "-lola", lon_spec, lat_spec, tmp_bin, "bin"]

      case System.cmd(wgrib2_path(), args, stderr_to_stdout: true) do
        {output, 0} ->
          messages = parse_wgrib2_inventory(output)

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

            {:error, :enoent} ->
              {: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_bin)
    end
  end

  # Cell-by-cell variant of parse_lola_binary. Instead of building the full
  # %{{lat,lon} => %{"VAR:LEVEL" => val}} map, iterates grid cells and for
  # each cell extracts its values from every message, passes the per-cell
  # map to the reducer, and keeps only the reduced output.
  defp parse_lola_binary_mapped(bin_data, messages, grid_spec, cell_reducer) 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
    record_overhead = 8
    stride = bytes_per_message + record_overhead

    # Pre-compute message metadata: {key, data_offset} for each message
    msg_meta =
      messages
      |> Enum.with_index()
      |> Enum.map(fn {msg, idx} ->
        {"#{msg.var}:#{msg.level}", idx * stride + 4}
      end)
      |> Enum.filter(fn {_key, offset} -> offset + bytes_per_message <= byte_size(bin_data) end)

    # Iterate each grid cell, extract all message values, reduce
    result =
      Enum.reduce(0..(ny - 1), %{}, fn j, acc_outer ->
        lat = Float.round(lat_start + j * lat_step, 3)

        Enum.reduce(0..(nx - 1), acc_outer, fn i, acc_inner ->
          cell_offset = (j * nx + i) * 4
          cell_data = extract_cell_values(bin_data, msg_meta, cell_offset)
          reduce_mapped_cell(cell_data, lat, lon_start + i * lon_step, cell_reducer, acc_inner)
        end)
      end)

    {:ok, result}
  end

  defp reduce_mapped_cell(cell_data, _lat, _raw_lon, _cell_reducer, acc) when cell_data == %{}, do: acc

  defp reduce_mapped_cell(cell_data, lat, raw_lon, cell_reducer, acc) do
    lon = Float.round(denormalize_lon(raw_lon), 3)
    reduced = cell_reducer.(cell_data)
    Map.put(acc, {lat, lon}, reduced)
  end

  defp extract_cell_values(bin_data, msg_meta, cell_offset) do
    Enum.reduce(msg_meta, %{}, fn {key, data_offset}, cell_acc ->
      offset = data_offset + cell_offset
      <<_::binary-size(^offset), value::float-little-32, _::binary>> = bin_data

      if value > @undefined_value / 2 do
        cell_acc
      else
        Map.put(cell_acc, key, value)
      end
    end)
  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, datetime: parse_wgrib2_date(date)}]

        _ ->
          Logger.warning("wgrib2: unparseable inventory line: #{String.slice(line, 0, 120)}")
          []
      end
    end)
  end

  # wgrib2 inventory date looks like "d=20250101000000" or "d=2025010100".
  # Returns a `DateTime` in UTC, or `nil` if it can't be parsed.
  defp parse_wgrib2_date("d=" <> digits), do: parse_wgrib2_date_digits(digits)
  defp parse_wgrib2_date(_), do: nil

  defp parse_wgrib2_date_digits(<<y::binary-4, m::binary-2, d::binary-2, h::binary-2, rest::binary>>) do
    {mi, s} =
      case rest do
        <<mm::binary-2, ss::binary-2>> -> {mm, ss}
        <<mm::binary-2>> -> {mm, "00"}
        _ -> {"00", "00"}
      end

    with {year, ""} <- Integer.parse(y),
         {month, ""} <- Integer.parse(m),
         {day, ""} <- Integer.parse(d),
         {hour, ""} <- Integer.parse(h),
         {minute, ""} <- Integer.parse(mi),
         {second, ""} <- Integer.parse(s),
         {:ok, naive} <- NaiveDateTime.new(year, month, day, hour, minute, second),
         {:ok, dt} <- DateTime.from_naive(naive, "Etc/UTC") do
      DateTime.truncate(dt, :second)
    else
      _ -> nil
    end
  end

  defp parse_wgrib2_date_digits(_), do: nil

  @doc """
  Like `extract_grid/3`, but returns a flat list of per-message results so
  the time dimension is preserved. Used when a single GRIB2 blob carries
  many timesteps.

  Each entry is `%{datetime: DateTime.t() | nil, var: String.t(),
  level: String.t(), values: %{{lat, lon} => float}}`.
  """
  @spec extract_grid_messages(binary(), String.t(), grid_spec()) ::
          {:ok,
           [
             %{
               datetime: DateTime.t() | nil,
               var: String.t(),
               level: String.t(),
               values: %{{float(), float()} => float()}
             }
           ]}
          | {:error, term()}
  def extract_grid_messages(grib_binary, match_pattern, grid_spec) do
    if available?() do
      tmp_grib = Path.join(System.tmp_dir!(), "wgrib2_#{System.unique_integer([:positive])}.grib2")

      try do
        File.write!(tmp_grib, grib_binary)
        extract_messages_with_wgrib2(tmp_grib, match_pattern, grid_spec)
      after
        File.rm(tmp_grib)
      end
    else
      {:error, :wgrib2_not_available}
    end
  end

  @doc """
  Like `extract_grid_messages/3`, but reads GRIB2 from a file on disk instead
  of a binary in memory. Saves 50-200 MB of heap per call for month-tile
  fetches by avoiding the intermediate `File.write!(tmp, binary)` step.
  """
  @spec extract_grid_messages_from_file(Path.t(), String.t(), grid_spec()) ::
          {:ok,
           [
             %{
               datetime: DateTime.t() | nil,
               var: String.t(),
               level: String.t(),
               values: %{{float(), float()} => float()}
             }
           ]}
          | {:error, term()}
  def extract_grid_messages_from_file(grib_path, match_pattern, grid_spec) do
    if available?() do
      extract_messages_with_wgrib2(grib_path, match_pattern, grid_spec)
    else
      {:error, :wgrib2_not_available}
    end
  end

  defp extract_messages_with_wgrib2(grib_path, 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_lon_start = normalize_lon(lon_start)
    lon_spec = "#{wgrib2_lon_start}:#{lon_count}:#{lon_step}"
    lat_spec = "#{lat_start}:#{lat_count}:#{lat_step}"

    tmp_bin = grib_path <> ".lola.#{System.unique_integer([:positive])}.bin"

    try do
      args = [grib_path, "-match", match_pattern, "-lola", lon_spec, lat_spec, tmp_bin, "bin"]

      case System.cmd(wgrib2_path(), args, stderr_to_stdout: true) do
        {output, 0} ->
          messages = parse_wgrib2_inventory(output)

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

            {:error, :enoent} ->
              {: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_bin)
    end
  end

  defp build_messages_per_message(bin_data, messages, grid_spec) do
    %{lon_count: nx, lat_count: ny} = grid_spec
    points_per_message = nx * ny
    bytes_per_message = points_per_message * 4

    # wgrib2 -lola ... bin writes Fortran unformatted binary: each
    # message is preceded by a 4-byte little-endian record length
    # and followed by a duplicate 4-byte record length. Same math as
    # parse_lola_binary/3 — see that function for context.
    record_overhead = 8
    stride = bytes_per_message + record_overhead

    messages
    |> Enum.with_index()
    |> Enum.flat_map(fn {msg, msg_idx} ->
      data_offset = msg_idx * stride + 4

      if data_offset + bytes_per_message <= byte_size(bin_data) do
        chunk = binary_part(bin_data, data_offset, bytes_per_message)
        values = extract_message_values(chunk, grid_spec)

        [Map.put(msg, :values, values)]
      else
        []
      end
    end)
  end

  defp extract_message_values(chunk, %{
         lon_count: nx,
         lat_count: ny,
         lon_start: lon_start,
         lon_step: lon_step,
         lat_start: lat_start,
         lat_step: lat_step
       }) do
    Enum.reduce(0..(ny - 1), %{}, fn j, outer ->
      lat = Float.round(lat_start + j * lat_step, 3)

      Enum.reduce(0..(nx - 1), outer, fn i, inner ->
        extract_grid_point(chunk, nx, j, i, lat, lon_start, lon_step, inner)
      end)
    end)
  end

  defp extract_grid_point(chunk, nx, j, i, lat, lon_start, lon_step, acc) do
    offset = (j * nx + i) * 4
    <<_::binary-size(^offset), value::float-little-32, _::binary>> = chunk

    if value > @undefined_value / 2 do
      acc
    else
      lon = Float.round(denormalize_lon(lon_start + i * lon_step), 3)
      Map.put(acc, {lat, lon}, value)
    end
  end

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

    points_per_message = nx * ny
    bytes_per_message = points_per_message * 4

    # wgrib2 -lola ... bin writes Fortran unformatted binary: each
    # message is preceded by a 4-byte little-endian record length
    # and followed by a duplicate 4-byte record length. So the
    # stride per message is 4 + data + 4 = data + 8.
    record_overhead = 8
    stride = bytes_per_message + record_overhead

    result =
      messages
      |> Enum.with_index()
      |> Enum.reduce(%{}, fn {msg, msg_idx}, acc ->
        # Skip the 4-byte header to reach the data
        data_offset = msg_idx * stride + 4
        key = "#{msg.var}:#{msg.level}"

        if data_offset + bytes_per_message <= byte_size(bin_data) do
          chunk = binary_part(bin_data, data_offset, bytes_per_message)
          merge_message_values(acc, key, chunk, grid_spec)
        else
          acc
        end
      end)

    {:ok, result}
  end

  defp merge_message_values(acc, key, chunk, 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

    # 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
        merge_keyed_grid_value(value, lat, lon_start + i * lon_step, acc_inner, key)
      end)
    end)
  end

  defp merge_keyed_grid_value(value, _lat, _raw_lon, acc, _key) when value > @undefined_value / 2 do
    acc
  end

  defp merge_keyed_grid_value(value, lat, raw_lon, acc, key) do
    lon = Float.round(denormalize_lon(raw_lon), 3)
    point = {lat, lon}
    existing = Map.get(acc, point, %{})
    Map.put(acc, point, Map.put(existing, key, value))
  end

  @doc """
  Extract values at specific `{lat, lon}` points from a GRIB2 file on disk
  using wgrib2 `-lon`. One file scan, text output, no binary grid — uses
  negligible BEAM memory regardless of point spread or message count.

  Returns `{:ok, %{{lat, lon} => %{"VAR:LEVEL" => float}}}` or `{:error, reason}`.
  """
  @spec extract_points_from_file(Path.t(), String.t(), [{float(), float()}]) ::
          {:ok, point_grid()} | {:error, term()}
  def extract_points_from_file(grib_path, match_pattern, points) when is_list(points) do
    if available?() do
      extract_points_with_wgrib2(grib_path, match_pattern, points)
    else
      {:error, :wgrib2_not_available}
    end
  end

  defp extract_points_with_wgrib2(grib_path, match_pattern, points) do
    # Build -lon args: -lon lon1 lat1 -lon lon2 lat2 ...
    lon_args =
      Enum.flat_map(points, fn {lat, lon} ->
        ["-lon", "#{normalize_lon(lon)}", "#{lat}"]
      end)

    args = [grib_path, "-s", "-match", match_pattern] ++ lon_args

    case System.cmd(wgrib2_path(), args, stderr_to_stdout: true) do
      {output, 0} ->
        {:ok, parse_lon_output(output, points)}

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

  # Parse wgrib2 -lon output. Each line looks like:
  # 1:0:d=2024092219:TMP:1 hybrid level:anl:lon=242.958,lat=32.938,val=306.5:lon=242.208,lat=33.604,val=302.1
  # Each -lon option appends a "lon=X,lat=Y,val=Z" segment.
  defp parse_lon_output(output, points) do
    output
    |> String.split("\n")
    |> Enum.reject(&(&1 == ""))
    |> Enum.reduce(%{}, fn line, acc ->
      parse_lon_line(line, acc, points)
    end)
  end

  defp parse_lon_line(line, acc, points) do
    parts = String.split(line, ":")

    case extract_var_level(parts) do
      [var, level] ->
        key = "#{var}:#{level}"

        Enum.reduce(parts, acc, fn segment, inner_acc ->
          merge_lon_val_segment(segment, inner_acc, key, points)
        end)

      _ ->
        acc
    end
  end

  defp merge_lon_val_segment(segment, acc, key, points) do
    case parse_lon_val_segment(segment) do
      {lat, lon, val} ->
        point = snap_to_nearest(lat, lon, points)

        if point do
          existing = Map.get(acc, point, %{})
          Map.put(acc, point, Map.put(existing, key, val))
        else
          acc
        end

      nil ->
        acc
    end
  end

  defp extract_var_level(parts) when length(parts) >= 5 do
    [Enum.at(parts, 3), Enum.at(parts, 4)]
  end

  defp extract_var_level(_), do: nil

  @doc false
  # Parse "lon=242.958,lat=32.938,val=306.5". Exposed (with @doc false)
  # so the integer-lon edge case (wgrib2 occasionally drops the trailing
  # `.0`) can be asserted directly.
  def parse_lon_val_segment(segment) do
    case Regex.run(~r/lon=([\d.]+),lat=([\d.]+),val=([\d.eE+-]+)/, segment) do
      [_, lon_str, lat_str, val_str] ->
        # `Float.parse/1` accepts both `"243"` and `"243.0"`, where
        # `String.to_float/1` only accepts the latter. wgrib2 normally
        # emits the trailing `.0` but a build or locale flip that drops
        # it would crash the entire chain step.
        with {lon_val, ""} <- Float.parse(lon_str),
             {lat, ""} <- Float.parse(lat_str),
             {val, ""} <- Float.parse(val_str) do
          {Float.round(lat, 3), Float.round(denormalize_lon(lon_val), 3), val}
        else
          _ -> nil
        end

      _ ->
        nil
    end
  end

  # Find the requested point nearest to the wgrib2-reported lat/lon
  # (wgrib2 snaps to nearest grid cell, so reported coords may differ slightly)
  defp snap_to_nearest(lat, lon, points) do
    Enum.min_by(points, fn {plat, plon} ->
      :math.pow(plat - lat, 2) + :math.pow(plon - lon, 2)
    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