// Tests for genetic.h - genetic algorithm generation step
typedef char *ga_individual_t;
#define GENETIC_IMPLEMENTATION
#include "../include/genetic.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <float.h>

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

static int failures = 0;

#define TARGET "HELLO"
#define POP_N 20

static void crossover_fn(ga_individual_t *child, ga_individual_t p1, ga_individual_t p2) {
  if (!child || !p1 || !p2)
    return;
  uint32_t l1 = strlen(p1);
  uint32_t l2 = strlen(p2);
  if (l1 == 0 || l2 == 0)
    return;
  uint32_t c1 = rand() % l1;
  uint32_t c2 = rand() % l2;
  uint32_t len = c1 + (l2 - c2);
  char *s = malloc(len + 1);
  if (!s) {
    free(*child);
    *child = NULL;
    return;
  }
  memcpy(s, p1, c1);
  memcpy(s + c1, p2 + c2, l2 - c2);
  s[len] = '\0';
  free(*child);
  *child = s;
}

static void mutation_fn(ga_individual_t *ind) {
  if (!ind || !*ind)
    return;
  uint32_t l = strlen(*ind);
  if (l == 0)
    return;
  (*ind)[rand() % l] = (char)(rand() % 26 + 'A');
}

static double fitness_fn(ga_individual_t ind, void const *data) {
  if (!ind)
    return DBL_MAX;
  char const *target = (char const *)data;
  double error = 0;
  while (*ind && *target) {
    error += abs((int)(*target++ - *ind++));
  }
  while (*ind) error += abs((int)*ind++);
  while (*target) error += abs((int)*target++);
  return error;
}

static void free_population(char *population[], int n) {
  for (int i = 0; i < n; i++)
    free(population[i]);
}

void test_generation() {
  srand(42);
  char *population[POP_N];

  for (int i = 0; i < POP_N; i++) {
    uint32_t len = 4 + rand() % 3;
    char *s = malloc(len + 1);
    if (!s) { free_population(population, i); exit(1); }
    for (uint32_t j = 0; j < len; j++) {
      s[j] = (char)(rand() % 26 + 'A');
    }
    s[len] = '\0';
    population[i] = s;
  }

  double error = ga_generation(population, POP_N, 0.5, 1.0, fitness_fn, TARGET,
                                crossover_fn, mutation_fn);

  ASSERT(error >= 0, "error non-negative");

  for (int i = POP_N - 1; i >= 0; i--) {
    free(population[i]);
  }
}

void test_convergence() {
  srand(42);
  char *population[POP_N];

  for (int i = 0; i < POP_N; i++) {
    uint32_t len = strlen(TARGET);
    char *s = malloc(len + 1);
    if (!s) { free_population(population, i); exit(1); }
    for (uint32_t j = 0; j < len; j++) {
      s[j] = (char)(rand() % 26 + 'A');
    }
    s[len] = '\0';
    population[i] = s;
  }

  double error = 1;
  int generations = 0;
  while (error > 0 && generations < 500) {
    error = ga_generation(population, POP_N, 0.5, 1.0, fitness_fn, TARGET,
                          crossover_fn, mutation_fn);
    generations++;
  }

  ASSERT(generations < 500, "converges within 500 generations");

  for (int i = 0; i < POP_N; i++) {
    free(population[i]);
  }
}

void test_selection_rate() {
  srand(42);
  char *population[POP_N];

  for (int i = 0; i < POP_N; i++) {
    char *s = malloc(6);
    if (!s) { free_population(population, i); exit(1); }
    strcpy(s, "HELLO");
    population[i] = s;
  }

  double error = ga_generation(population, POP_N, 0.5, 0.1, fitness_fn, TARGET,
                                crossover_fn, mutation_fn);
  ASSERT(error >= 0, "low mutation rate returns non-negative error");
  ASSERT(error == 0, "all-population-matching-target has zero error");

  for (int i = 0; i < POP_N; i++) {
    free(population[i]);
  }
}

int main() {
  test_generation();
  test_convergence();
  test_selection_rate();

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