// Tests for ufo.h - Universal Function Originator evaluation and mutation
#include <stdio.h>
#define UFO_IMPLEMENTATION
#include "../include/ufo.h"

#include <math.h>

#define EPS 1e-6

#define ASSERT(cond, msg)                                                        \
  do {                                                                           \
    if (!(cond)) {                                                               \
      fprintf(stderr, "FAIL: %s (test_ufo.c:%d)\n", msg, __LINE__);             \
      failures++;                                                                \
    }                                                                            \
  } while (0)

#define ASSERT_FLOAT_EQ(a, b, msg)                                               \
  do {                                                                           \
    if (fabs((a) - (b)) > EPS) {                                                 \
      fprintf(stderr, "FAIL: %s (got %f, expected %f) (test_ufo.c:%d)\n",       \
              msg, (double)(a), (double)(b), __LINE__);                          \
      failures++;                                                                \
    }                                                                            \
  } while (0)

static int failures = 0;

void test_init() {
  ufo_t *ufo = ufo_init(2, 3);
  ASSERT(ufo != NULL, "init returns non-NULL");
  ASSERT(ufo->n == 2, "input count");
  ASSERT(ufo->v == 3, "block count");
  ASSERT(ufo->a != NULL, "weights allocated");
  ASSERT(ufo->b != NULL, "exponents allocated");
  ASSERT(ufo->w != NULL, "external weights allocated");
  ASSERT(ufo->c != NULL, "external exponents allocated");
  ASSERT(ufo->f != NULL, "unary functions allocated");
  ASSERT(ufo->fe != NULL, "external binary functions allocated");
  ASSERT(ufo->fi != NULL, "internal binary functions allocated");
  ufo_free(ufo);
}

void test_init_defaults() {
  ufo_t *ufo = ufo_init(1, 2);
  ASSERT(ufo != NULL, "init returns non-NULL for init_defaults");
  if (!ufo)
    return;

  ASSERT_FLOAT_EQ(ufo->w[2], 0, "bias is zero");

  for (uint32_t j = 0; j < ufo->v; j++) {
    ASSERT_FLOAT_EQ(ufo_ajk(ufo, j, ufo->n), 0, "bias weight is zero");
  }

  ufo_free(ufo);
}

void test_evaluate_default() {
  ufo_t *ufo = ufo_init(2, 3);
  f64 x[] = {1.0, 2.0};
  f64 result = ufo_evaluate(ufo, x);
  ASSERT(isfinite(result), "evaluate returns finite value");
  ASSERT(result > 0, "evaluate returns positive value with default weights");
  ufo_free(ufo);
}

void test_evaluate_multiple_inputs() {
  ufo_t *ufo = ufo_init(3, 2);
  f64 x[] = {1.0, 2.0, 3.0};
  f64 result = ufo_evaluate(ufo, x);
  ASSERT(isfinite(result), "multi-input evaluate returns finite");
  ASSERT(result > 0, "multi-input evaluate returns positive");
  ufo_free(ufo);
}

void test_free() {
  ufo_t *ufo = ufo_init(2, 5);
  ASSERT(ufo != NULL, "init returns non-NULL before free");
  ufo_free(ufo);
}

void test_different_configs() {
  ufo_t *u2 = ufo_init(5, 10);
  ufo_t *u3 = ufo_init(10, 2);
  ASSERT(u2 != NULL, "init config 2");
  ASSERT(u3 != NULL, "init config 3");

  f64 x2[5] = {1, 2, 3, 4, 5};
  f64 x3[10] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};

  f64 r2 = ufo_evaluate(u2, x2);
  f64 r3 = ufo_evaluate(u3, x3);

  ASSERT(isfinite(r2), "config 2 returns finite");
  ASSERT(isfinite(r3), "config 3 returns finite");
  ASSERT(r2 != r3, "different configs produce different results");

  ufo_free(u2);
  ufo_free(u3);
}

void test_weight_modification() {
  ufo_t *ufo = ufo_init(2, 3);
  ASSERT(ufo != NULL, "ufo_init returns non-NULL");
  if (!ufo) return;

  f64 x[] = {3.0, 4.0};
  f64 result_before = ufo_evaluate(ufo, x);

  ufo_ajk(ufo, 0, 0) = 0.0;
  ufo->w[0] = 0.0;
  ufo->w[1] = 0.0;
  ufo->w[2] = 0.0;

  f64 result_after = ufo_evaluate(ufo, x);
  ASSERT(isfinite(result_after), "modified weights evaluate finite");
  ASSERT(result_before != result_after, "modified weights change result");

  ufo_free(ufo);
}

void test_large_ufo() {
  ufo_t *ufo = ufo_init(4, 50);
  f64 x[4] = {0.1, 0.2, 0.3, 0.4};
  f64 result = ufo_evaluate(ufo, x);
  ASSERT(isfinite(result), "large UFO returns finite");
  ufo_free(ufo);
}

int main() {
  test_init();
  test_init_defaults();
  test_evaluate_default();
  test_evaluate_multiple_inputs();
  test_free();
  test_different_configs();
  test_weight_modification();
  test_large_ufo();

  printf("test_ufo: %d failures\n", failures);
  return failures;
}