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

  @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)

  ## 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)

    Logger.info("Recalibrator: loading contacts and computing factor vectors...")

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

    # Generate matched negatives
    negatives = generate_negative_samples(contacts, length(positives))
    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
    band_config = 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) do
    contacts =
      Contact
      |> where([c], not is_nil(c.pos1))
      |> where([c], c.distance_km < 3000)
      |> order_by([c], desc: c.qso_timestamp)
      |> limit(^sample_size)
      |> Repo.all()

    factor_vectors =
      contacts
      |> Enum.map(&contact_to_factors/1)
      |> Enum.reject(&is_nil/1)

    {factor_vectors, contacts}
  end

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

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

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

    baselines
    |> Enum.map(fn {lat, lon, time} ->
      case Weather.find_nearest_hrrr(lat, lon, time) do
        nil -> nil
        profile -> compute_factors(profile, time)
      end
    end)
    |> Enum.reject(&is_nil/1)
  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