defmodule Microwaveprop.Propagation.AsosNudge do
  @moduledoc """
  Blend live ASOS surface observations into the latest HRRR grid between
  hourly runs. HRRR analysis is already ~55 minutes old by the time the
  propagation worker scores it; fronts, humidity swings, and rain onset
  happening after that lag otherwise sit invisible until the next top-of-hour
  run. ASOS METARs publish every 5 minutes, so a 10-minute nudging cycle
  catches them.

  ## Pipeline

  1. `build_station_residuals/2` — for each ASOS observation, find the
     nearest HRRR grid cell (within one 0.125° step, otherwise the station
     is outside the grid domain and dropped). The residual is
     `(asos_value - hrrr_value)` for temp, dewpoint, and pressure, plus the
     raw ASOS rain rate.

  2. `nudge_profile/2` — for a given HRRR profile, compute an
     inverse-distance weighted (1/d²) blend of residuals from stations
     within 250 km. Only the surface fields are touched; refractivity
     gradient, PWAT, HPBL, duct metadata, and the full vertical profile
     all pass through unchanged so the HRRR upper-air signal is preserved.

  3. `compute/3` — run the pipeline over every HRRR profile that has at
     least one station within range, call
     `Microwaveprop.Propagation.score_grid_point/4` on the patched profile,
     and return score rows ready for `Propagation.upsert_scores/1`. Grid
     cells with no station in range are dropped from the output, leaving
     the existing HRRR scores untouched.

  The nudging module itself is pure — no Repo, no PubSub. The worker that
  schedules this (`Microwaveprop.Workers.AsosAdjustmentWorker`) handles all
  I/O.
  """

  alias Microwaveprop.Propagation
  alias Microwaveprop.Propagation.Grid

  @radius_km 250
  @min_station_weight_km 1.0
  @earth_radius_km 6371.0

  # Below this mm/hr threshold MRMS doesn't tell us anything useful — the
  # wet→dry direction of re-scoring would *lose* the wind/sky/native
  # gradient signal that lives in the original PropagationGridWorker run
  # but isn't persisted on HrrrProfile rows, so we'd be trading wet accuracy
  # for strictly worse dry accuracy at that cell.
  @min_mrms_rain_mmhr 0.1

  @type observation :: %{
          required(:lat) => float(),
          required(:lon) => float(),
          optional(:temp_f) => number() | nil,
          optional(:dewpoint_f) => number() | nil,
          optional(:sea_level_pressure_mb) => number() | nil,
          optional(:precip_1h_in) => number() | nil
        }

  @type residual :: %{
          lat: float(),
          lon: float(),
          dtemp_c: float() | nil,
          ddewpoint_c: float() | nil,
          dpressure_mb: float() | nil,
          asos_rain_mmhr: float()
        }

  @type hrrr_profile :: map()

  @type rain_grid :: %{{float(), float()} => float()}

  @type grid_score :: %{
          lat: float(),
          lon: float(),
          valid_time: DateTime.t(),
          band_mhz: non_neg_integer(),
          score: non_neg_integer(),
          factors: map()
        }

  @doc """
  Nudge every HRRR grid cell that has an ASOS station within 250 km *or*
  a meaningful rain signal from MRMS, then re-score. Cells that fail both
  filters are dropped so the existing HRRR scores stay in place.

  The MRMS rain grid is optional — pass an empty map to nudge on ASOS
  alone. `rain_grid` is a `%{{snapped_lat, snapped_lon} => mm_per_hour}`
  keyed on the 0.125° propagation grid.
  """
  @spec compute([observation()], DateTime.t(), [hrrr_profile()], rain_grid()) :: [grid_score()]
  def compute(observations, valid_time, profiles, rain_grid \\ %{})
  def compute(_observations, _valid_time, [], _rain_grid), do: []

  def compute([], _valid_time, _profiles, rain_grid) when map_size(rain_grid) == 0, do: []

  def compute(observations, valid_time, profiles, rain_grid) do
    residuals = build_station_residuals(observations, profiles)

    profiles
    |> Enum.filter(fn profile ->
      any_station_within_radius?(profile, residuals) or
        significant_mrms_rain?(profile, rain_grid)
    end)
    |> Enum.flat_map(fn profile ->
      profile
      |> nudge_profile(residuals)
      |> patch_mrms_rain(rain_grid)
      |> score_point(valid_time)
    end)
  end

  @doc """
  For each observation with a valid HRRR anchor cell, compute the
  `(asos - hrrr)` residual at that cell. Used by `nudge_profile/2` to
  spread the bias field across the grid.
  """
  @spec build_station_residuals([observation()], [hrrr_profile()]) :: [residual()]
  def build_station_residuals(observations, profiles) do
    lookup = Map.new(profiles, fn p -> {snap_to_grid(p.lat, p.lon), p} end)

    observations
    |> Enum.map(&station_residual(&1, lookup))
    |> Enum.reject(&is_nil/1)
  end

  @doc """
  Return a patched HRRR profile with `surface_temp_c`, `surface_dewpoint_c`,
  and `surface_pressure_mb` shifted by the IDW-weighted residuals from
  stations within 250 km. If no station is close enough, the profile is
  returned unchanged.
  """
  @spec nudge_profile(hrrr_profile(), [residual()]) :: hrrr_profile()
  def nudge_profile(profile, residuals) do
    nearby =
      residuals
      |> Enum.map(fn r -> {r, distance_km(profile.lat, profile.lon, r.lat, r.lon)} end)
      |> Enum.filter(fn {_r, d} -> d <= @radius_km end)

    case nearby do
      [] ->
        profile

      _ ->
        profile
        |> Map.put(:surface_temp_c, apply_idw(profile.surface_temp_c, nearby, :dtemp_c))
        |> Map.put(:surface_dewpoint_c, apply_idw(profile.surface_dewpoint_c, nearby, :ddewpoint_c))
        |> Map.put(:surface_pressure_mb, apply_idw(profile.surface_pressure_mb, nearby, :dpressure_mb))
    end
  end

  # -- internals ---------------------------------------------------------

  defp station_residual(obs, lookup) do
    with {:ok, temp_c} <- f_to_c(obs[:temp_f]),
         {:ok, dewpoint_c} <- f_to_c(obs[:dewpoint_f]),
         snapped = snap_to_grid(obs.lat, obs.lon),
         hrrr when not is_nil(hrrr) <- Map.get(lookup, snapped),
         true <- within_one_grid_step?(obs, hrrr) do
      %{
        lat: obs.lat,
        lon: obs.lon,
        dtemp_c: delta(temp_c, hrrr.surface_temp_c),
        ddewpoint_c: delta(dewpoint_c, hrrr.surface_dewpoint_c),
        dpressure_mb: delta(obs[:sea_level_pressure_mb], hrrr.surface_pressure_mb),
        asos_rain_mmhr: precip_to_mmhr(obs[:precip_1h_in])
      }
    else
      _ -> nil
    end
  end

  defp f_to_c(nil), do: :error
  defp f_to_c(f) when is_number(f), do: {:ok, (f - 32) * 5 / 9}
  defp f_to_c(_), do: :error

  defp delta(nil, _), do: nil
  defp delta(_, nil), do: nil
  defp delta(a, b), do: a - b

  defp precip_to_mmhr(nil), do: 0.0
  defp precip_to_mmhr(inches) when inches <= 0, do: 0.0
  defp precip_to_mmhr(inches), do: inches * 25.4

  defp within_one_grid_step?(obs, hrrr) do
    distance_km(obs.lat, obs.lon, hrrr.lat, hrrr.lon) <= Grid.step() * 111.0
  end

  defp any_station_within_radius?(profile, residuals) do
    Enum.any?(residuals, fn r ->
      distance_km(profile.lat, profile.lon, r.lat, r.lon) <= @radius_km
    end)
  end

  defp significant_mrms_rain?(profile, rain_grid) do
    case Map.get(rain_grid, {profile.lat, profile.lon}) do
      rate when is_number(rate) and rate >= @min_mrms_rain_mmhr -> true
      _ -> false
    end
  end

  defp patch_mrms_rain(profile, rain_grid) do
    case Map.get(rain_grid, {profile.lat, profile.lon}) do
      rate when is_number(rate) and rate >= @min_mrms_rain_mmhr ->
        # Scorer reads `hrrr_profile[:precip_mm]` and treats the value as
        # mm/hr directly (see Scorer.precip_to_rate_mmhr/1), so storing the
        # MRMS rate on that key Just Works without a units hack.
        Map.put(profile, :precip_mm, rate)

      _ ->
        profile
    end
  end

  defp apply_idw(base, _nearby, _key) when is_nil(base), do: nil

  defp apply_idw(base, nearby, key) do
    contributing =
      Enum.filter(nearby, fn {residual, _d} -> not is_nil(Map.get(residual, key)) end)

    case contributing do
      [] ->
        base

      _ ->
        {num, den} =
          Enum.reduce(contributing, {0.0, 0.0}, fn {residual, d_km}, {num, den} ->
            w = 1.0 / :math.pow(max(d_km, @min_station_weight_km), 2)
            {num + w * Map.get(residual, key), den + w}
          end)

        base + num / den
    end
  end

  defp score_point(profile, valid_time) do
    profile
    |> Propagation.score_grid_point(valid_time, profile.lat, profile.lon)
    |> Enum.map(fn row ->
      row
      |> Map.put(:lat, profile.lat)
      |> Map.put(:lon, profile.lon)
      |> Map.put(:valid_time, valid_time)
    end)
  end

  defp snap_to_grid(lat, lon) do
    step = Grid.step()

    {
      Float.round(Float.round(lat / step) * step, 3),
      Float.round(Float.round(lon / step) * step, 3)
    }
  end

  defp distance_km(lat1, lon1, lat2, lon2) do
    lat1_rad = lat1 * :math.pi() / 180
    lat2_rad = lat2 * :math.pi() / 180
    dlat = (lat2 - lat1) * :math.pi() / 180
    dlon = (lon2 - lon1) * :math.pi() / 180

    a =
      :math.sin(dlat / 2) * :math.sin(dlat / 2) +
        :math.cos(lat1_rad) * :math.cos(lat2_rad) *
          :math.sin(dlon / 2) * :math.sin(dlon / 2)

    c = 2 * :math.atan2(:math.sqrt(a), :math.sqrt(1 - a))
    @earth_radius_km * c
  end
end