prop/lib/microwaveprop/weather/frontal_analysis.ex

defmodule Microwaveprop.Weather.FrontalAnalysis do
  @moduledoc """
  Detects surface frontal boundaries from HRRR surface grids using
  the Thermal Front Parameter (TFP).

  The TFP (Renard & Clarke, 1965) locates the warm side of frontal
  zones by finding where the gradient of potential temperature
  magnitude is maximized:

      TFP = −∇|∇θ| · ∇θ/|∇θ|

  The most negative TFP values mark cold fronts; values near zero
  are non-frontal. We threshold at an empirically determined value
  and produce a boolean front mask plus a bearing per pixel.

  The meteorologist's key claim: propagation is best parallel to and
  south of an advancing cold front. Paths that cross a front are
  turbulently mixed and go dead. This module enables the geometric
  features (distance_to_front, path_front_angle) that Phase 5 needs.

  All grid math uses Nx for vectorized 2D operations.
  """

  # Nx is `only: [:dev, :test]` in mix.exs. This module is a research
  # tool — no prod code path invokes it — so silence the prod-compile
  # warnings about Nx.* being undefined.
  @compile {:no_warn_undefined, Nx}

  @doc """
  Detect fronts from a 2D grid of surface temperature and pressure.

  ## Arguments

  - `temp_grid` - 2D Nx tensor of surface temperature (K or °C),
    shape `{ny, nx}`, row-major (latitude varies slowest).
  - `pressure_grid` - 2D Nx tensor of surface pressure (Pa),
    same shape. Used to compute potential temperature.
  - `grid_spec` - `%{lat_start, lat_step, lon_start, lon_step}`.

  ## Returns

  `%{mask: Nx.tensor (boolean), bearing_deg: Nx.tensor (float),
     front_points: [{lat, lon, bearing}]}`

  where `front_points` is a list of the strongest frontal pixels
  for easy lookup.
  """
  @spec detect_fronts(Nx.Tensor.t(), Nx.Tensor.t(), map()) :: %{
          tfp: Nx.Tensor.t(),
          mask: Nx.Tensor.t(),
          bearing_deg: Nx.Tensor.t(),
          front_points: [{float(), float(), float()}]
        }
  def detect_fronts(temp_grid, pressure_grid, grid_spec) do
    # Potential temperature: θ = T * (100000/P)^0.286
    theta = Nx.multiply(temp_grid, Nx.pow(Nx.divide(100_000.0, pressure_grid), 0.286))

    # Gradient of theta: dθ/dx (east-west) and dθ/dy (north-south)
    # Use central differences (Nx doesn't have gradient, so manual)
    {ny, nx} = Nx.shape(theta)
    {grad_y, grad_x} = central_gradient(theta)

    # Magnitude of gradient
    grad_mag = Nx.sqrt(Nx.add(Nx.pow(grad_x, 2), Nx.pow(grad_y, 2)))
    # Avoid division by zero
    grad_mag_safe = Nx.max(grad_mag, 1.0e-10)

    # Gradient of gradient magnitude (the TFP numerator uses this)
    {grad_mag_y, grad_mag_x} = central_gradient(grad_mag)

    # TFP = −(∇|∇θ| · ∇θ/|∇θ|)
    # Dot product of grad(|grad_theta|) with the unit gradient of theta
    unit_x = Nx.divide(grad_x, grad_mag_safe)
    unit_y = Nx.divide(grad_y, grad_mag_safe)

    dot = Nx.add(Nx.multiply(grad_mag_x, unit_x), Nx.multiply(grad_mag_y, unit_y))
    tfp = Nx.negate(dot)

    # Threshold: most negative TFP = strongest fronts
    # Empirical threshold calibrated against NWS surface analysis
    threshold = -2.0e-10
    mask = Nx.less(tfp, threshold)

    # Bearing of the front at each pixel (perpendicular to the temperature gradient)
    bearing_rad = Nx.atan2(grad_x, grad_y)
    bearing_deg = Nx.multiply(bearing_rad, 180.0 / :math.pi())

    # Extract front points as a list for spatial queries
    front_points = extract_front_points(mask, bearing_deg, grid_spec, ny, nx)

    %{
      tfp: tfp,
      mask: mask,
      bearing_deg: bearing_deg,
      front_points: front_points
    }
  end

  @doc """
  Compute distance (km) and angle to the nearest front point from a given lat/lon.
  """
  @spec nearest_front([{float(), float(), float()}], float(), float()) ::
          %{distance_km: float(), front_bearing_deg: float(), front_lat: float(), front_lon: float()} | nil
  def nearest_front(front_points, lat, lon) when is_list(front_points) do
    if front_points == [] do
      nil
    else
      front_points
      |> Enum.min_by(fn {flat, flon, _bearing} ->
        dlat = flat - lat
        dlon = (flon - lon) * :math.cos(lat * :math.pi() / 180.0)
        dlat * dlat + dlon * dlon
      end)
      |> then(fn {flat, flon, bearing} ->
        dlat = flat - lat
        dlon = (flon - lon) * :math.cos(lat * :math.pi() / 180.0)
        distance_km = :math.sqrt(dlat * dlat + dlon * dlon) * 111.0

        %{
          distance_km: distance_km,
          front_bearing_deg: bearing,
          front_lat: flat,
          front_lon: flon
        }
      end)
    end
  end

  @doc """
  Angle between a QSO path and the nearest front (0° = parallel, 90° = crosses).
  """
  @spec path_front_angle(float(), float()) :: number()
  def path_front_angle(path_bearing_deg, front_bearing_deg) do
    diff = abs(path_bearing_deg - front_bearing_deg)
    diff = if diff > 180, do: 360 - diff, else: diff
    if diff > 90, do: 180 - diff, else: diff
  end

  # 2D central differences (interior points only; edges use forward/backward)
  defp central_gradient(tensor) do
    {ny, nx} = Nx.shape(tensor)

    # dy: grad along rows (north-south)
    top = Nx.slice(tensor, [0, 0], [ny - 2, nx])
    bottom = Nx.slice(tensor, [2, 0], [ny - 2, nx])
    dy_interior = Nx.divide(Nx.subtract(bottom, top), 2.0)
    # Pad edges with forward/backward diff
    dy_top = Nx.subtract(Nx.slice(tensor, [1, 0], [1, nx]), Nx.slice(tensor, [0, 0], [1, nx]))
    dy_bottom = Nx.subtract(Nx.slice(tensor, [ny - 1, 0], [1, nx]), Nx.slice(tensor, [ny - 2, 0], [1, nx]))
    dy = Nx.concatenate([dy_top, dy_interior, dy_bottom], axis: 0)

    # dx: grad along columns (east-west)
    left = Nx.slice(tensor, [0, 0], [ny, nx - 2])
    right = Nx.slice(tensor, [0, 2], [ny, nx - 2])
    dx_interior = Nx.divide(Nx.subtract(right, left), 2.0)
    dx_left = Nx.subtract(Nx.slice(tensor, [0, 1], [ny, 1]), Nx.slice(tensor, [0, 0], [ny, 1]))
    dx_right = Nx.subtract(Nx.slice(tensor, [0, nx - 1], [ny, 1]), Nx.slice(tensor, [0, nx - 2], [ny, 1]))
    dx = Nx.concatenate([dx_left, dx_interior, dx_right], axis: 1)

    {dy, dx}
  end

  defp extract_front_points(mask, bearing_deg, grid_spec, _ny, nx) do
    mask_binary = Nx.to_flat_list(mask)
    bearing_list = Nx.to_flat_list(bearing_deg)

    mask_binary
    |> Enum.zip(bearing_list)
    |> Enum.with_index()
    |> Enum.flat_map(fn {{is_front, bearing}, idx} ->
      if is_front == 1 do
        row = div(idx, nx)
        col = rem(idx, nx)
        lat = grid_spec.lat_start + row * grid_spec.lat_step
        lon = grid_spec.lon_start + col * grid_spec.lon_step
        [{lat, lon, bearing}]
      else
        []
      end
    end)
  end
end