defmodule Microwaveprop.Propagation.RecalibratorTest do
  use Microwaveprop.DataCase, async: true

  alias Microwaveprop.Propagation.Recalibrator
  alias Microwaveprop.Radio.Contact
  alias Microwaveprop.Weather.HrrrProfile

  @factor_keys ~w(humidity time_of_day td_depression refractivity sky season wind rain pwat pressure)a

  defp create_contact(attrs) do
    ts = Map.get(attrs, :qso_timestamp, ~U[2024-07-15 06:00:00Z])
    lat = Map.get(attrs, :lat, 32.9)
    lon = Map.get(attrs, :lon, -97.0)

    contact_attrs = %{
      station1: Map.fetch!(attrs, :station1),
      station2: "K5TR",
      qso_timestamp: ts,
      mode: "CW",
      band: Decimal.new("10000"),
      grid1: "EM12",
      grid2: "EM00",
      pos1: %{"lat" => lat, "lon" => lon},
      pos2: %{"lat" => 30.3, "lon" => -97.7},
      distance_km: Decimal.new("295")
    }

    {:ok, contact} =
      %Contact{}
      |> Contact.changeset(contact_attrs)
      |> Repo.insert()

    contact
  end

  defp create_hrrr_profile(attrs) do
    {:ok, profile} =
      %HrrrProfile{}
      |> HrrrProfile.changeset(%{
        valid_time: Map.fetch!(attrs, :valid_time),
        lat: Map.fetch!(attrs, :lat),
        lon: Map.fetch!(attrs, :lon),
        surface_temp_c: Map.get(attrs, :surface_temp_c, 28.0),
        surface_dewpoint_c: Map.get(attrs, :surface_dewpoint_c, 22.0),
        surface_pressure_mb: Map.get(attrs, :surface_pressure_mb, 1005.0),
        min_refractivity_gradient: Map.get(attrs, :min_refractivity_gradient, -120.0),
        hpbl_m: Map.get(attrs, :hpbl_m, 800.0),
        pwat_mm: Map.get(attrs, :pwat_mm, 35.0)
      })
      |> Repo.insert()

    profile
  end

  # Synthetic factor vectors for training tests.
  # Positives: high scores (good propagation conditions).
  # Negatives: low scores (poor propagation conditions).
  defp synthetic_positives do
    for _ <- 1..20 do
      [85, 90, 80, 75, 88, 70, 95, 100, 75, 82]
      |> Enum.map(fn base -> base + :rand.uniform(10) - 5 end)
      |> Enum.map(&min(100, max(0, &1)))
    end
  end

  defp synthetic_negatives do
    for _ <- 1..20 do
      [40, 35, 45, 50, 30, 55, 50, 60, 50, 45]
      |> Enum.map(fn base -> base + :rand.uniform(10) - 5 end)
      |> Enum.map(&min(100, max(0, &1)))
    end
  end

  describe "train/3" do
    test "returns a weights map with all 10 factor keys summing to ~1.0" do
      result = Recalibrator.train(synthetic_positives(), synthetic_negatives(), epochs: 50, learning_rate: 0.01)

      assert is_map(result.weights)

      for key <- @factor_keys do
        assert Map.has_key?(result.weights, key), "missing weight for #{key}"
        weight = result.weights[key]
        assert is_float(weight), "weight for #{key} should be float, got #{inspect(weight)}"
        assert weight >= 0.0, "weight for #{key} should be non-negative"
      end

      sum = result.weights |> Map.values() |> Enum.sum()
      assert_in_delta sum, 1.0, 0.001, "weights should sum to 1.0, got #{sum}"
    end

    test "train_loss decreases from initial loss" do
      result = Recalibrator.train(synthetic_positives(), synthetic_negatives(), epochs: 200, learning_rate: 0.01)

      assert is_float(result.train_loss)
      assert is_float(result.val_loss)
      assert is_float(result.initial_loss)

      assert result.train_loss < result.initial_loss,
             "train_loss (#{result.train_loss}) should be less than initial_loss (#{result.initial_loss})"
    end
  end

  describe "fit/1" do
    test "returns valid result with insufficient HRRR data (falls back to current weights)" do
      # Contacts without matching HRRR profiles for baselines
      create_contact(%{station1: "W5AA"})
      create_contact(%{station1: "W5BB", lat: 33.0, lon: -96.0, qso_timestamp: ~U[2024-07-16 06:00:00Z]})

      result = Recalibrator.fit(sample_size: 5, epochs: 20, learning_rate: 0.01)

      assert is_map(result.weights)
      assert map_size(result.weights) == 10

      sum = result.weights |> Map.values() |> Enum.sum()
      assert_in_delta sum, 1.0, 0.01
    end
  end

  describe "compute_factors/2" do
    test "builds a factor vector from an HRRR profile and timestamp" do
      profile =
        create_hrrr_profile(%{
          valid_time: ~U[2024-07-15 06:00:00Z],
          lat: 32.9,
          lon: -97.0,
          surface_temp_c: 28.0,
          surface_dewpoint_c: 22.0,
          surface_pressure_mb: 1005.0,
          min_refractivity_gradient: -120.0,
          hpbl_m: 800.0,
          pwat_mm: 35.0
        })

      factors = Recalibrator.compute_factors(profile, ~U[2024-07-15 06:00:00Z])

      assert is_list(factors)
      assert length(factors) == 10

      Enum.each(factors, fn f ->
        assert is_number(f)
        assert f >= 0 and f <= 100
      end)
    end
  end
end