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

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

  @doc """
  Extract a single value from GRIB2 complex-packed data with spatial differencing
  (Template 5.3) at the given grid index.

  Decodes all values since spatial differencing requires sequential access,
  then returns the value at the target index.
  """
  @spec extract_value(map(), binary(), non_neg_integer()) :: {:ok, float()} | {:error, term()}
  def extract_value(params, data, index) do
    if index < 0 or index >= params.num_data_points do
      {:error, :index_out_of_range}
    else
      case decode_all(params, data) do
        {:ok, values} -> {:ok, :array.get(index, values)}
        {:error, _} = err -> err
      end
    end
  end

  @doc """
  Extract multiple values from GRIB2 complex-packed data at the given grid indices.

  Decodes all values (since spatial differencing requires sequential access),
  then returns the values at the requested indices.

  Returns `{:ok, %{index => value}}`.
  """
  @spec extract_values(map(), binary(), [non_neg_integer()]) ::
          {:ok, %{non_neg_integer() => float()}} | {:error, term()}
  def extract_values(params, data, indices) do
    case decode_all(params, data) do
      {:ok, array} ->
        results = Map.new(indices, fn index -> {index, :array.get(index, array)} end)
        {:ok, results}

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

  @doc """
  Decode all values from complex-packed data with spatial differencing.

  Returns an Erlang array of floats for O(1) index access.
  """
  @spec decode_all(map(), binary()) :: {:ok, :array.array(float())} | {:error, String.t()}
  def decode_all(params, data) do
    %{
      reference_value: ref,
      binary_scale: e,
      decimal_scale: d,
      bits_per_value: nbits,
      num_groups: num_groups,
      ref_group_widths: ref_gw,
      nbits_group_widths: nbits_gw,
      ref_group_lengths: ref_gl,
      length_increment: len_inc,
      last_group_length: last_gl,
      nbits_group_lengths: nbits_gl,
      spatial_order: spatial_order,
      num_extra_octets: num_extra_octets
    } = params

    # Step 1: Extract spatial differencing initial values
    octets_per_val = num_extra_octets
    num_init_vals = spatial_order + 1

    {init_vals, rest} = extract_init_values(data, num_init_vals, octets_per_val)

    {spatial_init, [overall_min]} = Enum.split(init_vals, spatial_order)

    # Step 2: Extract group reference values (byte-padded per GRIB2 spec)
    {group_refs, rest} = extract_n_values_array(rest, num_groups, nbits)

    # Step 3: Extract group widths (byte-padded)
    {group_widths_arr, rest} = extract_n_values_array(rest, num_groups, nbits_gw)

    # Step 4: Extract group lengths (byte-padded)
    {group_lengths_arr, rest} = extract_n_values_array(rest, num_groups, nbits_gl)

    # Step 5: Build group info lists
    group_widths =
      for g <- 0..(num_groups - 1) do
        get_val(group_widths_arr, g) + ref_gw
      end

    group_lengths =
      for g <- 0..(num_groups - 1) do
        if g == num_groups - 1 do
          last_gl
        else
          get_val(group_lengths_arr, g) * len_inc + ref_gl
        end
      end

    group_refs_list =
      for g <- 0..(num_groups - 1) do
        get_val(group_refs, g)
      end

    # Step 6: Decode packed data for each group using bitstring operations
    raw_values = decode_groups_bitwise(rest, group_refs_list, group_widths, group_lengths)

    # Step 7: Apply spatial differencing in reverse
    undiffed = apply_spatial_differencing(spatial_init, overall_min, raw_values)

    # Step 8: Apply scaling formula: value = (R + X * 2^E) * 10^(-D)
    factor_2e = :math.pow(2, e)
    factor_10d = :math.pow(10, -d)

    result =
      undiffed
      |> :array.to_list()
      |> Enum.map(fn x -> (ref + x * factor_2e) * factor_10d end)
      |> :array.from_list()

    {:ok, result}
  rescue
    e in [MatchError, ArgumentError] ->
      {:error, "GRIB2 complex packing decode failed: #{inspect(e)}"}
  end

  # --- Private helpers ---

  defp extract_init_values(data, count, octets_per_val) do
    total_bytes = count * octets_per_val
    <<init_bytes::binary-size(^total_bytes), rest::binary>> = data

    values =
      for i <- 0..(count - 1) do
        offset = i * octets_per_val
        val_bytes = binary_part(init_bytes, offset, octets_per_val)
        decode_signed_big(val_bytes)
      end

    {values, rest}
  end

  defp decode_signed_big(bytes) do
    bits = byte_size(bytes) * 8
    <<sign::1, magnitude::size(^bits - 1)>> = bytes
    if sign == 1, do: -magnitude, else: magnitude
  end

  defp extract_n_values_array(data, _count, 0) do
    {nil, data}
  end

  defp extract_n_values_array(data, count, nbits) do
    total_bits = count * nbits
    total_bytes = div(total_bits + 7, 8)
    <<chunk::binary-size(^total_bytes), rest::binary>> = data

    vals = consume_bits_simple(<<chunk::binary>>, nbits, [])
    # consume_bits_simple prepends (reversed); reverse first, then take count to discard padding
    arr = vals |> Enum.reverse() |> Enum.take(count) |> :array.from_list()

    {arr, rest}
  end

  defp consume_bits_simple(bits, nbits, acc) when bit_size(bits) < nbits, do: acc

  defp consume_bits_simple(bits, nbits, acc) do
    <<val::size(^nbits)-unsigned-big, rest::bitstring>> = bits
    consume_bits_simple(rest, nbits, [val | acc])
  end

  defp get_val(nil, _index), do: 0
  defp get_val(arr, index), do: :array.get(index, arr)

  # Decode all groups from packed data using a bitstring cursor.
  # Returns an Erlang array of raw (pre-differencing) integer values.
  defp decode_groups_bitwise(data, group_refs, group_widths, group_lengths) do
    {all_values_reversed, _remaining_bits} =
      [group_refs, group_widths, group_lengths]
      |> Enum.zip()
      |> Enum.reduce({[], data}, fn {gref, width, length}, {acc, remaining} ->
        if width == 0 do
          # All values equal gref — prepend in reverse (same values, order irrelevant)
          new_acc = prepend_n(acc, gref, length)
          {new_acc, remaining}
        else
          {new_acc, rest} = consume_group_bits(remaining, width, gref, length, acc)
          {new_acc, rest}
        end
      end)

    all_values_reversed
    |> Enum.reverse()
    |> :array.from_list()
  end

  # Consume `count` values of `width` bits each from bitstring, prepending to acc.
  # Values are prepended in reverse order (last value first).
  # Handles trailing padding: if fewer bits remain than needed, pad with gref.
  defp consume_group_bits(bits, width, gref, count, acc) do
    total_bits = count * width
    available = bit_size(bits)

    if available >= total_bits do
      <<chunk::bitstring-size(^total_bits), rest::bitstring>> = bits
      new_acc = consume_bits(chunk, width, gref, acc)
      {new_acc, rest}
    else
      # Extract what we can, fill remainder with gref
      new_acc = consume_bits(bits, width, gref, acc)
      extracted = div(available, width)
      remaining = count - extracted
      new_acc = prepend_n(new_acc, gref, remaining)
      {new_acc, <<>>}
    end
  end

  defp consume_bits(bits, width, _gref, acc) when bit_size(bits) < width, do: acc

  defp consume_bits(bits, width, gref, acc) do
    <<val::size(^width)-unsigned-big, rest::bitstring>> = bits
    consume_bits(rest, width, gref, [val + gref | acc])
  end

  defp prepend_n(acc, _value, 0), do: acc
  defp prepend_n(acc, value, n), do: prepend_n([value | acc], value, n - 1)

  defp apply_spatial_differencing(spatial_init, overall_min, raw_values) do
    raw_list = :array.to_list(raw_values)

    case spatial_init do
      [ival1] ->
        {result_reversed, _} =
          Enum.reduce(tl(raw_list), {[ival1], ival1}, fn raw, {acc, prev} ->
            current = prev + raw + overall_min
            {[current | acc], current}
          end)

        :array.from_list(Enum.reverse(result_reversed))

      [ival1, ival2] ->
        {result_reversed, _, _} =
          raw_list
          |> Enum.drop(2)
          |> Enum.reduce({[ival2, ival1], ival2, ival1}, fn raw, {acc, prev1, prev2} ->
            current = raw + overall_min + 2 * prev1 - prev2
            {[current | acc], current, prev1}
          end)

        :array.from_list(Enum.reverse(result_reversed))
    end
  end
end