defmodule Mix.Tasks.PropagationTrain do
  @shortdoc "Train ML model from QSO-HRRR data with algorithm pre-training"

  @moduledoc """
  Two-phase ML training for propagation prediction:

  Phase 1 (pre-train): Learn broad atmospheric patterns from 500K+ algorithm
  scores covering all seasons, times, and locations across CONUS.

  Phase 2 (fine-tune): Calibrate against real QSO outcomes (57K+ contacts
  matched to HRRR conditions). Target is within-band distance percentile.

  Features: 15 inputs (10 atmospheric + latitude + 4 cyclical temporal + frequency).

  ## Usage

      mix propagation_train
      mix propagation_train --pretrain-epochs 50 --finetune-epochs 200
  """

  use Mix.Task

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

  @query_timeout 600_000

  # Per-band normalization for QSO distance percentile computation
  @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: [
          pretrain_epochs: :integer,
          finetune_epochs: :integer,
          pretrain_sample: :integer,
          batch_size: :integer
        ]
      )

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

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

    pretrain_epochs = Keyword.get(opts, :pretrain_epochs, 50)
    finetune_epochs = Keyword.get(opts, :finetune_epochs, 200)
    pretrain_sample = Keyword.get(opts, :pretrain_sample, 500_000)
    batch_size = Keyword.get(opts, :batch_size, 256)

    "=" |> String.duplicate(70) |> IO.puts()
    IO.puts("PROPAGATION MODEL TRAINING (Two-Phase)")
    "=" |> String.duplicate(70) |> IO.puts()
    IO.puts("Phase 1: #{pretrain_epochs} epochs on #{pretrain_sample} algorithm scores")
    IO.puts("Phase 2: #{finetune_epochs} epochs on real QSO data")
    IO.puts("Started: #{DateTime.to_string(DateTime.utc_now())}\n")

    # ── Phase 1: Pre-train on algorithm scores ──────────────────────
    IO.puts("=" <> String.duplicate("─", 50))
    IO.puts("PHASE 1: Pre-training on algorithm scores")
    IO.puts("=" <> String.duplicate("─", 50))

    {pt_features, pt_targets, pt_counts} = load_pretrain_data(pretrain_sample)
    pt_n = elem(Nx.shape(pt_features), 0)
    IO.puts("Loaded #{pt_n} algorithm score samples:")
    print_band_counts(pt_counts)

    {pt_features, pt_targets} = shuffle(pt_features, pt_targets, pt_n)
    {pt_train_x, pt_train_y, pt_val_x, pt_val_y, _, _} = split(pt_features, pt_targets, pt_n)

    IO.puts("\nPre-training for #{pretrain_epochs} epochs (lr=0.001)...")

    {pretrained_state, pt_metrics} =
      Model.train(pt_train_x, pt_train_y, epochs: pretrain_epochs, batch_size: batch_size)

    print_loss(pt_metrics.final_loss)
    print_eval("Pre-train val", Model.evaluate(pretrained_state, pt_val_x, pt_val_y))

    # ── Phase 2: Fine-tune on real QSO data ─────────────────────────
    IO.puts("\n" <> "=" <> String.duplicate("─", 50))
    IO.puts("PHASE 2: Fine-tuning on real QSO-HRRR data")
    IO.puts("=" <> String.duplicate("─", 50))

    {ft_features, ft_targets, ft_counts} = load_qso_data()
    ft_n = elem(Nx.shape(ft_features), 0)
    IO.puts("Loaded #{ft_n} QSO training samples:")
    print_band_counts(ft_counts)

    {ft_features, ft_targets} = shuffle(ft_features, ft_targets, ft_n)
    {ft_train_x, ft_train_y, ft_val_x, ft_val_y, ft_test_x, ft_test_y} = split(ft_features, ft_targets, ft_n)

    IO.puts("\nFine-tuning for #{finetune_epochs} epochs (lr=0.0003)...")

    {trained_state, ft_metrics} =
      Model.train(ft_train_x, ft_train_y,
        epochs: finetune_epochs,
        batch_size: batch_size,
        learning_rate: 0.0003,
        initial_state: pretrained_state
      )

    print_loss(ft_metrics.final_loss)
    print_eval("Fine-tune val", Model.evaluate(trained_state, ft_val_x, ft_val_y))
    print_eval("Fine-tune test", Model.evaluate(trained_state, ft_test_x, ft_test_y))

    # ── Save ────────────────────────────────────────────────────────
    IO.puts("\nSaving 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

  # ── Phase 1 data: algorithm scores from propagation_scores ────────

  defp load_pretrain_data(sample_size) do
    bands = BandConfig.all_freqs()
    per_band = div(sample_size, length(bands))

    IO.puts("  Loading algorithm scores (stratified #{per_band}/band)...")

    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,
          h.surface_refractivity,
          ps.lat,
          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,
                       refractivity,
                       lat,
                       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),
            surface_refractivity: to_float(refractivity),
            latitude: to_float(lat),
            utc_hour: to_float(utc_hour),
            month: trunc(to_float(month)),
            longitude: to_float(lon),
            freq_mhz: band_mhz
          })

        {features, [score / 100.0]}
      end)
      |> Enum.unzip()

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

  # ── Phase 2 data: real QSO-HRRR matches ──────────────────────────

  defp load_qso_data do
    IO.puts("  Running QSO-HRRR join query...")

    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,
      ((q.pos1->>'lat')::float + (q.pos2->>'lat')::float) / 2.0 AS avg_lat,
      (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,
      (COALESCE(h1.surface_refractivity, h2.surface_refractivity) +
       COALESCE(h2.surface_refractivity, h1.surface_refractivity)) / 2.0 AS avg_refractivity
    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
    """

    %{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)

    # Compute within-band distance percentile
    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_rows =
      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 = Enum.count(distances, &(&1 <= distance_km))
          {row, rank / n}
        end)
      end)

    {feature_rows, target_rows} =
      percentile_rows
      |> Enum.map(fn {[
                        band_mhz,
                        _dist,
                        utc_hour,
                        month,
                        lon,
                        lat,
                        temp,
                        dewpoint,
                        pressure,
                        grad,
                        hpbl,
                        pwat,
                        refractivity
                      ], 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),
            surface_refractivity: to_float(refractivity),
            latitude: to_float(lat),
            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")

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

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

  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_band_counts(counts) do
    counts
    |> Enum.sort_by(fn {band, _} -> band end)
    |> Enum.each(fn {band, count} ->
      IO.puts("  #{div(band, 1000)} GHz: #{count}")
    end)
  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