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

defmodule Microwaveprop.Weather.Grib2.Extractor do
  @moduledoc false

  alias Microwaveprop.Weather.Grib2.ComplexPacking
  alias Microwaveprop.Weather.Grib2.LambertConformal
  alias Microwaveprop.Weather.Grib2.Section
  alias Microwaveprop.Weather.Grib2.SimplePacking

  @doc """
  Extract weather values from a GRIB2 binary blob at the given lat/lon.

  Returns `{:ok, %{"VAR:LEVEL" => float}}` or `{:error, term}`.
  """
  def extract_points(binary, lat, lon) do
    messages = split_messages(binary)

    results =
      Enum.reduce_while(messages, {:ok, %{}}, fn msg, {:ok, acc} ->
        case extract_single(msg, lat, lon) do
          {:ok, key, value} -> {:cont, {:ok, Map.put(acc, key, value)}}
          {:error, :outside_grid} -> {:halt, {:error, :outside_grid}}
          {:error, reason} -> {:halt, {:error, reason}}
        end
      end)

    results
  end

  @doc """
  Split a binary blob into individual GRIB2 messages by scanning for "GRIB" magic bytes
  and reading the total length from the indicator section.
  """
  def split_messages(binary), do: split_messages(binary, [])

  @doc """
  Compute linear index from grid (i, j) and scan mode.

  For HRRR scan_mode=64 (j positive, i positive, i consecutive): index = j * nx + i
  """
  def linear_index({i, j}, nx, _scan_mode) do
    j * nx + i
  end

  @doc """
  Extract weather values from a GRIB2 binary blob at multiple lat/lon points.

  For each message in the binary, converts all points to grid indices, then
  batch-extracts values. Points outside the grid are silently skipped.

  Returns `{:ok, %{{lat, lon} => %{"VAR:LEVEL" => float}}}` or `{:error, term}`.
  """
  def extract_grid(binary, points) when is_list(points) do
    messages = split_messages(binary)

    Enum.reduce_while(messages, {:ok, %{}}, fn msg, {:ok, acc} ->
      case extract_grid_single(msg, points) do
        {:ok, key, point_values} ->
          {:cont, {:ok, merge_point_values(acc, key, point_values)}}

        {:error, reason} ->
          {:halt, {:error, reason}}
      end
    end)
  end

  # --- Private ---

  defp split_messages(<<>>, acc), do: Enum.reverse(acc)

  defp split_messages(<<"GRIB", _::32, total_length::64-big, _rest::binary>> = binary, acc) do
    if total_length > byte_size(binary) do
      Enum.reverse(acc)
    else
      msg = binary_part(binary, 0, total_length)
      remaining = binary_part(binary, total_length, byte_size(binary) - total_length)
      split_messages(remaining, [msg | acc])
    end
  end

  defp split_messages(<<_::8, rest::binary>>, acc) do
    # Skip non-GRIB bytes (padding between messages)
    split_messages(rest, acc)
  end

  defp extract_single(msg, lat, lon) do
    with {:ok, parsed} <- Section.parse_message(msg),
         %{grid_params: grid, product: prod, packing_params: packing, data: data} = parsed,
         key = "#{prod.var}:#{prod.level}",
         {:ok, {i, j}} <- LambertConformal.to_grid_index(grid, lat, lon),
         index = linear_index({i, j}, grid.nx, grid.scan_mode),
         {:ok, value} <- unpack_value(packing, data, index) do
      {:ok, key, value}
    end
  rescue
    e -> {:error, "GRIB2 extraction failed: #{inspect(e)}"}
  end

  defp extract_grid_single(msg, points) do
    with {:ok, parsed} <- Section.parse_message(msg) do
      %{grid_params: grid, product: prod, packing_params: packing, data: data} = parsed
      key = "#{prod.var}:#{prod.level}"
      point_indices = resolve_grid_indices(grid, points)
      batch_extract_points(packing, data, key, point_indices)
    end
  rescue
    e -> {:error, "GRIB2 grid extraction failed: #{inspect(e)}"}
  end

  defp resolve_grid_indices(grid, points) do
    points
    |> Enum.reduce([], fn {lat, lon} = point, acc ->
      case LambertConformal.to_grid_index(grid, lat, lon) do
        {:ok, {i, j}} ->
          index = linear_index({i, j}, grid.nx, grid.scan_mode)
          [{point, index} | acc]

        {:error, :outside_grid} ->
          acc
      end
    end)
    |> Enum.reverse()
  end

  defp batch_extract_points(_packing, _data, key, []), do: {:ok, key, []}

  defp batch_extract_points(packing, data, key, point_indices) do
    unique_indices = point_indices |> Enum.map(&elem(&1, 1)) |> Enum.uniq()

    case unpack_values(packing, data, unique_indices) do
      {:ok, values_map} ->
        point_values = Enum.map(point_indices, fn {point, index} -> {point, Map.fetch!(values_map, index)} end)
        {:ok, key, point_values}

      {:error, _} = err ->
        err
    end
  end

  defp merge_point_values(acc, key, point_values) do
    Enum.reduce(point_values, acc, fn {point, value}, a ->
      existing = Map.get(a, point, %{})
      Map.put(a, point, Map.put(existing, key, value))
    end)
  end

  defp unpack_value(%{template: 3} = params, data, index) do
    ComplexPacking.extract_value(params, data, index)
  end

  defp unpack_value(params, data, index) do
    SimplePacking.extract_value(params, data, index)
  end

  defp unpack_values(%{template: 3} = params, data, indices) do
    ComplexPacking.extract_values(params, data, indices)
  end

  defp unpack_values(params, data, indices) do
    SimplePacking.extract_values(params, data, indices)
  end
end