defmodule Mix.Tasks.PropagationTrain do
  @shortdoc "Train ML model from propagation scores + HRRR data"

  @moduledoc """
  Builds training data from propagation_scores joined to hrrr_profiles,
  trains the Axon model, evaluates on held-out test set, and saves weights.

  Uses stratified sampling across all 8 bands so the model learns
  frequency-specific scoring behavior via the log_freq_mhz feature.

  ## Usage

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

  use Mix.Task

  alias Microwaveprop.Propagation.BandConfig
  alias Microwaveprop.Propagation.Model
  alias Microwaveprop.Repo

  @query_timeout 600_000

  @impl Mix.Task
  def run(args) do
    {opts, _, _} =
      OptionParser.parse(args,
        strict: [epochs: :integer, batch_size: :integer, sample: :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, 50)
    batch_size = Keyword.get(opts, :batch_size, 256)
    sample_size = Keyword.get(opts, :sample, 500_000)

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

    IO.puts("Loading training data (stratified, sample=#{sample_size})...")
    {features, targets, band_counts} = load_training_data(sample_size)
    n = elem(Nx.shape(features), 0)

    IO.puts("Loaded #{n} samples across #{map_size(band_counts)} bands:")

    band_counts
    |> Enum.sort_by(fn {band, _} -> band end)
    |> Enum.each(fn {band, count} ->
      IO.puts("  #{div(band, 1000)} GHz: #{count} samples")
    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)
    IO.puts("Training complete. Final loss: #{Float.round(train_metrics.final_loss, 6)}\n")

    # 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, 4)}")
    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, 4)}")
    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(sample_size) do
    bands = BandConfig.all_freqs()
    per_band = div(sample_size, length(bands))

    band_queries =
      Enum.map_join(bands, "\nUNION ALL\n", fn band_mhz ->
        """
        (SELECT
          ps.score, ps.band_mhz,
          h.surface_temp_c, h.surface_dewpoint_c, h.surface_pressure_mb,
          h.min_refractivity_gradient, h.hpbl_m, h.pwat_mm,
          EXTRACT(HOUR FROM ps.valid_time) AS utc_hour,
          EXTRACT(MONTH FROM ps.valid_time) AS month,
          ps.lon
        FROM propagation_scores ps
        JOIN hrrr_profiles h
          ON h.lat = ps.lat AND h.lon = ps.lon AND h.valid_time = ps.valid_time
        WHERE ps.band_mhz = #{band_mhz}
          AND h.surface_temp_c IS NOT NULL
          AND h.surface_dewpoint_c IS NOT NULL
          AND h.surface_pressure_mb IS NOT NULL
        ORDER BY RANDOM()
        LIMIT #{per_band})
        """
      end)

    %{rows: rows} = Repo.query!(band_queries, [], timeout: @query_timeout)

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

    {feature_rows, target_rows} =
      rows
      |> Enum.map(fn [score, band_mhz, temp_c, dewpoint_c, pressure_mb, grad, hpbl, pwat, utc_hour, month, lon] ->
        features =
          Model.encode_features(%{
            surface_temp_c: to_float(temp_c),
            surface_dewpoint_c: to_float(dewpoint_c),
            surface_pressure_mb: to_float(pressure_mb),
            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
          })

        target = [score / 100.0]

        {features, target}
      end)
      |> Enum.unzip()

    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