defmodule Microwaveprop.Propagation.Recalibrator do
  @moduledoc """
  Fits new scoring weights via gradient descent on the historical QSO corpus.

  Uses a logistic regression model: `p = sigmoid(w · factors)` where
  `factors` is a 10-element vector of individual factor scores (0-100)
  from the existing `Scorer` functions, and `w` is the weight vector
  to learn.

  QSOs are labeled as positive examples (1.0). Matched random baseline
  samples (same locations, ±30 days via `Backtest.random_baseline/2`)
  serve as negative examples (0.0).

  The loss is binary cross-entropy. Training uses simple gradient descent
  on CPU — fine for 10 parameters.

  Output: a normalized weights map `%{humidity: 0.XX, ...}` summing to 1.0,
  along with train and validation losses.
  """

  import Ecto.Query

  alias Microwaveprop.Backtest
  alias Microwaveprop.Propagation.BandConfig
  alias Microwaveprop.Propagation.Scorer
  alias Microwaveprop.Radio.Contact
  alias Microwaveprop.Repo
  alias Microwaveprop.Weather

  require Logger

  # Nx is `only: [:dev, :test]` in mix.exs because the recalibrator is a
  # research / weight-fitting tool that never runs in prod (no scheduled
  # job, no cron, no LiveView entry point). Silence the prod-compile
  # warnings about Nx.* being undefined — this module wouldn't be called
  # without a developer explicitly invoking it.
  @compile {:no_warn_undefined, Nx}

  @factor_keys ~w(humidity time_of_day td_depression refractivity sky season wind rain pwat pressure)a

  @doc """
  Fits new weights using gradient descent on the QSO corpus.

  ## Options

    * `:sample_size` - max contacts to load (default: 5000)
    * `:learning_rate` - gradient descent step size (default: 0.01)
    * `:epochs` - number of training iterations (default: 2000)
    * `:band_mhz` - restrict contacts + factor scoring to a specific band
      (default: unrestricted; factor scoring uses 10 GHz config as it did
      historically)

  ## Returns

  A map with `:weights`, `:train_loss`, `:val_loss`, and `:initial_loss`.
  """
  @spec fit(keyword()) :: %{weights: map(), train_loss: float(), val_loss: float(), initial_loss: float()}
  def fit(opts \\ []) do
    sample_size = Keyword.get(opts, :sample_size, 5000)
    learning_rate = Keyword.get(opts, :learning_rate, 0.01)
    epochs = Keyword.get(opts, :epochs, 2000)
    band_mhz = Keyword.get(opts, :band_mhz)

    band_label = if band_mhz, do: "#{band_mhz} MHz", else: "all bands"
    Logger.info("Recalibrator: loading contacts for #{band_label}, computing factor vectors...")

    # Load contacts with HRRR profiles
    {positives, contacts} = load_positive_samples(sample_size, band_mhz)
    Logger.info("Recalibrator: #{length(positives)} positive samples from #{length(contacts)} contacts")

    # Generate matched negatives
    negatives = generate_negative_samples(contacts, length(positives), band_mhz)
    Logger.info("Recalibrator: #{length(negatives)} negative samples")

    if positives == [] or negatives == [] do
      Logger.warning("Recalibrator: insufficient data, returning current weights")
      weights = BandConfig.weights()

      %{
        weights: weights,
        train_loss: 0.0,
        val_loss: 0.0,
        initial_loss: 0.0
      }
    else
      run_training(positives, negatives, learning_rate, epochs)
    end
  end

  @doc """
  Computes the 10-element factor score vector from an HRRR profile and timestamp.

  Uses `Scorer` functions with a default 10 GHz band config (beneficial humidity).
  Returns a list of 10 floats, each in [0, 100].
  """
  @spec compute_factors(map(), DateTime.t()) :: [number()]
  def compute_factors(profile, timestamp), do: compute_factors(profile, timestamp, 10_000)

  @doc """
  Band-aware variant of `compute_factors/2`.

  Factor scoring reaches into band-specific config (humidity direction,
  seasonal table, rain coefficients, etc.) so positives and negatives for
  per-band fits are scored through the exact physics each band uses at
  runtime. `band_mhz` must be a supported band in `BandConfig`.
  """
  @spec compute_factors(map(), DateTime.t(), pos_integer()) :: [number()]
  def compute_factors(profile, timestamp, band_mhz) do
    band_config = BandConfig.get(band_mhz) || BandConfig.get(10_000)

    temp_f = Scorer.c_to_f(profile.surface_temp_c)
    dewpoint_f = Scorer.c_to_f(profile.surface_dewpoint_c)

    abs_humidity =
      if profile.surface_temp_c && profile.surface_dewpoint_c do
        Scorer.absolute_humidity(profile.surface_temp_c, profile.surface_dewpoint_c)
      else
        10.0
      end

    {tod_score, _label} =
      Scorer.score_time_of_day(
        timestamp.hour,
        timestamp.minute,
        timestamp.month,
        profile.lon || -97.0
      )

    [
      Scorer.score_humidity(abs_humidity, band_config),
      tod_score,
      Scorer.score_td_depression(temp_f || 70, dewpoint_f || 60, band_config),
      Scorer.score_refractivity(profile.min_refractivity_gradient, profile.hpbl_m, band_config),
      Scorer.score_sky(nil),
      Scorer.score_season(timestamp.month, profile.lat, profile.lon, band_config),
      Scorer.score_wind(nil),
      Scorer.score_rain(0.0, band_config),
      Scorer.score_pwat(profile.pwat_mm, band_config),
      Scorer.score_pressure(profile.surface_pressure_mb, nil)
    ]
  end

  @doc """
  Runs gradient descent on pre-built factor vectors.

  `positives` and `negatives` are lists of 10-element factor score lists
  (each score in [0, 100]). Returns the same result shape as `fit/1`.

  This is the core training function, separated from data loading so it
  can be tested independently.
  """
  @spec train([list()], [list()], keyword()) :: %{
          weights: map(),
          train_loss: float(),
          val_loss: float(),
          initial_loss: float()
        }
  def train(positives, negatives, opts \\ []) do
    learning_rate = Keyword.get(opts, :learning_rate, 0.01)
    epochs = Keyword.get(opts, :epochs, 2000)
    run_training(positives, negatives, learning_rate, epochs)
  end

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

  defp load_positive_samples(sample_size, band_mhz) do
    base_query =
      Contact
      |> where([c], not is_nil(c.pos1))
      |> where([c], c.distance_km < 3000)
      |> order_by([c], desc: c.qso_timestamp)
      |> limit(^sample_size)

    query =
      if band_mhz do
        from c in base_query, where: c.band == ^Decimal.new(band_mhz)
      else
        base_query
      end

    contacts = Repo.all(query)

    factor_band = band_mhz || 10_000

    # Each contact_to_factors/2 calls `Weather.find_nearest_hrrr/3` —
    # serial execution over ~5000 contacts burns a minute+ just on
    # query round-trips. Parallelize at 20 concurrent queries, well
    # under the :db pool size (30 in prod), so the recalibration
    # finishes in seconds instead of minutes.
    supervisor = {:via, PartitionSupervisor, {Microwaveprop.TaskSupervisor, self()}}

    factor_vectors =
      supervisor
      |> Task.Supervisor.async_stream_nolink(contacts, &contact_to_factors(&1, factor_band),
        max_concurrency: 20,
        timeout: 30_000,
        on_timeout: :kill_task
      )
      |> Enum.flat_map(fn
        {:ok, nil} ->
          []

        {:ok, vec} ->
          [vec]

        {:exit, reason} ->
          Logger.error("Recalibrator contact_to_factors crashed: reason=#{inspect(reason)}")
          []
      end)

    {factor_vectors, contacts}
  end

  defp contact_to_factors(%Contact{pos1: pos, qso_timestamp: ts}, band_mhz) do
    lat = pos["lat"]
    lon = pos["lon"]

    if lat && lon do
      case Weather.find_nearest_hrrr(lat, lon, ts) do
        nil -> nil
        profile -> compute_factors(profile, ts, band_mhz)
      end
    end
  end

  defp generate_negative_samples(contacts, n, band_mhz) do
    baselines = Backtest.random_baseline(n, sample_size: length(contacts))

    factor_band = band_mhz || 10_000

    supervisor = {:via, PartitionSupervisor, {Microwaveprop.TaskSupervisor, self()}}

    supervisor
    |> Task.Supervisor.async_stream_nolink(
      baselines,
      fn {lat, lon, time} ->
        case Weather.find_nearest_hrrr(lat, lon, time) do
          nil -> nil
          profile -> compute_factors(profile, time, factor_band)
        end
      end,
      max_concurrency: 20,
      timeout: 30_000,
      on_timeout: :kill_task
    )
    |> Enum.flat_map(fn
      {:ok, nil} ->
        []

      {:ok, vec} ->
        [vec]

      {:exit, reason} ->
        Logger.error("Recalibrator negative-sample factor compute crashed: reason=#{inspect(reason)}")
        []
    end)
  end

  defp run_training(positives, negatives, learning_rate, epochs) do
    # Build feature matrix and label vector
    all_features = positives ++ negatives
    labels = List.duplicate(1.0, length(positives)) ++ List.duplicate(0.0, length(negatives))

    n = length(all_features)

    # Normalize factor scores to [0, 1]
    x = all_features |> Nx.tensor(type: :f32) |> Nx.divide(100.0)
    y = labels |> Nx.tensor(type: :f32) |> Nx.reshape({n, 1})

    # Cross-validation: split by month index
    # Use indices to split - train on ~80%, validate on ~20%
    {train_x, train_y, val_x, val_y} = split_train_val(x, y, all_features, positives, negatives)

    # Initialize weights uniformly (10 factors, each 0.1)
    w = 0.1 |> List.duplicate(10) |> Nx.tensor(type: :f32) |> Nx.reshape({10, 1})
    bias = Nx.tensor([0.0], type: :f32)

    # Compute initial loss
    initial_loss = compute_loss(train_x, train_y, w, bias)

    # Training loop
    {final_w, final_bias, train_loss} =
      Enum.reduce(1..epochs, {w, bias, initial_loss}, fn epoch, {w_acc, b_acc, _prev_loss} ->
        # Forward pass
        logits = Nx.add(Nx.dot(train_x, w_acc), b_acc)
        preds = Nx.sigmoid(logits)

        # Gradient of BCE loss
        error = Nx.subtract(preds, train_y)
        n_train = Nx.axis_size(train_x, 0)

        grad_w = train_x |> Nx.transpose() |> Nx.dot(error) |> Nx.divide(n_train)
        grad_b = Nx.mean(error)

        # Update
        new_w = Nx.subtract(w_acc, Nx.multiply(learning_rate, grad_w))
        new_b = Nx.subtract(b_acc, Nx.multiply(learning_rate, grad_b))

        loss = compute_loss(train_x, train_y, new_w, new_b)

        if rem(epoch, 500) == 0 do
          Logger.info("Recalibrator: epoch #{epoch}, train_loss=#{loss |> Nx.to_number() |> Float.round(6)}")
        end

        {new_w, new_b, loss}
      end)

    # Validation loss
    val_loss = compute_loss(val_x, val_y, final_w, final_bias)

    # Extract and normalize weights
    raw_weights = final_w |> Nx.abs() |> Nx.squeeze()
    weight_sum = Nx.sum(raw_weights)
    normalized = Nx.divide(raw_weights, weight_sum)

    weight_values = Nx.to_flat_list(normalized)

    weights_map =
      @factor_keys
      |> Enum.zip(weight_values)
      |> Map.new(fn {k, v} -> {k, Float.round(v, 4)} end)

    %{
      weights: weights_map,
      train_loss: train_loss |> Nx.to_number() |> to_float(),
      val_loss: val_loss |> Nx.to_number() |> to_float(),
      initial_loss: initial_loss |> Nx.to_number() |> to_float()
    }
  end

  defp compute_loss(x, y, w, bias) do
    logits = Nx.add(Nx.dot(x, w), bias)
    preds = Nx.sigmoid(logits)

    # Binary cross-entropy: -mean(y * log(p) + (1-y) * log(1-p))
    eps = 1.0e-7
    preds_clipped = Nx.clip(preds, eps, 1.0 - eps)

    Nx.negate(
      Nx.mean(
        Nx.add(
          Nx.multiply(y, Nx.log(preds_clipped)),
          Nx.multiply(Nx.subtract(1.0, y), Nx.log(Nx.subtract(1.0, preds_clipped)))
        )
      )
    )
  end

  defp split_train_val(x, y, _all_features, positives, negatives) do
    # Simple 80/20 split
    n_pos = length(positives)
    n_neg = length(negatives)

    train_pos = max(1, round(n_pos * 0.8))
    train_neg = max(1, round(n_neg * 0.8))

    # Positives come first in the combined tensor, then negatives.
    # Take 80% of positives and 80% of negatives for training.
    train_pos_x = Nx.slice(x, [0, 0], [train_pos, 10])
    train_neg_x = Nx.slice(x, [n_pos, 0], [train_neg, 10])
    train_x = Nx.concatenate([train_pos_x, train_neg_x], axis: 0)

    train_pos_y = Nx.slice(y, [0, 0], [train_pos, 1])
    train_neg_y = Nx.slice(y, [n_pos, 0], [train_neg, 1])
    train_y = Nx.concatenate([train_pos_y, train_neg_y], axis: 0)

    val_pos_count = n_pos - train_pos
    val_neg_count = n_neg - train_neg

    if val_pos_count <= 0 or val_neg_count <= 0 do
      # Not enough data for validation — use training data for both
      {train_x, train_y, train_x, train_y}
    else
      val_pos_x = Nx.slice(x, [train_pos, 0], [val_pos_count, 10])
      val_neg_x = Nx.slice(x, [n_pos + train_neg, 0], [val_neg_count, 10])
      val_x = Nx.concatenate([val_pos_x, val_neg_x], axis: 0)

      val_pos_y = Nx.slice(y, [train_pos, 0], [val_pos_count, 1])
      val_neg_y = Nx.slice(y, [n_pos + train_neg, 0], [val_neg_count, 1])
      val_y = Nx.concatenate([val_pos_y, val_neg_y], axis: 0)

      {train_x, train_y, val_x, val_y}
    end
  end

  defp to_float(n) when is_float(n), do: n
  defp to_float(n) when is_integer(n), do: n * 1.0
  defp to_float(_), do: 0.0
end