prop/lib/microwaveprop/weather/sgrid.ex

defmodule Microwaveprop.Weather.Sgrid do
  @moduledoc """
  Reader for the `.sgrid` derived-weather-scalar format written by
  `prop-grid-rs` (`rust/prop_grid_rs/src/sgrid.rs`).

  Replaces the 5°×5° chunked 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.
    * reading a viewport is one contiguous `pread` per grid row.

  Same header layout as `Microwaveprop.Propagation.Pgrid` with magic
  `SGRD`. The field table names map 1:1 to the atom keys
  `Microwaveprop.Weather.ScalarFile` already produces, so callers see
  the same map shape regardless of the underlying artifact format.

  ## Layout (little-endian)

      magic       4   "SGRD"
      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`.
  """

  @magic "SGRD"
  @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. Matches
  # `Microwaveprop.Weather.ScalarFile`'s `@atom_keys` whitelist
  # and `prop_grid_rs::sgrid::SCALAR_FIELDS`.
  @field_atoms %{
    "temperature" => :temperature,
    "dewpoint_depression" => :dewpoint_depression,
    "surface_rh" => :surface_rh,
    "surface_pressure_mb" => :surface_pressure_mb,
    "surface_refractivity" => :surface_refractivity,
    "refractivity_gradient" => :refractivity_gradient,
    "bl_height" => :bl_height,
    "pwat" => :pwat,
    "temp_850mb" => :temp_850mb,
    "dewpoint_850mb" => :dewpoint_850mb,
    "temp_700mb" => :temp_700mb,
    "dewpoint_700mb" => :dewpoint_700mb,
    "lapse_rate" => :lapse_rate,
    "mid_lapse_rate" => :mid_lapse_rate,
    "inversion_strength" => :inversion_strength,
    "inversion_base_m" => :inversion_base_m,
    "ducting" => :ducting,
    "duct_base_m" => :duct_base_m,
    "duct_strength" => :duct_strength,
    "duct_cutoff_ghz" => :duct_cutoff_ghz
  }

  defmodule Header do
    @moduledoc "Parsed `.sgrid` 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
    ]
  end

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

  @doc "Absolute path for the `.sgrid` 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}.sgrid")
  end

  @doc "HRDPS sibling path — `<base>/weather_scalars/<iso>.hrdps.sgrid`."
  @spec path_for_hrdps(DateTime.t()) :: String.t()
  def path_for_hrdps(%DateTime{} = valid_time) do
    iso = valid_time |> DateTime.truncate(:second) |> DateTime.to_iso8601()
    Path.join(base_dir(), "#{iso}.hrdps.sgrid")
  end

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

  @doc "Whether an HRDPS `.sgrid` exists for `valid_time`."
  @spec exists_hrdps?(DateTime.t()) :: boolean()
  def exists_hrdps?(%DateTime{} = valid_time) do
    File.exists?(path_for_hrdps(valid_time))
  end

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

  Opens the file, `pread`s the header and then one record. Returns `nil`
  when the file is missing, the point falls outside the grid, or the
  cell has no surface temperature.
  """
  @spec read_point(DateTime.t(), number(), number()) :: map() | nil
  def read_point(%DateTime{} = valid_time, lat, lon) do
    # Try HRRR first, then HRDPS for Canadian cells.
    Enum.find_value(
      [path_for(valid_time), path_for_hrdps(valid_time)],
      &read_point_from_path(&1, lat, lon)
    )
  end

  defp read_point_from_path(path, lat, lon) do
    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_row(header, bin, lat, lon)
    else
      _ -> nil
    end
  end

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

  Cells with no surface temperature are skipped. Returns `{:ok, grid}` or
  `{:error, :enoent}`.
  """
  @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 """
  Read every persisted row for `valid_time` whose lat/lon falls within
  `bounds`. Pass `nil` to read every cell. Returns `[]` if no file
  exists. Merges HRRR + HRDPS, preferring HRRR on overlap.
  """
  @spec read_bounds(DateTime.t(), ScalarFile.bounds() | nil) :: [ScalarFile.row()]
  def read_bounds(%DateTime{} = valid_time, bounds) do
    hrrr = read_bounds_from(path_for(valid_time), bounds)
    hrdps = read_bounds_from(path_for_hrdps(valid_time), bounds)

    case {hrrr, hrdps} do
      {[], []} -> []
      {h, []} -> h
      {[], c} -> c
      {h, c} -> merge_prefer_hrrr(h, c)
    end
  end

  @doc """
  Read only the HRDPS sibling rows for `valid_time` within `bounds`.
  Used by `/weather-ca` for Canadian-only views.
  """
  @spec read_bounds_hrdps(DateTime.t(), ScalarFile.bounds() | nil) :: [ScalarFile.row()]
  def read_bounds_hrdps(%DateTime{} = valid_time, bounds) do
    read_bounds_from(path_for_hrdps(valid_time), bounds)
  end

  @doc """
  Every `valid_time` with a `.sgrid` 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, ".sgrid"))
        |> Enum.reject(&String.ends_with?(&1, ".hrdps.sgrid"))
        |> Enum.map(&(&1 |> String.replace_suffix(".sgrid", "") |> parse_iso()))
        |> Enum.reject(&is_nil/1)
        |> Enum.sort(DateTime)

      _ ->
        []
    end
  end

  @doc """
  Like `list_valid_times/0` but only for HRDPS `.sgrid` files.
  """
  @spec list_valid_times_hrdps() :: [DateTime.t()]
  def list_valid_times_hrdps do
    case File.ls(base_dir()) do
      {:ok, names} ->
        names
        |> Enum.filter(&String.ends_with?(&1, ".hrdps.sgrid"))
        |> Enum.map(&(&1 |> String.replace_suffix(".hrdps.sgrid", "") |> parse_iso()))
        |> Enum.reject(&is_nil/1)
        |> Enum.sort(DateTime)

      _ ->
        []
    end
  end

  # ── Private ─────────────────────────────────────────────────────────

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

      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))
      |> maybe_filter_bounds(bounds)
    else
      _ -> []
    end
  end

  defp maybe_filter_bounds(cells, nil), do: cells

  defp maybe_filter_bounds(cells, %{"south" => s, "north" => n, "west" => w, "east" => e}) do
    Enum.filter(cells, fn {{lat, lon}, _row} ->
      lat >= s and lat <= n and lon >= w and lon <= e
    end)
  end

  defp populated_cell(header, {bin, cell}) do
    {lat, lon} = cell_latlon(header, cell)

    case to_row(header, bin, lat, lon) do
      nil -> []
      row -> [{{lat, lon}, row}]
    end
  end

  defp merge_prefer_hrrr(hrrr, hrdps) do
    hrrr_keys = MapSet.new(hrrr, fn {{lat, lon}, _row} -> {lat, lon} end)
    extras = Enum.reject(hrdps, fn {{lat, lon}, _row} -> MapSet.member?(hrrr_keys, {lat, lon}) end)
    hrrr ++ extras
  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
       }}
    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

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

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

  defp 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 — callers treat a
  # missing cell as "no data here".
  defp to_row(header, bin, lat, lon) do
    values = decode_values(bin)

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

      _ ->
        row =
          header.fields
          |> Enum.with_index()
          |> Enum.reduce(%{lat: lat, lon: lon, valid_time: header.valid_time}, fn {name, idx}, acc ->
            put_field(acc, name, Enum.at(values, idx))
          end)

        row
    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
      :error ->
        acc

      {:ok, key} when key == :ducting ->
        Map.put(acc, key, value == 1.0)

      {:ok, key} ->
        Map.put(acc, key, value)
    end
  end

  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: []
end