prop/lib/microwaveprop/propagation/pgrid.ex

defmodule Microwaveprop.Propagation.Pgrid do
  @moduledoc """
  Reader for the `.pgrid` profile format written by `prop-grid-rs`
  (`rust/prop_grid_rs/src/pgrid.rs`).

  Replaces the gzipped-MessagePack `.mp.gz` artifact. The format is a
  dense, **cell-major** `f32` record array on a fixed grid, so:

    * reading one cell is a single `:file.pread/3` of `n_fields * 4`
      bytes at a computed offset — no decompression, no term building.
      The `.mp.gz` path gunzipped, `Msgpax.unpack`ed and atomized the
      *entire* 95k-cell file to answer one point lookup, on every `/map`
      click and every Skew-T load.
    * reading a viewport is one contiguous `pread` per grid row.

  Layout (little-endian) — see the Rust module for the authoritative
  description:

      magic       4   "PGRD"
      version     1   0x01
      flags       1   bit0: 0 = hrrr, 1 = hrdps
      n_fields    2   u16
      valid_time  8   i64 unix seconds
      lat_start   8   f64
      lon_start   8   f64
      lat_step    8   f64
      lon_step    8   f64
      n_rows      2   u16
      n_cols      2   u16
      field_table n_fields * 32 (NUL-padded ASCII)
      body        n_rows*n_cols*n_fields * 4  f32, cell-major

  `NaN` is the missing-value sentinel. Erlang's float binary match
  rejects NaN, so `decode_f32/1` maps it to `nil` — which is exactly
  what the `.mp.gz` reader produced for an absent key.

  Callers get the same map shape `ProfilesFile.read_point/3` always
  returned: atom keys, a `:profile` list of per-level maps, and a
  nested `:commercial_link_degradation` map.
  """

  alias Microwaveprop.Propagation.Grid

  @magic "PGRD"
  @version 1
  @field_name_len 32
  @fixed_header_len 4 + 1 + 1 + 2 + 8 + 8 + 8 + 8 + 8 + 2 + 2

  # On-disk field name => the atom key callers expect.
  #
  # Written as literal atoms so they exist at compile time. The `.mp.gz`
  # reader used `String.to_existing_atom/1` behind a whitelist, which
  # only worked as long as some *other* module happened to have mentioned
  # each atom — `:wind_v` in particular did not, so that approach raised
  # at runtime. An explicit map removes the hidden dependency, and a
  # field the reader doesn't know (e.g. one added later on the Rust side)
  # is simply skipped instead of crashing or growing the atom table.
  @field_atoms %{
    "surface_temp_c" => :surface_temp_c,
    "surface_dewpoint_c" => :surface_dewpoint_c,
    "surface_pressure_mb" => :surface_pressure_mb,
    "hpbl_m" => :hpbl_m,
    "pwat_mm" => :pwat_mm,
    "wind_u" => :wind_u,
    "wind_v" => :wind_v,
    "cloud_cover_pct" => :cloud_cover_pct,
    "precip_mm" => :precip_mm,
    "native_min_gradient" => :native_min_gradient,
    "best_duct_freq_ghz" => :best_duct_freq_ghz,
    "max_duct_thickness_m" => :max_duct_thickness_m,
    "duct_count" => :duct_count,
    "nexrad_max_reflectivity_dbz" => :nexrad_max_reflectivity_dbz,
    "surface_refractivity" => :surface_refractivity,
    "min_refractivity_gradient" => :min_refractivity_gradient
  }

  # Deliberately absent from @field_atoms: the commercial_* fields are
  # folded into the nested :commercial_link_degradation map, and the
  # per-level hght_m_/tmpc_/dwpc_ triples are rebuilt into :profile.

  # Surfaced as an integer, as the msgpack writer did.
  @integer_fields ~w(duct_count)

  defmodule Header do
    @moduledoc "Parsed `.pgrid` header."

    @type t :: %__MODULE__{}

    defstruct [
      :hrdps?,
      :valid_time,
      :n_fields,
      :fields,
      :field_index,
      :lat_start,
      :lon_start,
      :lat_step,
      :lon_step,
      :n_rows,
      :n_cols,
      :body_offset,
      :levels
    ]
  end

  @doc "Base directory the profile store lives under."
  @spec base_dir() :: String.t()
  def base_dir do
    Path.join(
      Application.get_env(:microwaveprop, :propagation_scores_dir, "/data/scores"),
      "profiles"
    )
  end

  @doc "Absolute path for the `.pgrid` file covering `valid_time`."
  @spec path_for(DateTime.t()) :: String.t()
  def path_for(%DateTime{} = valid_time) do
    iso = valid_time |> DateTime.truncate(:second) |> DateTime.to_iso8601()
    Path.join(base_dir(), "#{iso}.pgrid")
  end

  @doc "Whether a `.pgrid` exists for `valid_time`."
  @spec exists?(DateTime.t()) :: boolean()
  def exists?(%DateTime{} = valid_time), do: File.exists?(path_for(valid_time))

  @doc """
  Read a single cell's profile for `(valid_time, lat, lon)`.

  Opens the file, `pread`s the header and then one record. Returns `nil`
  when the file is missing or the point falls outside the grid.
  """
  @spec read_point(DateTime.t(), number(), number()) :: map() | nil
  def read_point(%DateTime{} = valid_time, lat, lon) do
    path = path_for(valid_time)

    with {:ok, fd} <- :file.open(path, [:read, :binary, :raw]),
         {:ok, header} <- read_header(fd) do
      result = read_cell(fd, header, lat, lon)
      :file.close(fd)
      result
    else
      _ -> nil
    end
  end

  defp read_cell(fd, header, lat, lon) do
    with cell when is_integer(cell) <- cell_index(header, lat, lon),
         size = header.n_fields * 4,
         {:ok, bin} <- :file.pread(fd, record_offset(header, cell), size),
         true <- byte_size(bin) == size do
      to_profile(header, bin)
    else
      _ -> nil
    end
  end

  @doc """
  Read every populated cell as `%{{lat, lon} => profile}`.

  Used for the cold `GridCache` warm in `Microwaveprop.Weather.Grid`.
  Cells with no surface temperature are skipped, so the result matches
  what the `.mp.gz` reader produced (it only ever contained scored cells).
  """
  @spec read(DateTime.t()) :: {:ok, %{{float(), float()} => map()}} | {:error, :enoent}
  def read(%DateTime{} = valid_time) do
    path = path_for(valid_time)

    with {:ok, raw} <- File.read(path),
         {:ok, header} <- parse_header(raw) do
      record_bytes = header.n_fields * 4

      grid =
        raw
        |> binary_part(header.body_offset, byte_size(raw) - header.body_offset)
        |> chunk_records(record_bytes)
        |> Enum.with_index()
        |> Enum.flat_map(&populated_cell(header, &1))
        |> Map.new()

      {:ok, grid}
    else
      _ -> {:error, :enoent}
    end
  end

  @doc """
  Every `valid_time` with a `.pgrid` on disk, sorted ascending.
  """
  @spec list_valid_times() :: [DateTime.t()]
  def list_valid_times do
    case File.ls(base_dir()) do
      {:ok, names} ->
        names
        |> Enum.filter(&String.ends_with?(&1, ".pgrid"))
        |> Enum.map(&(&1 |> String.replace_suffix(".pgrid", "") |> parse_iso()))
        |> Enum.reject(&is_nil/1)
        |> Enum.sort(DateTime)

      _ ->
        []
    end
  end

  defp populated_cell(header, {bin, cell}) do
    case to_profile(header, bin) do
      nil -> []
      profile -> [{cell_latlon(header, cell), profile}]
    end
  end

  defp parse_iso(str) do
    case DateTime.from_iso8601(str) do
      {:ok, dt, _} -> dt
      _ -> nil
    end
  end

  # ── Header ─────────────────────────────────────────────────────────

  defp read_header(fd) do
    with {:ok, fixed} <- :file.pread(fd, 0, @fixed_header_len),
         {:ok, n_fields} <- peek_n_fields(fixed),
         {:ok, table} <- :file.pread(fd, @fixed_header_len, n_fields * @field_name_len) do
      parse_header(fixed <> table)
    else
      _ -> :error
    end
  end

  defp peek_n_fields(<<@magic, @version, _flags::8, n_fields::little-16, _rest::binary>>), do: {:ok, n_fields}

  defp peek_n_fields(_), do: :error

  @doc false
  @spec parse_header(binary()) :: {:ok, Header.t()} | :error
  def parse_header(
        <<@magic, @version, flags::8, n_fields::little-16, valid_unix::little-signed-64, lat_start::little-float-64,
          lon_start::little-float-64, lat_step::little-float-64, lon_step::little-float-64, n_rows::little-16,
          n_cols::little-16, rest::binary>>
      ) do
    table_len = n_fields * @field_name_len

    if byte_size(rest) < table_len do
      :error
    else
      fields =
        rest
        |> binary_part(0, table_len)
        |> chunk_records(@field_name_len)
        |> Enum.map(&trim_nul/1)

      field_index = fields |> Enum.with_index() |> Map.new()

      {:ok,
       %Header{
         hrdps?: Bitwise.band(flags, 1) == 1,
         valid_time: DateTime.from_unix!(valid_unix),
         n_fields: n_fields,
         fields: fields,
         field_index: field_index,
         lat_start: lat_start,
         lon_start: lon_start,
         lat_step: lat_step,
         lon_step: lon_step,
         n_rows: n_rows,
         n_cols: n_cols,
         body_offset: @fixed_header_len + table_len,
         levels: level_slots(fields, field_index)
       }}
    end
  end

  def parse_header(_), do: :error

  defp trim_nul(bin) do
    case :binary.match(bin, <<0>>) do
      {pos, _} -> binary_part(bin, 0, pos)
      :nomatch -> bin
    end
  end

  # Pressure levels present in the field table, as
  # {pres_mb, hght_idx, tmpc_idx, dwpc_idx}, ordered by the table.
  defp level_slots(fields, field_index) do
    fields
    |> Enum.flat_map(fn name ->
      case Regex.run(~r/^hght_m_(\d+)mb$/, name) do
        [_, p] -> [String.to_integer(p)]
        _ -> []
      end
    end)
    |> Enum.map(fn p ->
      {p, field_index["hght_m_#{p}mb"], field_index["tmpc_#{p}mb"], field_index["dwpc_#{p}mb"]}
    end)
  end

  # ── Cell addressing ────────────────────────────────────────────────

  defp record_offset(header, cell), do: header.body_offset + cell * header.n_fields * 4

  @doc false
  @spec cell_index(Header.t(), number(), number()) :: non_neg_integer() | nil
  def cell_index(header, lat, lon) do
    row = round((lat - header.lat_start) / header.lat_step)
    col = round((lon - header.lon_start) / header.lon_step)

    if row >= 0 and col >= 0 and row < header.n_rows and col < header.n_cols do
      row * header.n_cols + col
    end
  end

  defp cell_latlon(header, cell) do
    row = div(cell, header.n_cols)
    col = rem(cell, header.n_cols)

    {Float.round(header.lat_start + row * header.lat_step, 3), Float.round(header.lon_start + col * header.lon_step, 3)}
  end

  # ── Record decoding ────────────────────────────────────────────────

  # Returns nil for a cell with no surface temperature — the `.mp.gz`
  # file only ever contained scored cells, so callers already treat a
  # missing cell as "no data here".
  defp to_profile(header, bin) do
    values = decode_values(bin)

    case at(values, header, "surface_temp_c") do
      nil ->
        nil

      _ ->
        header.fields
        |> Enum.with_index()
        |> Enum.reduce(%{}, fn {name, idx}, acc ->
          put_field(acc, name, Enum.at(values, idx))
        end)
        |> put_commercial(values, header)
        |> put_levels(values, header)
    end
  end

  defp at(values, header, name) do
    case Map.fetch(header.field_index, name) do
      {:ok, idx} -> Enum.at(values, idx)
      :error -> nil
    end
  end

  defp put_field(acc, _name, nil), do: acc

  defp put_field(acc, name, value) do
    case Map.fetch(@field_atoms, name) do
      # Unknown field (level triple, commercial_*, or one added on the
      # Rust side after this reader was written) — handled elsewhere or
      # deliberately ignored.
      :error -> acc
      {:ok, key} when name in @integer_fields -> Map.put(acc, key, trunc(value))
      {:ok, key} -> Map.put(acc, key, value)
    end
  end

  defp put_commercial(acc, values, header) do
    case at(values, header, "commercial_degradation_db") do
      nil ->
        acc

      degradation ->
        Map.put(acc, :commercial_link_degradation, %{
          degradation_db: degradation,
          baseline_dbm: at(values, header, "commercial_baseline_dbm") || 0.0,
          current_dbm: at(values, header, "commercial_current_dbm") || 0.0,
          n_links: trunc(at(values, header, "commercial_n_links") || 0)
        })
    end
  end

  defp put_levels(acc, values, header) do
    levels = Enum.flat_map(header.levels, &build_level(&1, values))

    if levels == [], do: acc, else: Map.put(acc, :profile, levels)
  end

  defp build_level({pres_mb, hght_idx, tmpc_idx, dwpc_idx}, values) do
    do_build_level(
      pres_mb,
      Enum.at(values, hght_idx),
      Enum.at(values, tmpc_idx),
      dwpc_idx && Enum.at(values, dwpc_idx)
    )
  end

  # TMP and HGT are both required, matching the Rust-side filter; a level
  # missing either is dropped rather than surfaced half-populated.
  defp do_build_level(_pres_mb, nil, _tmpc, _dwpc), do: []
  defp do_build_level(_pres_mb, _hght, nil, _dwpc), do: []

  defp do_build_level(pres_mb, hght, tmpc, nil), do: [%{pres_mb: pres_mb * 1.0, hght_m: hght, tmpc: tmpc}]

  defp do_build_level(pres_mb, hght, tmpc, dwpc), do: [%{pres_mb: pres_mb * 1.0, hght_m: hght, tmpc: tmpc, dwpc: dwpc}]

  defp decode_values(bin), do: bin |> chunk_records(4) |> Enum.map(&decode_f32/1)

  # NaN is the missing-value sentinel. Erlang's float binary match
  # rejects NaN and Inf outright, so the fallback clause is what turns
  # "absent" into nil.
  defp decode_f32(<<v::float-little-32>>), do: v
  defp decode_f32(_), do: nil

  defp chunk_records(bin, size) when byte_size(bin) >= size do
    for <<chunk::binary-size(^size) <- bin>>, do: chunk
  end

  defp chunk_records(_bin, _size), do: []

  @doc """
  Snap a lat/lon to the propagation grid step. Kept in sync with
  `Microwaveprop.Propagation.ProfilesFile.snap/2` so a coordinate written
  by either side is reachable from the other.
  """
  @spec snap(number(), number()) :: {float(), float()}
  def snap(lat, lon), do: {snap_one(lat), snap_one(lon)}

  defp snap_one(coord) do
    step = Grid.step()
    Float.round(Float.round(coord / step) * step, 3)
  end
end