prop/lib/microwaveprop/weather/hrrr_native_client.ex

defmodule Microwaveprop.Weather.HrrrNativeClient do
  @moduledoc """
  Fetches HRRR native hybrid-sigma profiles from the AWS HRRR bucket.

  This is the companion to `HrrrClient`, which works against the
  surface and 25 hPa pressure-level products. The native file
  (`wrfnatf00.grib2`) carries all variables on the 50 hybrid-sigma
  levels native to the HRRR model grid. Vertical spacing near the
  surface is ~10-50 m instead of the ~250 m the pressure-level
  product gives us — crucial for resolving the ducts and
  boundary-layer inversions discussed in
  `docs/plans/2026-04-09-propagation-modeling-improvements.md`.

  ## Design: batch, not per-point

  Each native-level HRRR file is ~566 MB. Essential variables (TMP,
  SPFH, HGT, UGRD, VGRD, TKE, PRES on all 50 hybrid levels) span
  ~530 MB of that file. Per-point on-demand fetching is not viable.
  Instead, the worker fetches the file once per `(date, hour)`,
  extracts native profiles for every point of interest in one pass,
  and bulk-inserts them.

  See `docs/research/hrrr_native_levels.md` for the full analysis.
  """

  alias Microwaveprop.Weather.HrrrClient

  @native_levels 1..50

  @native_variables ~w(TMP SPFH HGT UGRD VGRD TKE PRES)

  @hrrr_base_default "https://noaa-hrrr-bdp-pds.s3.amazonaws.com"

  defp hrrr_base, do: Application.get_env(:microwaveprop, :hrrr_base_url, @hrrr_base_default)

  @doc "Number of native hybrid-sigma levels in HRRR (currently 50)."
  def native_level_count, do: Enum.count(@native_levels)

  @doc "The seven essential variables we extract on every native level."
  def native_variables, do: @native_variables

  @doc """
  The list of `%{var:, level:}` messages we extract from every native
  file. 7 vars × 50 levels = 350 messages, matching the spike in
  Task 1.1.
  """
  def native_messages do
    for level <- @native_levels, var <- @native_variables do
      %{var: var, level: "#{level} hybrid level"}
    end
  end

  @doc """
  Builds the AWS S3 URL for a native-level HRRR grib2 file.

  ## Examples

      iex> Microwaveprop.Weather.HrrrNativeClient.hrrr_native_url(~D[2026-04-09], 12)
      "https://noaa-hrrr-bdp-pds.s3.amazonaws.com/hrrr.20260409/conus/hrrr.t12z.wrfnatf00.grib2"
  """
  def hrrr_native_url(date, hour, forecast_hour \\ 0) do
    date_str = Calendar.strftime(date, "%Y%m%d")
    hour_str = hour |> Integer.to_string() |> String.pad_leading(2, "0")
    fh_str = forecast_hour |> Integer.to_string() |> String.pad_leading(2, "0")
    "#{hrrr_base()}/hrrr.#{date_str}/conus/hrrr.t#{hour_str}z.wrfnatf#{fh_str}.grib2"
  end

  @doc """
  Converts a parsed `%{"VAR:level" => value}` map into a
  `HrrrNativeProfile`-shaped map with level arrays sorted by ascending
  hybrid level (level 1 = surface).

  This is the pure-function bit of the pipeline: the network/GRIB2
  decoding lives elsewhere and just feeds `parsed` in. Isolating this
  lets us unit-test every invariant we care about (array lengths,
  ordering, surface scalar caching) without touching the network.
  """
  def build_native_profile(parsed) when is_map(parsed) do
    levels =
      Enum.map(@native_levels, fn level ->
        level_str = "#{level} hybrid level"

        %{
          level: level,
          hgt: parsed["HGT:#{level_str}"],
          tmp: parsed["TMP:#{level_str}"],
          spfh: parsed["SPFH:#{level_str}"],
          pres: parsed["PRES:#{level_str}"],
          ugrd: parsed["UGRD:#{level_str}"],
          vgrd: parsed["VGRD:#{level_str}"],
          tke: parsed["TKE:#{level_str}"]
        }
      end)
      |> Enum.reject(fn %{hgt: hgt, tmp: tmp} -> is_nil(hgt) or is_nil(tmp) end)
      |> Enum.sort_by(& &1.hgt)

    level_count = length(levels)

    %{
      level_count: level_count,
      heights_m: Enum.map(levels, & &1.hgt),
      temp_k: Enum.map(levels, & &1.tmp),
      spfh: Enum.map(levels, & &1.spfh),
      pressure_pa: Enum.map(levels, & &1.pres),
      u_wind_ms: Enum.map(levels, & &1.ugrd),
      v_wind_ms: Enum.map(levels, & &1.vgrd),
      tke_m2s2: Enum.map(levels, & &1.tke),
      surface_temp_k: parsed["TMP:surface"] || List.first(levels) |> safe_get(:tmp),
      surface_spfh: parsed["SPFH:2 m above ground"] || List.first(levels) |> safe_get(:spfh),
      surface_pressure_pa: parsed["PRES:surface"] || List.first(levels) |> safe_get(:pres)
    }
  end

  defp safe_get(nil, _key), do: nil
  defp safe_get(map, key), do: Map.get(map, key)

  @doc """
  Returns the list of byte ranges to download for the essentials in
  one native HRRR file. Used by the (still-to-be-built) grid worker.
  Wraps `HrrrClient.byte_ranges_for_messages/2` with our native
  message list so callers don't have to know both.
  """
  def essential_byte_ranges(idx_entries) do
    HrrrClient.byte_ranges_for_messages(idx_entries, native_messages())
  end

  @doc """
  Extract native profiles for a list of `{lat, lon}` points from a
  GRIB2 binary, using wgrib2 for speed.

  Under the hood this uses `-lola` on a bounding-box subgrid that
  covers all the requested points, then does nearest-neighbor lookup
  per point. Falls back to the pure-Elixir decoder if wgrib2 is not
  available (expect ~70s per point in that case).

  Returns `%{{lat, lon} => native_profile_map}` where each profile
  map has the shape expected by `HrrrNativeProfile.changeset/2`.
  """
  def extract_native_profiles(grib_binary, points) when is_list(points) do
    alias Microwaveprop.Weather.Grib2.Wgrib2

    # Match only hybrid-level messages — the simple var-name pattern
    # also hits surface/2m/10m messages which misalign the binary output.
    match_pattern = ":(#{Enum.join(@native_variables, "|")}):.*hybrid level:"

    if Wgrib2.available?() do
      grid_spec = bounding_grid(points)

      case Wgrib2.extract_grid(grib_binary, match_pattern, grid_spec) do
        {:ok, grid_data} ->
          # Nearest-neighbor lookup: for each requested point, find
          # the grid cell with the smallest (lat, lon) distance.
          result =
            Map.new(points, fn {lat, lon} ->
              nearest = nearest_grid_cell(grid_data, lat, lon)
              profile = if nearest, do: build_native_profile(nearest), else: %{level_count: 0}
              {{lat, lon}, profile}
            end)

          {:ok, result}

        error ->
          error
      end
    else
      # Fallback: pure Elixir (slow)
      alias Microwaveprop.Weather.Grib2.Extractor

      case Extractor.extract_grid(grib_binary, points) do
        {:ok, grid_data} ->
          result = Map.new(grid_data, fn {pt, parsed} -> {pt, build_native_profile(parsed)} end)
          {:ok, result}

        error ->
          error
      end
    end
  end

  # Build a -lola grid spec that covers all points with 0.03° padding.
  @grid_step 0.03
  defp bounding_grid(points) do
    lats = Enum.map(points, &elem(&1, 0))
    lons = Enum.map(points, &elem(&1, 1))

    lat_min = Enum.min(lats) - 0.1
    lat_max = Enum.max(lats) + 0.1
    lon_min = Enum.min(lons) - 0.1
    lon_max = Enum.max(lons) + 0.1

    lon_count = max(trunc(Float.ceil((lon_max - lon_min) / @grid_step)), 2)
    lat_count = max(trunc(Float.ceil((lat_max - lat_min) / @grid_step)), 2)

    %{
      lon_start: lon_min,
      lon_count: lon_count,
      lon_step: @grid_step,
      lat_start: lat_min,
      lat_count: lat_count,
      lat_step: @grid_step
    }
  end

  defp nearest_grid_cell(grid_data, lat, lon) do
    grid_data
    |> Enum.min_by(fn {{glat, glon}, _} ->
      :math.pow(glat - lat, 2) + :math.pow(glon - lon, 2)
    end, fn -> nil end)
    |> case do
      nil -> nil
      {_point, parsed} -> parsed
    end
  end
end