#include "activation_function.hpp"
#include "utility.hpp"
#include <cmath>
#include <gtest/gtest.h>

// Simple identity activation function for testing
struct Identity {
  void operator()(std::vector<float> &x) const {
    // Identity function - no change to values
    // Below statement is needed to remove compiler warning about unused var
    // warning
    (void)x;
  }
};

TEST(UtilityTest, FeedLayerIdentityTest) {
  // Test with identity activation function for simple verification
  // Input: [1, 2]
  // Weights: [0.5, -0.5, 1.0, -1.0]

  std::vector<float> weights = {0.5, -0.5, 1.0, -1.0};
  std::vector<float> input = {1.0, 2.0};
  Identity identity;

  auto output = Utilities::feed_layer<Identity>(weights.begin(), weights.end(),
                                                input, identity);

  ASSERT_EQ(output.size(), 4);
  EXPECT_NEAR(output[0], 1.5f, 1e-5);  // 1.0 * 0.5 + 2.0 * 0.5
  EXPECT_NEAR(output[1], -1.5f, 1e-5); // 1.0 * -0.5 + 2.0 * -0.5
  EXPECT_NEAR(output[2], 3.0f, 1e-5);  // 1.0 * 1.0 + 2.0 * 1.0
  EXPECT_NEAR(output[3], -3.0f, 1e-5); // 1.0 * -1.0 + 2.0 * -1.0
}

TEST(UtilityTest, FeedLayerSigmoidTest) {
  // Test with sigmoid activation
  // Input: [1]
  // Weights: [2, -2]
  std::vector<float> weights = {2.0, -2.0};
  std::vector<float> input = {1.0};
  Sigmoid sigmoid;

  auto output = Utilities::feed_layer<Sigmoid>(weights.begin(), weights.end(),
                                               input, sigmoid);

  ASSERT_EQ(output.size(), 2);
  // Note: Sigmoid is applied to the whole vector after matrix multiplication
  float expected0 = 2.0;  // 1.0 * 2.0
  float expected1 = -2.0; // 1.0 * -2.0
  EXPECT_NEAR(output[0], 1.0 / (1.0 + std::exp(-expected0)), 1e-5);
  EXPECT_NEAR(output[1], 1.0 / (1.0 + std::exp(-expected1)), 1e-5);
}

TEST(UtilityTest, FeedLayerSoftMaxTest) {
  // Test with softmax activation
  // Input: [1]
  // Weights: [2, 2]
  std::vector<float> weights = {2.0, 2.0};
  std::vector<float> input = {1.0};
  SoftMax softmax;

  auto output = Utilities::feed_layer<SoftMax>(weights.begin(), weights.end(),
                                               input, softmax);

  ASSERT_EQ(output.size(), 2);
  // Both outputs should be 0.5 since inputs to softmax are equal (both 2.0)
  EXPECT_NEAR(output[0], 0.5, 1e-5);
  EXPECT_NEAR(output[1], 0.5, 1e-5);
}

TEST(UtilityTest, FeedLayerEmptyInput) {
  std::vector<float> weights = {1.0, 1.0};
  std::vector<float> input = {};
  Identity identity;

  auto output = Utilities::feed_layer<Identity>(weights.begin(), weights.end(),
                                                input, identity);

  ASSERT_EQ(output.size(), 2);
  EXPECT_NEAR(output[0], 0.0f, 1e-5);
  EXPECT_NEAR(output[1], 0.0f, 1e-5);
}