defmodule Mix.Tasks.PropagationTrain do
  @shortdoc "Train propagation ML model from month-balanced HRRR + algorithm scores"

  @moduledoc """
  Trains the propagation Nx/Axon model to predict the physical algorithm
  score from weather conditions alone.

  ## Why not use QSO outcomes as the target?

  An earlier version of this task fine-tuned on "within-band distance
  percentile" from the `contacts` table. That target is structurally
  biased: non-contest months (Feb, Mar, Nov) have almost no QSOs, so the
  model learned "these conditions → QSO happened → good" and implicitly
  "no QSO happened → bad." The result was a model that systematically
  under-predicted propagation during months when nobody is on the air,
  which is exactly the wrong thing for a forecast.

  This task instead:

    1. Samples `hrrr_profiles` stratified *uniformly by calendar month*
       — every month contributes the same number of rows regardless of
       contest-driven density. Contest months still exist in the data,
       they just don't drown out the quiet months.
    2. Joins each sample to `solar_indices` (same date) and the nearest
       `soundings` row (±6 h) for K-index / lifted-index features.
    3. Computes the target via `Scorer.composite_score/2` — the *same*
       physical scoring algorithm that `PropagationGridWorker` uses in
       production. Sky / wind / rain fall back to neutral defaults
       because `hrrr_profiles` doesn't store them; the model learns the
       weather-only residual and the missing factors remain constant.
    4. Trains on every band for every sampled profile, producing one
       training row per (profile, band) pair.

  ## Usage

      mix propagation_train
      mix propagation_train --samples-per-month 20000 --epochs 100
  """

  use Mix.Task

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

  @query_timeout 600_000

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

    # Boot the app without starting Oban queues so training doesn't have
    # the HRRR / NARR / weather cron firing in the background against the
    # local DB we're reading from.
    Application.put_env(
      :microwaveprop,
      Oban,
      Keyword.put(Application.get_env(:microwaveprop, Oban, []), :queues, false)
    )

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

    samples_per_month = Keyword.get(opts, :samples_per_month, 20_000)
    epochs = Keyword.get(opts, :epochs, 100)
    batch_size = Keyword.get(opts, :batch_size, 256)
    learning_rate = Keyword.get(opts, :learning_rate, 0.001)

    print_header(samples_per_month, epochs, batch_size, learning_rate)

    {features, targets, month_counts} = load_training_data(samples_per_month)

    n = elem(Nx.shape(features), 0)
    IO.puts("Loaded #{n} training rows (one per profile × band).")
    print_month_counts(month_counts)

    {features, targets} = shuffle(features, targets, n)
    {train_x, train_y, val_x, val_y, test_x, test_y} = split(features, targets, n)

    IO.puts("\nTraining for #{epochs} epochs (lr=#{learning_rate}, batch=#{batch_size})...")

    {trained_state, metrics} =
      Model.train(train_x, train_y,
        epochs: epochs,
        batch_size: batch_size,
        learning_rate: learning_rate
      )

    print_loss(metrics.final_loss)
    print_eval("Validation", Model.evaluate(trained_state, val_x, val_y))
    print_eval("Test", Model.evaluate(trained_state, test_x, test_y))

    IO.puts("\nSaving model to priv/models/propagation_v1.nx...")
    Model.save(trained_state)
    IO.puts("Done.")

    _ = check_monthly_bias(trained_state)

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

  # ── Data loading ────────────────────────────────────────────────────

  # Returns {features_tensor, targets_tensor, %{month => count}}.
  # Stratifies by calendar month so every month contributes `per_month`
  # profile rows, then explodes each profile into one row per band.
  defp load_training_data(per_month) do
    IO.puts("Sampling hrrr_profiles stratified by month (#{per_month}/month)...")

    # The stratification is the whole point: random_row_in_month ranks each
    # profile within its (year, month) bucket by a random key, then we keep
    # the first `per_month` rows per month across all years. This is much
    # cheaper than ORDER BY random() over the whole 43M-row table because
    # the window function can stream.
    sql = """
    WITH ranked AS (
      SELECT
        h.valid_time,
        h.lat,
        h.lon,
        h.surface_temp_c,
        h.surface_dewpoint_c,
        h.surface_pressure_mb,
        h.min_refractivity_gradient,
        h.hpbl_m,
        h.pwat_mm,
        h.surface_refractivity,
        h.ducting_detected,
        extract(month from h.valid_time)::int AS month,
        extract(hour from h.valid_time)::int  AS utc_hour,
        ROW_NUMBER() OVER (
          PARTITION BY extract(month from h.valid_time)::int
          ORDER BY random()
        ) AS rn
      FROM hrrr_profiles h
      WHERE h.surface_temp_c IS NOT NULL
        AND h.surface_dewpoint_c IS NOT NULL
        AND h.surface_pressure_mb IS NOT NULL
    )
    SELECT
      r.valid_time, r.lat, r.lon,
      r.surface_temp_c, r.surface_dewpoint_c, r.surface_pressure_mb,
      r.min_refractivity_gradient, r.hpbl_m, r.pwat_mm,
      r.surface_refractivity, r.ducting_detected,
      r.month, r.utc_hour,
      si.sfi,
      COALESCE((SELECT MAX(v) FROM unnest(si.kp_values) AS v), 2.0) AS kp_max,
      snd.k_index,
      snd.lifted_index
    FROM ranked r
    LEFT JOIN solar_indices si ON si.date = r.valid_time::date
    LEFT JOIN LATERAL (
      SELECT s.k_index, s.lifted_index
      FROM soundings s
      WHERE s.observed_at BETWEEN r.valid_time - interval '6 hours'
                              AND r.valid_time + interval '6 hours'
      ORDER BY ABS(EXTRACT(EPOCH FROM s.observed_at - r.valid_time))
      LIMIT 1
    ) snd ON true
    WHERE r.rn <= $1
    """

    %{rows: rows} = Repo.query!(sql, [per_month], timeout: @query_timeout)
    IO.puts("  #{length(rows)} profile rows after stratified sample")

    month_counts = Enum.frequencies_by(rows, fn row -> Enum.at(row, 11) end)

    bands = BandConfig.all_freqs()

    # Each profile row turns into one training row per band.
    {feature_rows, target_rows} =
      rows
      |> Enum.flat_map(fn row -> encode_profile_rows(row, bands) end)
      |> Enum.unzip()

    {Nx.tensor(feature_rows, type: :f32), Nx.tensor(target_rows, type: :f32), month_counts}
  end

  # Builds one training row per (profile, band). Target is the algorithm
  # composite score normalized to [0, 1]; features are the ML model's
  # 20-element encoding of the same conditions.
  defp encode_profile_rows(row, bands) do
    [
      _valid_time,
      lat,
      lon,
      temp_c,
      dewpoint_c,
      pressure_mb,
      grad,
      hpbl_m,
      pwat_mm,
      refractivity,
      ducting,
      month,
      utc_hour,
      sfi,
      kp_max,
      k_index,
      lifted_index
    ] = row

    temp_c = to_float(temp_c)
    dewpoint_c = to_float(dewpoint_c)
    pressure_mb = to_float(pressure_mb)
    grad = to_float(grad)
    hpbl_m = to_float(hpbl_m)
    pwat_mm = to_float(pwat_mm)
    refractivity = to_float(refractivity)
    lat_f = to_float(lat)
    lon_f = to_float(lon)
    sfi_f = to_float(sfi)
    kp_f = to_float(kp_max)
    k_idx_f = to_float(k_index)
    li_f = to_float(lifted_index)

    scorer_conditions =
      build_scorer_conditions(
        temp_c,
        dewpoint_c,
        pressure_mb,
        grad,
        hpbl_m,
        pwat_mm,
        lat_f,
        lon_f,
        trunc(to_float(month)),
        trunc(to_float(utc_hour))
      )

    Enum.map(bands, fn freq_mhz ->
      band_config = BandConfig.get(freq_mhz)
      %{score: score} = Scorer.composite_score(scorer_conditions, band_config)

      features =
        Model.encode_features(%{
          surface_temp_c: temp_c,
          surface_dewpoint_c: dewpoint_c,
          surface_pressure_mb: pressure_mb,
          min_refractivity_gradient: grad,
          hpbl_m: hpbl_m,
          pwat_mm: pwat_mm,
          surface_refractivity: refractivity,
          latitude: lat_f,
          sfi: sfi_f,
          kp_max: kp_f,
          ducting_detected: ducting,
          k_index: k_idx_f,
          lifted_index: li_f,
          utc_hour: to_float(utc_hour),
          longitude: lon_f,
          month: trunc(to_float(month)),
          freq_mhz: freq_mhz
        })

      {features, [score / 100.0]}
    end)
  end

  # Builds a conditions map matching the shape Scorer.composite_score/2
  # expects. Sky / wind / rain default to nil → Scorer uses neutral
  # values (50, 50, 100). The ML model's feature set also omits those
  # channels, so the missing-factor contribution is constant across all
  # samples and the model learns the weather-derived residual cleanly.
  defp build_scorer_conditions(temp_c, dewpoint_c, pressure_mb, grad, hpbl_m, pwat_mm, lat, lon, month, utc_hour) do
    %{
      abs_humidity: Scorer.absolute_humidity(temp_c, dewpoint_c),
      temp_f: Scorer.c_to_f(temp_c),
      dewpoint_f: Scorer.c_to_f(dewpoint_c),
      pressure_mb: pressure_mb,
      prev_pressure_mb: nil,
      min_refractivity_gradient: grad,
      bl_depth_m: hpbl_m,
      pwat_mm: pwat_mm,
      sky_cover_pct: nil,
      wind_speed_kts: nil,
      rain_rate_mmhr: 0.0,
      utc_hour: utc_hour,
      utc_minute: 0,
      month: month,
      longitude: lon,
      latitude: lat,
      best_duct_band_ghz: nil
    }
  end

  # ── Post-training sanity check ──────────────────────────────────────

  # After training, evaluate the model on a synthetic set of
  # (temp, month) combinations at otherwise-identical conditions. The
  # goal is to confirm the model doesn't prefer / penalize specific
  # months independently of the weather they carry — i.e. that a
  # generic March day doesn't rank below a generic June day just
  # because the training data had fewer March contacts.
  defp check_monthly_bias(params) do
    IO.puts("\nMonthly sanity check at fixed conditions (10 GHz, 15°C, 10°C dewpoint)...")
    predict_fn = Model.compile_predict()

    for month <- 1..12 do
      conditions = %{
        surface_temp_c: 15.0,
        surface_dewpoint_c: 10.0,
        surface_pressure_mb: 1013.0,
        min_refractivity_gradient: -70.0,
        hpbl_m: 800.0,
        pwat_mm: 20.0,
        surface_refractivity: 320.0,
        latitude: 32.5,
        longitude: -97.0,
        sfi: 120.0,
        kp_max: 2.0,
        ducting_detected: false,
        k_index: 25.0,
        lifted_index: 0.0,
        utc_hour: 12,
        month: month,
        freq_mhz: 10_000
      }

      score = Model.predict_score(predict_fn, params, conditions)
      IO.puts("  month=#{month |> Integer.to_string() |> String.pad_leading(2)}: #{score}")
    end
  end

  # ── Helpers ──────────────────────────────────────────────────────────

  defp print_header(per_month, epochs, batch_size, lr) do
    "=" |> String.duplicate(70) |> IO.puts()
    IO.puts("PROPAGATION MODEL TRAINING — month-balanced HRRR → algorithm score")
    "=" |> String.duplicate(70) |> IO.puts()
    IO.puts("Samples per month: #{per_month}  epochs: #{epochs}  batch: #{batch_size}  lr: #{lr}")
    IO.puts("Started: #{DateTime.to_string(DateTime.utc_now())}\n")
  end

  defp print_month_counts(counts) do
    counts
    |> Enum.sort_by(fn {m, _} -> m end)
    |> Enum.each(fn {m, count} ->
      IO.puts("  month #{m |> Integer.to_string() |> String.pad_leading(2)}: #{count}")
    end)
  end

  defp shuffle(features, targets, n) do
    key = Nx.Random.key(System.os_time())
    {indices, _} = Nx.Random.shuffle(key, Nx.iota({n}))
    indices = Nx.as_type(indices, :s64)
    {Nx.take(features, indices), Nx.take(targets, indices)}
  end

  defp split(features, targets, n) do
    nc = elem(Nx.shape(features), 1)
    train_end = trunc(n * 0.8)
    val_end = trunc(n * 0.9)
    test_size = n - val_end

    IO.puts("  Split: train=#{train_end}, val=#{val_end - train_end}, test=#{test_size}")

    {
      Nx.slice(features, [0, 0], [train_end, nc]),
      Nx.slice(targets, [0, 0], [train_end, 1]),
      Nx.slice(features, [train_end, 0], [val_end - train_end, nc]),
      Nx.slice(targets, [train_end, 0], [val_end - train_end, 1]),
      Nx.slice(features, [val_end, 0], [test_size, nc]),
      Nx.slice(targets, [val_end, 0], [test_size, 1])
    }
  end

  defp print_loss(loss) do
    if is_float(loss) and loss == loss do
      IO.puts("  Final loss: #{Float.round(loss, 6)}")
    else
      IO.puts("  Final loss: NaN (training diverged)")
      System.halt(1)
    end
  end

  defp print_eval(label, metrics) do
    IO.puts("  #{label}: RMSE=#{Float.round(metrics.rmse * 100, 2)} pts, R²=#{Float.round(metrics.r_squared, 4)}")
  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