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

  require Logger

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

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

    Enum.reduce(messages, {:ok, %{}}, fn
      msg, {:ok, acc} ->
        case extract_single(msg, lat, lon) do
          {:ok, key, value} -> {:ok, Map.put(acc, key, value)}
          {:error, :outside_grid} -> {:error, :outside_grid}
          {:error, _reason} -> {:ok, acc}
        end

      _msg, error ->
        error
    end)
  end

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

  Takes a binary and a list of `{lat, lon}` tuples.
  Returns `{:ok, %{{lat, lon} => %{"VAR:LEVEL" => float}}}` or `{:error, term}`.
  """
  @spec extract_grid(binary(), [{float(), float()}]) ::
          {:ok, %{{float(), float()} => %{String.t() => float()}}} | {:error, term()}
  def extract_grid(binary, points) do
    messages = split_messages(binary)

    # Pre-compute grid indices once from the first message's grid params.
    # All messages in the same GRIB2 file share the same grid.
    grid_indices =
      case messages do
        [first_msg | _] ->
          case Section.parse_message(first_msg) do
            {:ok, %{grid_params: grid}} -> precompute_indices(grid, points)
            _ -> []
          end

        [] ->
          []
      end

    if grid_indices == [] do
      {:ok, %{}}
    else
      messages
      |> reduce_messages(grid_indices, init_grid(points))
      |> filter_empty_grid()
    end
  end

  defp reduce_messages(messages, grid_indices, initial_grid) do
    Enum.reduce(messages, {:ok, initial_grid}, fn msg, {:ok, acc} ->
      case extract_grid_with_indices(msg, grid_indices) do
        {:ok, key, values} ->
          {:ok, merge_grid_values(values, acc, key)}

        {:error, _reason} ->
          {:ok, acc}
      end
    end)
  end

  defp filter_empty_grid({:ok, grid}) do
    {:ok, Map.reject(grid, fn {_point, values} -> values == %{} end)}
  end

  defp filter_empty_grid(error), do: error

  @doc """
  Split a binary blob into individual GRIB2 messages by scanning for "GRIB" magic bytes
  and reading the total length from the indicator section.
  """
  @spec split_messages(binary()) :: [binary()]
  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
  """
  @spec linear_index({non_neg_integer(), non_neg_integer()}, pos_integer(), non_neg_integer()) ::
          number()
  def linear_index({i, j}, nx, _scan_mode) do
    j * nx + i
  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 init_grid(points) do
    Map.new(points, fn point -> {point, %{}} end)
  end

  # Pre-compute Lambert Conformal grid indices for all points (done once, reused per message)
  defp precompute_indices(grid, points) do
    Enum.flat_map(points, fn {lat, lon} = point ->
      case LambertConformal.to_grid_index(grid, lat, lon) do
        {:ok, {i, j}} ->
          [{point, linear_index({i, j}, grid.nx, grid.scan_mode)}]

        {:error, :outside_grid} ->
          Logger.warning("Extractor: point #{lat},#{lon} outside Lambert grid bounds")
          []
      end
    end)
  end

  # Extract values for all pre-computed indices from a single GRIB message
  defp extract_grid_with_indices(msg, grid_indices) do
    with {:ok, parsed} <- Section.parse_message(msg) do
      %{product: prod, packing_params: packing, data: data} = parsed
      key = "#{prod.var}:#{prod.level}"

      # For complex packing, decode_all is called once; for simple, each index is independent
      unique_indices = grid_indices |> Enum.map(&elem(&1, 1)) |> Enum.uniq()

      case batch_unpack(packing, data, unique_indices) do
        {:ok, index_values} ->
          {:ok, key, resolve_grid_values(grid_indices, index_values)}

        {:error, reason} ->
          {:error, reason}
      end
    end
  rescue
    e -> {:error, "GRIB2 grid extraction failed: #{inspect(e)}"}
  end

  defp merge_grid_values(values, grid, key) do
    Enum.reduce(values, grid, fn {point, value}, acc ->
      Map.update!(acc, point, &Map.put(&1, key, value))
    end)
  end

  defp resolve_grid_values(grid_indices, index_values) do
    Enum.flat_map(grid_indices, fn {point, index} ->
      case Map.get(index_values, index) do
        nil -> []
        value -> [{point, value}]
      end
    end)
  end

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

  defp batch_unpack(params, data, indices) do
    SimplePacking.extract_values(params, data, indices)
  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 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
end