defmodule Mix.Tasks.PropagationTrain do
  @shortdoc "Train ML model from QSO-HRRR matched data"

  @moduledoc """
  Trains the propagation ML model on real QSO outcomes matched to HRRR
  atmospheric conditions. The target is QSO distance normalized per-band
  (0-1), so the model learns what atmospheric conditions actually produce
  good propagation from historical contacts — not from the hand-tuned algorithm.

  Features: HRRR conditions averaged across both QSO endpoints + solar time + frequency.
  Target: distance_km / band_p99_distance (capped at 1.0).

  ## Usage

      mix propagation_train
      mix propagation_train --epochs 100 --batch-size 512
  """

  use Mix.Task

  alias Microwaveprop.Propagation.Model
  alias Microwaveprop.Repo

  @query_timeout 600_000

  # Per-band normalization: p99 distances from the QSO dataset.
  # Distance is divided by these values and capped at 1.0.
  @band_max_km %{
    10_000 => 800.0,
    24_000 => 350.0,
    47_000 => 200.0,
    75_000 => 200.0,
    122_000 => 140.0,
    134_000 => 160.0,
    241_000 => 115.0
  }

  @impl Mix.Task
  def run(args) do
    {opts, _, _} =
      OptionParser.parse(args,
        strict: [epochs: :integer, batch_size: :integer]
      )

    Application.put_env(
      :microwaveprop,
      Oban,
      Keyword.put(
        Application.get_env(:microwaveprop, Oban, []),
        :queues,
        false
      )
    )

    Mix.Task.run("app.start")

    epochs = Keyword.get(opts, :epochs, 100)
    batch_size = Keyword.get(opts, :batch_size, 256)

    "=" |> String.duplicate(70) |> IO.puts()
    IO.puts("PROPAGATION MODEL TRAINING (QSO-HRRR)")
    "=" |> String.duplicate(70) |> IO.puts()
    IO.puts("Config: epochs=#{epochs}, batch_size=#{batch_size}")
    IO.puts("Started: #{DateTime.to_string(DateTime.utc_now())}\n")

    IO.puts("Loading QSO-HRRR training data...")
    {features, targets, band_counts} = load_training_data()
    n = elem(Nx.shape(features), 0)

    IO.puts("\nLoaded #{n} training samples from real QSOs:")

    band_counts
    |> Enum.sort_by(fn {band, _} -> band end)
    |> Enum.each(fn {band, count} ->
      max_km = Map.get(@band_max_km, band)
      suffix = if max_km, do: " (norm: #{trunc(max_km)} km)", else: " (excluded)"
      IO.puts("  #{div(band, 1000)} GHz: #{count} QSOs#{suffix}")
    end)

    IO.puts("")

    # Shuffle the dataset
    IO.puts("Shuffling and splitting dataset (80/10/10)...")
    key = Nx.Random.key(System.os_time())
    {indices, _} = Nx.Random.shuffle(key, Nx.iota({n}))
    indices = Nx.as_type(indices, :s64)
    features = Nx.take(features, indices)
    targets = Nx.take(targets, indices)

    # 80/10/10 split
    train_end = trunc(n * 0.8)
    val_end = trunc(n * 0.9)

    train_features = Nx.slice(features, [0, 0], [train_end, 13])
    train_targets = Nx.slice(targets, [0, 0], [train_end, 1])
    val_features = Nx.slice(features, [train_end, 0], [val_end - train_end, 13])
    val_targets = Nx.slice(targets, [train_end, 0], [val_end - train_end, 1])
    test_features = Nx.slice(features, [val_end, 0], [n - val_end, 13])
    test_targets = Nx.slice(targets, [val_end, 0], [n - val_end, 1])

    IO.puts("  Train: #{train_end}, Val: #{val_end - train_end}, Test: #{n - val_end}\n")

    # Train
    IO.puts("Training for #{epochs} epochs...")
    {trained_state, train_metrics} = Model.train(train_features, train_targets, epochs: epochs, batch_size: batch_size)
    final_loss = train_metrics.final_loss

    if is_float(final_loss) and not (final_loss != final_loss) do
      IO.puts("Training complete. Final loss: #{Float.round(final_loss, 6)}\n")
    else
      IO.puts("Training complete. Final loss: NaN (training diverged)\n")
      IO.puts("Try: --sample 100000 or reduce epochs")
      System.halt(1)
    end

    # Evaluate on validation set
    val_metrics = Model.evaluate(trained_state, val_features, val_targets)
    IO.puts("Validation metrics:")
    IO.puts("  RMSE:      #{Float.round(val_metrics.rmse * 100, 2)} points (on 0-100 scale)")
    IO.puts("  R-squared: #{Float.round(val_metrics.r_squared, 4)}\n")

    # Evaluate on test set
    test_metrics = Model.evaluate(trained_state, test_features, test_targets)
    IO.puts("Test metrics:")
    IO.puts("  RMSE:      #{Float.round(test_metrics.rmse * 100, 2)} points (on 0-100 scale)")
    IO.puts("  R-squared: #{Float.round(test_metrics.r_squared, 4)}\n")

    # Save model
    IO.puts("Saving model to priv/models/propagation_v1.nx...")
    Model.save(trained_state)
    IO.puts("Done!")

    "-" |> String.duplicate(70) |> IO.puts()
    IO.puts("Finished: #{DateTime.to_string(DateTime.utc_now())}")
  end

  defp load_training_data do
    # Join QSOs to HRRR at both endpoints, average conditions.
    # Same query as the analysis task but returns raw values for encoding.
    sql = """
    SELECT
      q.band::integer AS band_mhz,
      q.distance_km::float AS distance_km,
      EXTRACT(HOUR FROM q.qso_timestamp)::int AS utc_hour,
      EXTRACT(MONTH FROM q.qso_timestamp)::int AS month,
      ((q.pos1->>'lng')::float + (q.pos2->>'lng')::float) / 2.0 AS avg_lon,
      (COALESCE(h1.surface_temp_c, h2.surface_temp_c) +
       COALESCE(h2.surface_temp_c, h1.surface_temp_c)) / 2.0 AS avg_temp,
      (COALESCE(h1.surface_dewpoint_c, h2.surface_dewpoint_c) +
       COALESCE(h2.surface_dewpoint_c, h1.surface_dewpoint_c)) / 2.0 AS avg_dewpoint,
      (COALESCE(h1.surface_pressure_mb, h2.surface_pressure_mb) +
       COALESCE(h2.surface_pressure_mb, h1.surface_pressure_mb)) / 2.0 AS avg_pressure,
      (COALESCE(h1.min_refractivity_gradient, h2.min_refractivity_gradient) +
       COALESCE(h2.min_refractivity_gradient, h1.min_refractivity_gradient)) / 2.0 AS avg_gradient,
      (COALESCE(h1.hpbl_m, h2.hpbl_m) +
       COALESCE(h2.hpbl_m, h1.hpbl_m)) / 2.0 AS avg_hpbl,
      (COALESCE(h1.pwat_mm, h2.pwat_mm) +
       COALESCE(h2.pwat_mm, h1.pwat_mm)) / 2.0 AS avg_pwat
    FROM qsos q
    INNER JOIN hrrr_profiles h1
      ON h1.lat = ROUND((q.pos1->>'lat')::numeric * 8) / 8
      AND h1.lon = ROUND((q.pos1->>'lng')::numeric * 8) / 8
      AND h1.valid_time = date_trunc('hour', q.qso_timestamp)
    INNER JOIN hrrr_profiles h2
      ON h2.lat = ROUND((q.pos2->>'lat')::numeric * 8) / 8
      AND h2.lon = ROUND((q.pos2->>'lng')::numeric * 8) / 8
      AND h2.valid_time = date_trunc('hour', q.qso_timestamp)
    WHERE q.distance_km > 0
      AND q.distance_km < 3000
      AND q.qso_timestamp >= '2016-06-30'
      AND q.pos1->>'lng' IS NOT NULL
      AND q.pos2->>'lng' IS NOT NULL
    """

    IO.puts("  Running QSO-HRRR join query...")
    %{rows: rows} = Repo.query!(sql, [], timeout: @query_timeout)
    IO.puts("  Matched #{length(rows)} QSOs to HRRR conditions")

    band_counts = Enum.frequencies_by(rows, fn row -> Enum.at(row, 0) end)

    # Group by band and compute within-band distance percentile as target.
    # This converts raw distance to "how good were conditions relative to
    # what's possible on this band" — reduces noise from operator/equipment.
    rows_by_band =
      rows
      |> Enum.filter(fn [band_mhz | _] -> Map.has_key?(@band_max_km, band_mhz) end)
      |> Enum.group_by(fn [band_mhz | _] -> band_mhz end)

    percentile_lookup =
      Enum.flat_map(rows_by_band, fn {_band, band_rows} ->
        distances = band_rows |> Enum.map(fn [_, d | _] -> d end) |> Enum.sort()
        n = length(distances)

        Enum.map(band_rows, fn [_, distance_km | _] = row ->
          # Rank this distance within its band (0.0 to 1.0)
          rank = Enum.count(distances, &(&1 <= distance_km))
          pct = rank / n
          {row, pct}
        end)
      end)

    {feature_rows, target_rows} =
      percentile_lookup
      |> Enum.map(fn {[band_mhz, _distance_km, utc_hour, month, lon, temp, dewpoint, pressure, grad, hpbl, pwat],
                      percentile} ->
        features =
          Model.encode_features(%{
            surface_temp_c: to_float(temp),
            surface_dewpoint_c: to_float(dewpoint),
            surface_pressure_mb: to_float(pressure),
            min_refractivity_gradient: to_float(grad),
            hpbl_m: to_float(hpbl),
            pwat_mm: to_float(pwat),
            utc_hour: to_float(utc_hour),
            month: trunc(to_float(month)),
            longitude: to_float(lon),
            freq_mhz: band_mhz
          })

        {features, [percentile]}
      end)
      |> Enum.unzip()

    IO.puts("  After filtering: #{length(feature_rows)} training samples")

    features_tensor = Nx.tensor(feature_rows, type: :f32)
    targets_tensor = Nx.tensor(target_rows, type: :f32)

    {features_tensor, targets_tensor, band_counts}
  end

  defp to_float(nil), do: 0.0
  defp to_float(%Decimal{} = d), do: Decimal.to_float(d)
  defp to_float(v) when is_float(v), do: v
  defp to_float(v) when is_integer(v), do: v / 1
end