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

  alias Microwaveprop.Propagation.Model

  describe "build/0" do
    test "returns an Axon model" do
      model = Model.build()
      assert %Axon{} = model
    end
  end

  describe "init/0" do
    test "returns initialized model state with parameters" do
      params = Model.init()
      assert %Axon.ModelState{} = params
      assert Map.has_key?(params.data, "hidden_1")
      assert Map.has_key?(params.data, "hidden_2")
      assert Map.has_key?(params.data, "hidden_3")
      assert Map.has_key?(params.data, "output")
    end
  end

  describe "save/2 and load/1" do
    @tag :tmp_dir
    test "round-trips parameters to disk", %{tmp_dir: tmp_dir} do
      path = Path.join(tmp_dir, "test_model.nx")
      params = Model.init()

      assert :ok = Model.save(params, path)
      assert File.exists?(path)

      assert {:ok, loaded} = Model.load(path)

      # Predict with both and compare
      input = Nx.broadcast(0.5, {1, 20})
      original = Model.predict(params, input)
      reloaded = Model.predict(loaded, input)

      assert Nx.to_flat_list(original) == Nx.to_flat_list(reloaded)
    end

    test "load returns :error for missing file" do
      assert :error = Model.load("/tmp/nonexistent_model.nx")
    end
  end

  describe "load_or_init/1" do
    @tag :tmp_dir
    test "loads from file when it exists", %{tmp_dir: tmp_dir} do
      path = Path.join(tmp_dir, "test_model.nx")
      params = Model.init()
      Model.save(params, path)

      loaded = Model.load_or_init(path)
      input = Nx.broadcast(0.5, {1, 20})
      assert Nx.to_flat_list(Model.predict(params, input)) == Nx.to_flat_list(Model.predict(loaded, input))
    end

    test "initializes when file missing" do
      params = Model.load_or_init("/tmp/nonexistent_model.nx")
      assert %Axon.ModelState{} = params
    end
  end

  describe "predict/2" do
    test "produces output in [0, 1] range" do
      params = Model.init()
      # Single sample, 20 features
      input = Nx.broadcast(0.5, {1, 20})
      output = Model.predict(params, input)

      assert Nx.shape(output) == {1, 1}

      value = output |> Nx.to_flat_list() |> hd()
      assert value >= 0.0 and value <= 1.0
    end

    test "handles batch of multiple samples" do
      params = Model.init()
      input = Nx.broadcast(0.5, {4, 20})
      output = Model.predict(params, input)

      assert Nx.shape(output) == {4, 1}
    end
  end

  describe "encode_features/1" do
    test "returns 20-element list" do
      features =
        Model.encode_features(%{
          surface_temp_c: 25.0,
          surface_dewpoint_c: 15.0,
          surface_pressure_mb: 1015.0,
          min_refractivity_gradient: -80.0,
          hpbl_m: 600.0,
          pwat_mm: 25.0,
          utc_hour: 18,
          month: 7,
          freq_mhz: 10_000
        })

      assert length(features) == 20
      assert Enum.all?(features, &is_float/1)
    end

    test "uses defaults for missing values" do
      features = Model.encode_features(%{})
      assert length(features) == 20
      assert Enum.all?(features, &is_float/1)
    end

    test "cyclical encoding: hour 0 and hour 24 produce same sin/cos" do
      f0 = Model.encode_features(%{utc_hour: 0})
      f24 = Model.encode_features(%{utc_hour: 24})

      # sin/cos at indices 8, 9
      assert_in_delta Enum.at(f0, 8), Enum.at(f24, 8), 0.001
      assert_in_delta Enum.at(f0, 9), Enum.at(f24, 9), 0.001
    end

    test "log_freq_mhz is last feature" do
      features = Model.encode_features(%{freq_mhz: 10_000})
      assert_in_delta List.last(features), :math.log(10_000) / 13.0, 0.001
    end
  end

  describe "train/3" do
    test "trains on synthetic data and returns state with metrics" do
      # Generate small synthetic dataset: 64 samples, 13 features
      key = Nx.Random.key(42)
      {features, key} = Nx.Random.uniform(key, shape: {64, 20}, type: :f32)
      {targets, _key} = Nx.Random.uniform(key, shape: {64, 1}, type: :f32)

      {trained_state, metrics} = Model.train(features, targets, epochs: 2, batch_size: 32)

      assert %Axon.ModelState{} = trained_state
      assert is_float(metrics.final_loss)
      assert metrics.final_loss >= 0.0
    end
  end

  describe "evaluate/3" do
    test "returns rmse and r_squared" do
      params = Model.init()
      key = Nx.Random.key(42)
      {features, key} = Nx.Random.uniform(key, shape: {32, 20}, type: :f32)
      {targets, _key} = Nx.Random.uniform(key, shape: {32, 1}, type: :f32)

      metrics = Model.evaluate(params, features, targets)

      assert is_float(metrics.rmse)
      assert metrics.rmse >= 0.0
      assert is_float(metrics.r_squared)
    end
  end

  describe "predict_score/3" do
    test "returns integer between 0 and 100" do
      params = Model.init()
      predict_fn = Model.compile_predict()

      score =
        Model.predict_score(predict_fn, params, %{
          surface_temp_c: 25.0,
          surface_dewpoint_c: 15.0,
          surface_pressure_mb: 1015.0,
          min_refractivity_gradient: -80.0,
          hpbl_m: 600.0,
          pwat_mm: 25.0,
          utc_hour: 18,
          month: 7,
          freq_mhz: 10_000
        })

      assert is_integer(score)
      assert score >= 0
      assert score <= 100
    end

    test "returns different scores for different conditions" do
      params = Model.init()
      predict_fn = Model.compile_predict()

      score_hot =
        Model.predict_score(predict_fn, params, %{
          surface_temp_c: 35.0,
          surface_dewpoint_c: 25.0,
          surface_pressure_mb: 1010.0,
          min_refractivity_gradient: -300.0,
          hpbl_m: 200.0,
          pwat_mm: 40.0,
          utc_hour: 3,
          month: 7,
          freq_mhz: 10_000
        })

      score_cold =
        Model.predict_score(predict_fn, params, %{
          surface_temp_c: -10.0,
          surface_dewpoint_c: -20.0,
          surface_pressure_mb: 1030.0,
          min_refractivity_gradient: -40.0,
          hpbl_m: 1500.0,
          pwat_mm: 5.0,
          utc_hour: 14,
          month: 1,
          freq_mhz: 10_000
        })

      # With random init, scores likely differ (not guaranteed but very probable)
      assert is_integer(score_hot)
      assert is_integer(score_cold)
    end
  end

  describe "feature_names/0" do
    test "returns 20 feature names" do
      assert length(Model.feature_names()) == 20
    end
  end

  describe "explain_prediction/3" do
    test "returns score and one contribution entry per feature" do
      params = Model.init()
      predict_fn = Model.compile_predict()

      result =
        Model.explain_prediction(predict_fn, params, %{
          surface_temp_c: 25.0,
          surface_dewpoint_c: 18.0,
          surface_pressure_mb: 1013.0,
          min_refractivity_gradient: -250.0,
          hpbl_m: 800.0,
          pwat_mm: 30.0,
          surface_refractivity: 340.0,
          latitude: 32.5,
          longitude: -97.0,
          sfi: 130.0,
          kp_max: 2.0,
          ducting_detected: true,
          k_index: 25.0,
          lifted_index: -2.0,
          utc_hour: 23,
          month: 7,
          freq_mhz: 10_000
        })

      assert is_integer(result.score)
      assert result.score >= 0 and result.score <= 100
      assert length(result.contributions) == 20

      Enum.each(result.contributions, fn c ->
        assert is_atom(c.feature)
        assert is_float(c.normalized) or is_integer(c.normalized)
        assert is_float(c.sensitivity)
        assert is_float(c.contribution)
      end)
    end

    test "contributions are sorted by absolute contribution descending" do
      params = Model.init()
      predict_fn = Model.compile_predict()

      result =
        Model.explain_prediction(predict_fn, params, %{
          surface_temp_c: 15.0,
          surface_dewpoint_c: 5.0,
          freq_mhz: 24_000
        })

      magnitudes = Enum.map(result.contributions, &abs(&1.contribution))
      assert magnitudes == Enum.sort(magnitudes, :desc)
    end

    test "every emitted feature name matches feature_names/0" do
      params = Model.init()
      predict_fn = Model.compile_predict()

      result = Model.explain_prediction(predict_fn, params, %{freq_mhz: 47_000})

      emitted = result.contributions |> Enum.map(& &1.feature) |> Enum.sort()
      expected = Enum.sort(Model.feature_names())
      assert emitted == expected
    end

    test "perturbing a feature changes the predicted score" do
      # Verify the sensitivity output is non-trivial: the model must
      # actually move when at least one feature moves. With random init
      # this is essentially guaranteed unless every weight collapses.
      params = Model.init()
      predict_fn = Model.compile_predict()

      result =
        Model.explain_prediction(predict_fn, params, %{
          surface_temp_c: 20.0,
          surface_dewpoint_c: 10.0,
          surface_pressure_mb: 1013.0,
          freq_mhz: 10_000
        })

      max_sensitivity = result.contributions |> Enum.map(&abs(&1.sensitivity)) |> Enum.max()
      assert max_sensitivity > 0.0
    end
  end
end