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.
  """
  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}}`.
  """
  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.
  """
  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