// Tests for ann.h - neural network init, forward/backward pass, activations
#include <stdio.h>
#define ANN_IMPLEMENTATION
#include "../include/ann.h"

#include <math.h>
#include <time.h>

#define EPS 1e-4

#define ASSERT(cond, msg)                                                        \
  do {                                                                           \
    if (!(cond)) {                                                               \
      fprintf(stderr, "FAIL: %s (test_ann.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_ann.c:%d)\n",       \
              msg, (double)(a), (double)(b), __LINE__);                          \
      failures++;                                                                \
    }                                                                            \
  } while (0)

static int failures = 0;

void test_init() {
  ann_t ann = {0};
  u64 layers[] = {2, 3, 1};
  ASSERT(ann_init(&ann, 3, layers), "ann_init succeeds");

  ASSERT(ann.layer_n == 3, "layer count");
  ASSERT(ann.buffer != NULL, "buffer allocated");
  ASSERT(ann.layer_neuron_n != NULL, "layer_neuron_n set");
  ASSERT(ann.layer_neuron_n[0] == 2, "input layer size");
  ASSERT(ann.layer_neuron_n[1] == 3, "hidden layer size");
  ASSERT(ann.layer_neuron_n[2] == 1, "output layer size");
  ASSERT(ann.neuron != NULL, "neuron pointer set");
  ASSERT(ann.weight != NULL, "weight pointer set");

  ann_free(&ann);
}

void test_free() {
  ann_t ann = {0};
  u64 layers[] = {2, 4, 2};
  ASSERT(ann_init(&ann, 3, layers), "ann_init succeeds");
  ann_free(&ann);
  ASSERT(ann.buffer == NULL, "buffer freed");

  ann_free(&ann);
}

void test_copy() {
  ann_t src = {0};
  u64 layers[] = {2, 3, 1};
  ASSERT(ann_init(&src, 3, layers), "ann_init succeeds");
  ann_random(&src);

  ann_t dst = {0};
  ann_copy(&dst, &src);

  ASSERT(dst.layer_n == src.layer_n, "copy layer_n");
  ASSERT(dst.neuron_n == src.neuron_n, "copy neuron_n");
  ASSERT(dst.weight_n == src.weight_n, "copy weight_n");
  ASSERT(dst.buffer != src.buffer, "copy has separate buffer");

  ann_free(&src);
  ann_free(&dst);
}

void test_random() {
  srand(42);

 ann_t ann = {0};
  u64 layers[] = {2, 2, 1};
  ASSERT(ann_init(&ann, 3, layers), "ann_init succeeds");
  ann_random(&ann);

  int has_nonzero = 0;
  for (u64 i = 0; i < ann.weight_n; i++) {
    if (ann.weight[i] != 0) {
      has_nonzero = 1;
      break;
    }
  }
  ASSERT(has_nonzero, "random sets non-zero weights");

  ann_free(&ann);
}

void test_forward() {
  ann_t ann = {0};
  u64 layers[] = {2, 2, 1};
  ASSERT(ann_init(&ann, 3, layers), "ann_init succeeds");

  f64 *w = ann.weight;
  for (u64 i = 0; i < ann.weight_n; i++) {
    w[i] = 0;
  }

  f64 input[2] = {1, 1};
  f64 output[1] = {0};
  ann_propagation_forward(&ann, input, output);

  ASSERT(output[0] >= 0 && output[0] <= 1, "forward output in sigmoid range");

  ann_free(&ann);
}

void test_backward() {
  ann_t ann = {0};
  u64 layers[] = {2, 2, 1};
  ASSERT(ann_init(&ann, 3, layers), "ann_init succeeds");
  ann_random(&ann);

  f64 input[2] = {0.5, 0.3};
  f64 output[1] = {0};
  f64 target[1] = {1};

  f64 weight_before = ann.weight[0];

  ann_propagation_forward(&ann, input, output);

  f64 output_before = output[0];

  ann_propagation_backward(&ann, input, output, target, 0.1);

  f64 weight_after = ann.weight[0];
  ASSERT(weight_before != weight_after, "backward changes weights");

  ann_propagation_forward(&ann, input, output);
  ASSERT(output[0] != output_before, "output changed after weight update");

  ann_free(&ann);
}

void test_error() {
  f64 err = ann_error_total((f64[]){0.9}, (f64[]){1.0}, 1);
  ASSERT(err > 0, "error is positive");

  f64 err_zero = ann_error_total((f64[]){1.0}, (f64[]){1.0}, 1);
  ASSERT_FLOAT_EQ(err_zero, 0, "error zero when output == target");

  f64 err1 = ann_error(0.5, 0.5);
  ASSERT_FLOAT_EQ(err1, 0, "single error zero");

  f64 err2 = ann_error(1.0, 0.0);
  ASSERT_FLOAT_EQ(err2, 0.5, "single error 0.5");
}

void test_activations() {
  ann_t ann = {0};
  u64 layers[] = {2, 2, 1};
  ASSERT(ann_init(&ann, 3, layers), "ann_init succeeds");

  ann_set_activation(&ann, SIGMOID, SIGMOID);
  f64 input[2] = {0.5, 0.5};
  f64 output[1] = {0};
  ann_propagation_forward(&ann, input, output);
  ASSERT(output[0] >= 0 && output[0] <= 1, "sigmoid output range");

  ann_set_activation(&ann, RELU, RELU);
  ann_propagation_forward(&ann, input, output);
  ASSERT(output[0] >= 0, "relu output non-negative");

  ann_set_activation(&ann, TANH, TANH);
  ann_propagation_forward(&ann, input, output);
  ASSERT(output[0] >= -1 && output[0] <= 1, "tanh output range");

  ann_set_activation(&ann, IDENTITY, IDENTITY);
  ann_propagation_forward(&ann, input, output);

  ann_set_activation(&ann, BINARY, BINARY);
  ann_propagation_forward(&ann, input, output);
  ASSERT(output[0] == 0 || output[0] == 1, "binary output is 0 or 1");

  ann_set_activation(&ann, ELU, ELU);
  ann_propagation_forward(&ann, input, output);
  ASSERT(isfinite(output[0]), "elu output is finite");

  ann_set_activation(&ann, LRELU, LRELU);
  ann_propagation_forward(&ann, input, output);
  ASSERT(isfinite(output[0]), "lrelu output is finite");

  ann_free(&ann);
}

void test_train_numeric() {
  srand(42);

  ann_t ann = {0};
  u64 layers[] = {2, 4, 1};
  ASSERT(ann_init(&ann, 3, layers), "ann_init succeeds");
  ann_random(&ann);

  f64 input[2] = {0, 1};
  f64 target[1] = {1};
  f64 output_before[1] = {0};
  ann_propagation_forward(&ann, input, output_before);

  ann_train_numeric(&ann, input, target, 1, 0.01);

  f64 output_after[1] = {0};
  ann_propagation_forward(&ann, input, output_after);
  ASSERT(output_before[0] != output_after[0], "numeric training changes output");
  ASSERT(isfinite(output_after[0]), "output is finite after training");

  ann_free(&ann);
}

void test_xor_training() {
  srand(42);

  ann_t ann = {0};
  u64 layers[] = {2, 4, 1};
  ASSERT(ann_init(&ann, 3, layers), "ann_init succeeds");
  ann_set_activation(&ann, SIGMOID, SIGMOID);
  ann_random(&ann);

  f64 inputs[4][2] = {{0, 0}, {0, 1}, {1, 0}, {1, 1}};
  f64 targets[4] = {0, 1, 1, 0};
  f64 output[1] = {0};

  for (int iter = 0; iter < 5000; iter++) {
    for (int i = 0; i < 4; i++) {
      ann_propagation_forward(&ann, inputs[i], output);
      ann_propagation_backward(&ann, inputs[i], output, &targets[i], 0.5);
    }
  }

  f64 outputs[4];
  for (int i = 0; i < 4; i++) {
    ann_propagation_forward(&ann, inputs[i], &outputs[i]);
  }

  ASSERT(outputs[0] < 0.5, "XOR(0,0) < 0.5");
  ASSERT(outputs[1] > 0.5, "XOR(0,1) > 0.5");
  ASSERT(outputs[2] > 0.5, "XOR(1,0) > 0.5");
  ASSERT(outputs[3] < 0.5, "XOR(1,1) < 0.5");

  ann_free(&ann);
}

void test_multi_layer() {
  ann_t ann = {0};
  u64 layers[] = {3, 4, 5, 2};
  ASSERT(ann_init(&ann, 4, layers), "ann_init succeeds");
  ann_random(&ann);

  ASSERT(ann.layer_neuron_n[0] == 3, "input 3");
  ASSERT(ann.layer_neuron_n[1] == 4, "hidden1 4");
  ASSERT(ann.layer_neuron_n[2] == 5, "hidden2 5");
  ASSERT(ann.layer_neuron_n[3] == 2, "output 2");

  f64 input[3] = {0.1, 0.2, 0.3};
  f64 output[2] = {0, 0};
  ann_propagation_forward(&ann, input, output);

  ASSERT(isfinite(output[0]), "multi-layer output[0] is finite");
  ASSERT(isfinite(output[1]), "multi-layer output[1] is finite");

  ann_free(&ann);
}

int main() {
  test_init();
  test_free();
  test_copy();
  test_random();
  test_forward();
  test_backward();
  test_error();
  test_activations();
  test_train_numeric();
  test_xor_training();
  test_multi_layer();

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