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 using all available data:

  Phase 1 (pre-train): Learn broad atmospheric patterns from 500K+ algorithm
  scores covering all seasons, times, and locations across CONUS. Joined to
  solar indices and nearest sounding for stability indices.

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

  Features (20): atmospheric (temp, dewpoint, pressure, abs humidity, td depression,
  refractivity gradient, HPBL, PWAT, surface refractivity), geographic (latitude),
  solar (SFI, Kp max), HRRR-derived (ducting), sounding-derived (K-index,
  lifted index), temporal (solar hour sin/cos, month sin/cos), frequency (log MHz).

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

  @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,
          no_pretrain: :boolean
        ]
      )

    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)
    skip_pretrain = Keyword.get(opts, :no_pretrain, false)

    "=" |> String.duplicate(70) |> IO.puts()
    IO.puts("PROPAGATION MODEL TRAINING (20 features)")
    "=" |> String.duplicate(70) |> IO.puts()

    pretrained_state =
      if skip_pretrain do
        IO.puts("Skipping pre-training (--no-pretrain)")
        nil
      else
        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 + solar + soundings")
        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} 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)...")

        {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(state, pt_val_x, pt_val_y))
        state
      end

    # ── QSO training ────────────────────────────────────────────────
    IO.puts("\n" <> "=" <> String.duplicate("─", 50))

    if pretrained_state do
      IO.puts("PHASE 2: Fine-tuning on QSO + HRRR + solar + soundings")
    else
      IO.puts("TRAINING: QSO + HRRR + solar + soundings (direct, lr=0.001)")
    end

    IO.puts("=" <> String.duplicate("─", 50))

    {ft_features, ft_targets, ft_counts} = load_contact_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)

    lr = if pretrained_state, do: 0.0003, else: 0.001
    IO.puts("\nTraining for #{finetune_epochs} epochs (lr=#{lr})...")

    train_opts = [epochs: finetune_epochs, batch_size: batch_size, learning_rate: lr]
    train_opts = if pretrained_state, do: [{:initial_state, pretrained_state} | train_opts], else: train_opts

    {trained_state, ft_metrics} = Model.train(ft_train_x, ft_train_y, train_opts)

    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 + solar + nearest sounding ─────

  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)...")

    # Join propagation_scores → hrrr_profiles → solar_indices → nearest sounding
    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, h.ducting_detected,
          ps.lat, ps.lon,
          EXTRACT(HOUR FROM ps.valid_time) AS utc_hour,
          EXTRACT(MONTH FROM ps.valid_time) AS month,
          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 propagation_scores ps
        JOIN hrrr_profiles h
          ON h.lat = ps.lat AND h.lon = ps.lon AND h.valid_time = ps.valid_time
        LEFT JOIN solar_indices si ON si.date = ps.valid_time::date
        LEFT JOIN LATERAL (
          SELECT s.k_index, s.lifted_index
          FROM soundings s
          WHERE s.observed_at BETWEEN ps.valid_time - interval '6 hours'
                                  AND ps.valid_time + interval '6 hours'
          ORDER BY ABS(EXTRACT(EPOCH FROM s.observed_at - ps.valid_time))
          LIMIT 1
        ) snd ON true
        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(&encode_pretrain_row/1)
      |> Enum.unzip()

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

  defp encode_pretrain_row([
         score,
         band_mhz,
         temp_c,
         dewpoint_c,
         pressure_mb,
         grad,
         hpbl,
         pwat,
         refractivity,
         ducting,
         lat,
         lon,
         utc_hour,
         month,
         sfi,
         kp_max,
         k_index,
         lifted_index
       ]) do
    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),
        sfi: to_float(sfi),
        kp_max: to_float(kp_max),
        ducting_detected: ducting,
        k_index: to_float(k_index),
        lifted_index: to_float(lifted_index),
        utc_hour: to_float(utc_hour),
        month: trunc(to_float(month)),
        longitude: to_float(lon),
        freq_mhz: band_mhz
      })

    {features, [score / 100.0]}
  end

  # ── Phase 2 data: QSO-HRRR + solar + sounding ────────────────────

  defp load_contact_data do
    IO.puts("  Running QSO-HRRR-solar-sounding 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,
      COALESCE(h1.ducting_detected, false) OR COALESCE(h2.ducting_detected, false) AS ducting_either,
      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 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)
    LEFT JOIN solar_indices si ON si.date = q.qso_timestamp::date
    LEFT JOIN LATERAL (
      SELECT s.k_index, s.lifted_index
      FROM soundings s
      WHERE s.observed_at BETWEEN q.qso_timestamp - interval '6 hours'
                              AND q.qso_timestamp + interval '6 hours'
      ORDER BY ABS(EXTRACT(EPOCH FROM s.observed_at - q.qso_timestamp))
      LIMIT 1
    ) snd ON true
    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+solar+sounding")

    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 {row, percentile} -> encode_contact_row(row, 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

  defp encode_contact_row(
         [
           band_mhz,
           _dist,
           utc_hour,
           month,
           lon,
           lat,
           temp,
           dewpoint,
           pressure,
           grad,
           hpbl,
           pwat,
           refractivity,
           ducting,
           sfi,
           kp_max,
           k_index,
           lifted_index
         ],
         percentile
       ) do
    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),
        sfi: to_float(sfi),
        kp_max: to_float(kp_max),
        ducting_detected: ducting,
        k_index: to_float(k_index),
        lifted_index: to_float(lifted_index),
        utc_hour: to_float(utc_hour),
        month: trunc(to_float(month)),
        longitude: to_float(lon),
        freq_mhz: band_mhz
      })

    {features, [percentile]}
  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