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, "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, 13})
      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, 13})
      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, 13 features
      input = Nx.broadcast(0.5, {1, 13})
      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, 13})
      output = Model.predict(params, input)

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

  describe "encode_features/1" do
    test "returns 13-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) == 13
      assert Enum.all?(features, &is_float/1)
    end

    test "uses defaults for missing values" do
      features = Model.encode_features(%{})
      assert length(features) == 13
      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), 0.001
    end
  end

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