defmodule Microwaveprop.Propagation.AsosNudgeTest do
  use ExUnit.Case, async: true

  alias Microwaveprop.Propagation.AsosNudge

  @valid_time ~U[2026-04-12 18:00:00Z]

  defp profile(lat, lon, overrides \\ []) do
    Map.merge(
      %{
        valid_time: @valid_time,
        lat: lat,
        lon: lon,
        surface_temp_c: 20.0,
        surface_dewpoint_c: 10.0,
        surface_pressure_mb: 1013.0,
        hpbl_m: 1000.0,
        pwat_mm: 25.0,
        min_refractivity_gradient: -40.0,
        surface_refractivity: 320.0,
        profile: []
      },
      Map.new(overrides)
    )
  end

  defp station(lat, lon, overrides \\ []) do
    Map.merge(
      %{
        lat: lat,
        lon: lon,
        temp_f: 68.0,
        dewpoint_f: 50.0,
        sea_level_pressure_mb: 1013.0,
        precip_1h_in: 0.0
      },
      Map.new(overrides)
    )
  end

  describe "build_station_residuals/2" do
    test "station co-located with a HRRR grid cell yields residuals against that cell" do
      station_obs = station(32.0, -97.0, temp_f: 77.0, dewpoint_f: 59.0, sea_level_pressure_mb: 1020.0)
      hrrr = profile(32.0, -97.0, surface_temp_c: 20.0, surface_dewpoint_c: 10.0, surface_pressure_mb: 1013.0)

      [residual] = AsosNudge.build_station_residuals([station_obs], [hrrr])

      assert_in_delta residual.dtemp_c, 5.0, 0.1
      assert_in_delta residual.ddewpoint_c, 5.0, 0.1
      assert_in_delta residual.dpressure_mb, 7.0, 0.1
      assert residual.lat == 32.0
      assert residual.lon == -97.0
    end

    test "station with no HRRR cell within 1 grid step is dropped" do
      obs = station(32.0, -97.0, temp_f: 77.0)
      far_hrrr = profile(40.0, -85.0, surface_temp_c: 20.0)

      assert AsosNudge.build_station_residuals([obs], [far_hrrr]) == []
    end

    test "station with nil temp is dropped" do
      obs = station(32.0, -97.0, temp_f: nil)
      hrrr = profile(32.0, -97.0)

      assert AsosNudge.build_station_residuals([obs], [hrrr]) == []
    end

    test "station with nil pressure still contributes temp and dewpoint deltas" do
      obs = station(32.0, -97.0, temp_f: 77.0, dewpoint_f: 59.0, sea_level_pressure_mb: nil)
      hrrr = profile(32.0, -97.0, surface_temp_c: 20.0, surface_dewpoint_c: 10.0, surface_pressure_mb: 1013.0)

      [residual] = AsosNudge.build_station_residuals([obs], [hrrr])

      assert_in_delta residual.dtemp_c, 5.0, 0.1
      assert_in_delta residual.ddewpoint_c, 5.0, 0.1
      assert residual.dpressure_mb == nil
    end
  end

  describe "nudge_profile/2" do
    test "single station at the same grid cell applies (almost) the full residual" do
      residuals = [
        %{lat: 32.0, lon: -97.0, dtemp_c: 5.0, ddewpoint_c: 3.0, dpressure_mb: 2.0, asos_rain_mmhr: 0.0}
      ]

      hrrr = profile(32.0, -97.0, surface_temp_c: 20.0, surface_dewpoint_c: 10.0, surface_pressure_mb: 1013.0)
      nudged = AsosNudge.nudge_profile(hrrr, residuals)

      assert_in_delta nudged.surface_temp_c, 25.0, 0.1
      assert_in_delta nudged.surface_dewpoint_c, 13.0, 0.1
      assert_in_delta nudged.surface_pressure_mb, 1015.0, 0.1
    end

    test "station farther than 250 km contributes nothing" do
      residuals = [
        %{lat: 40.0, lon: -97.0, dtemp_c: 10.0, ddewpoint_c: 10.0, dpressure_mb: 5.0, asos_rain_mmhr: 0.0}
      ]

      hrrr = profile(32.0, -97.0, surface_temp_c: 20.0)
      nudged = AsosNudge.nudge_profile(hrrr, residuals)

      assert nudged.surface_temp_c == 20.0
      assert nudged.surface_dewpoint_c == 10.0
      assert nudged.surface_pressure_mb == 1013.0
    end

    test "two equidistant stations with opposite residuals average to near zero" do
      residuals = [
        %{lat: 32.0, lon: -95.93, dtemp_c: 5.0, ddewpoint_c: 0.0, dpressure_mb: 0.0, asos_rain_mmhr: 0.0},
        %{lat: 32.0, lon: -98.07, dtemp_c: -5.0, ddewpoint_c: 0.0, dpressure_mb: 0.0, asos_rain_mmhr: 0.0}
      ]

      hrrr = profile(32.0, -97.0, surface_temp_c: 20.0)
      nudged = AsosNudge.nudge_profile(hrrr, residuals)

      assert_in_delta nudged.surface_temp_c, 20.0, 0.1
    end

    test "nil residual components (missing pressure) leave the HRRR field unchanged" do
      residuals = [
        %{lat: 32.0, lon: -97.0, dtemp_c: 5.0, ddewpoint_c: 3.0, dpressure_mb: nil, asos_rain_mmhr: 0.0}
      ]

      hrrr = profile(32.0, -97.0, surface_temp_c: 20.0, surface_dewpoint_c: 10.0, surface_pressure_mb: 1013.0)
      nudged = AsosNudge.nudge_profile(hrrr, residuals)

      assert_in_delta nudged.surface_temp_c, 25.0, 0.1
      assert_in_delta nudged.surface_dewpoint_c, 13.0, 0.1
      assert nudged.surface_pressure_mb == 1013.0
    end

    test "upper-air fields pass through unchanged regardless of residuals" do
      residuals = [
        %{lat: 32.0, lon: -97.0, dtemp_c: 10.0, ddewpoint_c: 10.0, dpressure_mb: 20.0, asos_rain_mmhr: 0.0}
      ]

      hrrr = profile(32.0, -97.0, min_refractivity_gradient: -250.0, pwat_mm: 35.0, hpbl_m: 500.0)
      nudged = AsosNudge.nudge_profile(hrrr, residuals)

      assert nudged.min_refractivity_gradient == -250.0
      assert nudged.pwat_mm == 35.0
      assert nudged.hpbl_m == 500.0
    end
  end

  describe "compute/3" do
    test "returns empty list when observations list is empty" do
      assert AsosNudge.compute([], @valid_time, [profile(32.0, -97.0)]) == []
    end

    test "returns empty list when no HRRR profiles provided" do
      assert AsosNudge.compute([station(32.0, -97.0)], @valid_time, []) == []
    end

    test "returns grid score maps for nudged profiles with band_mhz, score, factors" do
      hrrr = profile(32.0, -97.0, surface_temp_c: 20.0)
      obs = station(32.0, -97.0, temp_f: 77.0)

      results = AsosNudge.compute([obs], @valid_time, [hrrr])

      assert results != []

      for result <- results do
        assert result.lat == 32.0
        assert result.lon == -97.0
        assert result.valid_time == @valid_time
        assert is_integer(result.band_mhz)
        assert is_number(result.score)
        assert is_map(result.factors)
      end
    end

    test "drops grid points with no station within 250 km from the output" do
      nudged_cell = profile(32.0, -97.0, surface_temp_c: 20.0)
      far_cell = profile(40.0, -85.0, surface_temp_c: 20.0)
      obs = station(32.0, -97.0, temp_f: 77.0)

      results = AsosNudge.compute([obs], @valid_time, [nudged_cell, far_cell])

      lat_lons = results |> Enum.map(&{&1.lat, &1.lon}) |> Enum.uniq()
      assert {32.0, -97.0} in lat_lons
      refute {40.0, -85.0} in lat_lons
    end

    test "re-scores grid cells with MRMS rain even when no ASOS station is nearby" do
      far_cell = profile(40.0, -85.0, surface_temp_c: 20.0)
      rain_grid = %{{40.0, -85.0} => 12.5}

      results = AsosNudge.compute([], @valid_time, [far_cell], rain_grid)

      assert results != []
      lat_lons = results |> Enum.map(&{&1.lat, &1.lon}) |> Enum.uniq()
      assert {40.0, -85.0} in lat_lons
    end

    test "ignores MRMS rain below the 0.1 mm/hr threshold so dry cells keep raw HRRR scores" do
      cell = profile(40.0, -85.0, surface_temp_c: 20.0)
      rain_grid = %{{40.0, -85.0} => 0.05}

      assert AsosNudge.compute([], @valid_time, [cell], rain_grid) == []
    end

    test "MRMS rain is applied through to the scored profile when above threshold" do
      cell = profile(40.0, -85.0, surface_temp_c: 20.0)
      rain_grid_dry = %{}
      rain_grid_wet = %{{40.0, -85.0} => 15.0}
      obs = station(40.0, -85.0, temp_f: 68.0)

      [dry_result | _] = AsosNudge.compute([obs], @valid_time, [cell], rain_grid_dry)
      [wet_result | _] = AsosNudge.compute([obs], @valid_time, [cell], rain_grid_wet)

      # Heavy rain should drop the score meaningfully vs the same cell with
      # no rain. (Exact delta depends on the band weights but the 0→15 mm/hr
      # transition always reduces the rain factor score.)
      assert wet_result.factors.rain < dry_result.factors.rain
    end
  end
end