prop/lib/microwaveprop/propagation.ex

defmodule Microwaveprop.Propagation do
  @moduledoc false

  import Ecto.Query

  alias Microwaveprop.Propagation.BandConfig
  alias Microwaveprop.Propagation.Grid
  alias Microwaveprop.Propagation.GridScore
  alias Microwaveprop.Propagation.Scorer
  alias Microwaveprop.Repo
  alias Microwaveprop.Weather.SoundingParams

  require Logger

  @ml_key :propagation_ml
  @ml_module Microwaveprop.Propagation.Model

  @doc """
  Loads the ML model from disk, compiles the predict function, and caches both
  in persistent_term. No-op if the model file doesn't exist or ML deps unavailable.
  """
  def load_ml_model do
    if Code.ensure_loaded?(@ml_module) do
      case apply(@ml_module, :load, []) do
        {:ok, params} ->
          predict_fn = apply(@ml_module, :compile_predict, [])
          :persistent_term.put(@ml_key, {predict_fn, params})
          Logger.info("PropagationML: model loaded and compiled")
          :ok

        :error ->
          Logger.info("PropagationML: no model file found, using algorithm scorer only")
          :ok
      end
    else
      Logger.info("PropagationML: ML dependencies not available")
      :ok
    end
  end

  @doc "Returns cached {predict_fn, params} tuple, or nil if not loaded."
  def ml_model do
    :persistent_term.get(@ml_key, nil)
  end

  @doc """
  Score a single grid point across all bands using HRRR profile data.
  Uses ML model if loaded, falls back to algorithm scorer.
  Returns a list of %{band_mhz, score, factors} maps.
  """
  def score_grid_point(hrrr_profile, valid_time, latitude, longitude) do
    derived = derive_from_hrrr(hrrr_profile)

    temp_c = hrrr_profile.surface_temp_c
    dewpoint_c = hrrr_profile.surface_dewpoint_c

    # Skip points with missing or physically impossible surface data
    if is_nil(temp_c) or is_nil(dewpoint_c) or temp_c < -80 or temp_c > 60 or
         dewpoint_c < -80 or dewpoint_c > 50 do
      []
    else
      score_grid_point_with_data(hrrr_profile, valid_time, temp_c, dewpoint_c, derived, latitude, longitude)
    end
  end

  defp score_grid_point_with_data(hrrr_profile, valid_time, temp_c, dewpoint_c, derived, latitude, longitude) do
    # Algorithm is the primary scorer — always used for the map score.
    # ML score stored in factors as :ml_score for comparison/analysis.
    score_with_algorithm(hrrr_profile, valid_time, temp_c, dewpoint_c, derived, latitude, longitude)
  end

  defp score_with_algorithm(hrrr_profile, valid_time, temp_c, dewpoint_c, derived, _latitude, longitude) do
    temp_f = Scorer.c_to_f(temp_c)
    dewpoint_f = Scorer.c_to_f(dewpoint_c)

    conditions = %{
      abs_humidity: Scorer.absolute_humidity(temp_c, dewpoint_c),
      temp_f: temp_f,
      dewpoint_f: dewpoint_f,
      wind_speed_kts: Scorer.wind_speed_kts(hrrr_profile[:wind_u], hrrr_profile[:wind_v]),
      sky_cover_pct: hrrr_profile[:cloud_cover_pct],
      utc_hour: valid_time.hour,
      utc_minute: valid_time.minute,
      month: valid_time.month,
      longitude: longitude,
      pressure_mb: hrrr_profile.surface_pressure_mb,
      prev_pressure_mb: nil,
      rain_rate_mmhr: Scorer.precip_to_rate_mmhr(hrrr_profile[:precip_mm]),
      min_refractivity_gradient: derived[:min_refractivity_gradient],
      bl_depth_m: hrrr_profile[:hpbl_m],
      pwat_mm: hrrr_profile[:pwat_mm]
    }

    Enum.map(BandConfig.all_bands(), fn band_config ->
      result = Scorer.composite_score(conditions, band_config)
      Map.put(result, :band_mhz, band_config.freq_mhz)
    end)
  end

  @doc """
  Upsert propagation scores in batches within a transaction so readers see all-or-nothing.

  Options:
    * `:prune` - whether to prune old scores after upsert (default true)
  """
  def upsert_scores(scores, opts \\ []) do
    now = DateTime.truncate(DateTime.utc_now(), :second)

    entries =
      Enum.map(scores, fn s ->
        %{
          id: Ecto.UUID.generate(),
          lat: s.lat,
          lon: s.lon,
          valid_time: s.valid_time,
          band_mhz: s.band_mhz,
          score: s.score,
          factors: s.factors,
          inserted_at: now,
          updated_at: now
        }
      end)

    result =
      Repo.transaction(
        fn ->
          entries
          |> Enum.chunk_every(500)
          |> Enum.reduce(0, fn chunk, acc ->
            {count, _} =
              Repo.insert_all(GridScore, chunk,
                on_conflict:
                  from(g in GridScore,
                    update: [
                      set: [
                        score: fragment("EXCLUDED.score"),
                        factors: fragment("EXCLUDED.factors"),
                        updated_at: fragment("EXCLUDED.updated_at")
                      ]
                    ],
                    where: g.score != fragment("EXCLUDED.score")
                  ),
                conflict_target: [:lat, :lon, :valid_time, :band_mhz]
              )

            acc + count
          end)
        end,
        timeout: 600_000
      )

    if Keyword.get(opts, :prune, true) do
      case result do
        {:ok, _count} -> prune_old_scores()
        _ -> :ok
      end
    end

    result
  end

  @doc "Remove scores with valid_times older than 2 hours."
  def prune_old_scores do
    cutoff = DateTime.add(DateTime.utc_now(), -2, :hour)

    {deleted, _} = Repo.delete_all(from(gs in GridScore, where: gs.valid_time < ^cutoff), timeout: 120_000)

    if deleted > 0 do
      Logger.info("PropagationScores: pruned #{deleted} old scores (before #{cutoff})")
    end
  end

  @doc """
  Returns distinct valid_times for a band, ordered ascending.
  Filters out times more than 1 hour in the past, but always includes
  the most recent valid_time so there's always data to display.
  """
  def available_valid_times(band_mhz) do
    cutoff = DateTime.add(DateTime.utc_now(), -3600, :second)

    times =
      Repo.all(
        from(gs in GridScore,
          where: gs.band_mhz == ^band_mhz and gs.valid_time >= ^cutoff,
          select: gs.valid_time,
          distinct: gs.valid_time,
          order_by: [asc: gs.valid_time]
        )
      )

    if times == [] do
      # No future times — return the single most recent as fallback
      case latest_valid_time(band_mhz) do
        nil -> []
        t -> [t]
      end
    else
      times
    end
  end

  @doc """
  Get scores for a band at a specific valid_time, optionally within a bounding box.
  If valid_time is nil, uses the earliest available (current analysis hour).
  Excludes factors for performance.
  """
  def scores_at(band_mhz, valid_time, bounds \\ nil) do
    time = valid_time || earliest_valid_time(band_mhz)

    case time do
      nil ->
        []

      _ ->
        from(gs in GridScore,
          where: gs.band_mhz == ^band_mhz and gs.valid_time == ^time,
          select: %{lat: gs.lat, lon: gs.lon, score: gs.score, valid_time: gs.valid_time}
        )
        |> maybe_filter_bounds(bounds)
        |> Repo.all()
    end
  end

  @doc "Get the latest scores for a band (alias for scores_at with earliest valid_time)."
  def latest_scores(band_mhz, bounds \\ nil) do
    scores_at(band_mhz, nil, bounds)
  end

  defp earliest_valid_time(band_mhz) do
    Repo.one(from(gs in GridScore, where: gs.band_mhz == ^band_mhz, select: min(gs.valid_time)))
  end

  @doc "Get scores across all forecast hours for a single grid point (for sparkline)."
  def point_forecast(band_mhz, lat, lon) do
    step = Grid.step()
    snapped_lat = Float.round(Float.round(lat / step) * step, 3)
    snapped_lon = Float.round(Float.round(lon / step) * step, 3)
    now = DateTime.utc_now()

    Repo.all(
      from(gs in GridScore,
        where:
          gs.band_mhz == ^band_mhz and
            gs.lat == ^snapped_lat and
            gs.lon == ^snapped_lon and
            gs.valid_time >= ^now,
        select: %{valid_time: gs.valid_time, score: gs.score},
        order_by: [asc: gs.valid_time]
      )
    )
  end

  @doc "Get the full score and factors for a specific grid point, snapped to nearest grid."
  def point_detail(band_mhz, lat, lon, valid_time \\ nil) do
    step = Grid.step()
    snapped_lat = Float.round(Float.round(lat / step) * step, 3)
    snapped_lon = Float.round(Float.round(lon / step) * step, 3)

    time = valid_time || latest_valid_time(band_mhz)

    case time do
      nil ->
        nil

      _ ->
        Repo.one(
          from(gs in GridScore,
            where:
              gs.band_mhz == ^band_mhz and
                gs.valid_time == ^time and
                gs.lat == ^snapped_lat and
                gs.lon == ^snapped_lon,
            select: %{
              lat: gs.lat,
              lon: gs.lon,
              score: gs.score,
              factors: gs.factors,
              valid_time: gs.valid_time
            }
          )
        )
    end
  end

  defp maybe_filter_bounds(query, nil), do: query

  defp maybe_filter_bounds(query, %{"south" => s, "north" => n, "west" => w, "east" => e}) do
    from(gs in query,
      where: gs.lat >= ^s and gs.lat <= ^n and gs.lon >= ^w and gs.lon <= ^e
    )
  end

  @doc "Get the latest valid_time across all scores."
  def latest_valid_time do
    Repo.one(from(gs in GridScore, select: max(gs.valid_time)))
  end

  @doc "Get the latest valid_time for a specific band."
  def latest_valid_time(band_mhz) do
    Repo.one(from(gs in GridScore, where: gs.band_mhz == ^band_mhz, select: max(gs.valid_time)))
  end

  defp derive_from_hrrr(%{profile: profile}) when is_list(profile) and length(profile) >= 3 do
    case SoundingParams.derive(profile) do
      nil -> %{}
      derived -> %{min_refractivity_gradient: derived.min_refractivity_gradient}
    end
  end

  defp derive_from_hrrr(_), do: %{}
end