// genetic.h - A simple genetic algorithm implementation
//
// Thread safety: Not thread-safe. External synchronization required.
// Behavior: rate_selection must be in (0, 1). A value of 1.0 is rejected
//   because it results in no evolution (all individuals survive).
//   rate_mutation must be >= 0. Returns -1.0 on any parameter error.

#ifndef GENETIC_H
#define GENETIC_H

#include "type.h"

// ga_individual_t must be typedef'd by the user before including this header.

// ga_crossover_t
//
// Type definition for a crossover function
//
// ga_individual_t * - The address for the child individual
// ga_individual_t - Parent individual 0
// ga_individual_t - Parent individual 1

typedef void (*ga_crossover_t)(ga_individual_t *, ga_individual_t,
                                ga_individual_t);

// ga_mutation_t
//
// Type definition for a mutation function
//
// ga_individual_t * - The individual being mutated

typedef void (*ga_mutation_t)(ga_individual_t *);

// ga_fitness_t()
//
// Type definition for a fitness function
//
// uintptr_t - The individual on which to run the fitness function
// void const * - The user defined data pointer

typedef f64 (*ga_fitness_t)(ga_individual_t, void const *);

f64 ga_generation(ga_individual_t *, u32, f64, f64, ga_fitness_t,
                    void const *, ga_crossover_t, ga_mutation_t);

#endif

#ifdef GENETIC_IMPLEMENTATION // IMPLEMENTATION

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <tgmath.h>

// ga_fitness_sort()
//
// Run the fitness function on
//
// population - An array of pointers to each individual
// population_n - The size of the population array
// fitness - The function used to evaluate the fitness of an individual
// fitness_data - Additional data required for the fitness function
//
// return - The error for the best performing individual

static f64 ga_fitness_sort(ga_individual_t *population,
                                 u32 population_n, ga_fitness_t fitness,
                                 void const *fitness_data) {
  if (!population || !fitness || population_n == 0)
    return -1;

  f64 *fitness_buffer = calloc((size_t)population_n, sizeof(f64));
  if (!fitness_buffer)
    return -1;

  for (u32 i = 0; i < population_n; i++) {
    fitness_buffer[i] = fabs(fitness(population[i], fitness_data));
  }

  // Index array for indirect sort
  u32 *indices = malloc((size_t)population_n * sizeof(u32));
  if (!indices) {
    free(fitness_buffer);
    return -1;
  }
  for (u32 i = 0; i < population_n; i++)
    indices[i] = i;

  // Shell sort on indices by fitness value (indirect sort, only rearrange indices)
  {
    u32 gap = population_n / 2;
    while (gap > 0) {
      for (u32 i = gap; i < population_n; i++) {
        u32 temp = indices[i];
        f64 temp_f = fitness_buffer[temp];
        u32 j;
        for (j = i; j >= gap && fitness_buffer[indices[j - gap]] > temp_f; j -= gap) {
          indices[j] = indices[j - gap];
        }
        indices[j] = temp;
      }
      gap /= 2;
    }
  }

  // Copy sorted individuals back
  ga_individual_t *sorted = calloc((size_t)population_n, sizeof(ga_individual_t));
  if (!sorted) {
    free(fitness_buffer);
    free(indices);
    return -1;
  }
  for (u32 i = 0; i < population_n; i++)
    sorted[i] = population[indices[i]];
  memcpy(population, sorted, (size_t)population_n * sizeof(ga_individual_t));

  f64 error = fitness_buffer[indices[0]];
  free(sorted);
  free(fitness_buffer);
  free(indices);

  return error;
}

// ga_generation()
//
// Run a single generation of the genetic algorithm
//
// population - An array of pointers to each individual
// population_n - The size of the population array
// rate_selection - A number [0,1] denoting the percentage of the population
//                  which survives to the next generation
// rate_mutation - A number [0,1] denoting the percentage of new individuals
//                 which will recieve mutations
// fitness - The function used to evaluate the fitness of an individual
// fitness_data - Additional data required for the fitness function
// crossover - The function used to perform repopulation via the
//              combination of two parents
// mutation - The function used to perform mutations on a given individual
//
// return - The error of the best performing individual

f64 ga_generation(ga_individual_t *population, u32 population_n,
                       f64 rate_selection, f64 rate_mutation,
                       ga_fitness_t fitness, void const *fitness_data,
                       ga_crossover_t crossover, ga_mutation_t mutation) {
  if (!population || !fitness || !crossover || !mutation)
    return -1;
  if (rate_selection <= 0 || rate_selection >= 1.0)
    return -1;
  if (rate_mutation < 0)
    return -1;
  if (population_n < 3)
    return -1;

  u32 selection_n = (u32)(rate_selection * (f64)population_n);
  if (selection_n >= population_n)
    return -1;

  // Crossover
  for (u32 i = selection_n; i < population_n; i++) {
    u32 parent_i = (u32)((u64)rand() * (u64)selection_n / ((u64)RAND_MAX + 1));
    u32 parent_j = (u32)((u64)rand() * (u64)selection_n / ((u64)RAND_MAX + 1));

    crossover(&population[i], population[parent_i], population[parent_j]);
  }

  // Mutation count
  u32 mutation_n = (u32)(rate_mutation * (f64)(population_n - selection_n));
  for (u32 i = 0; i < mutation_n; i++) {
    u32 range = population_n - selection_n;
    u32 mutation_i = selection_n + (u32)((u64)rand() * (u64)range / ((u64)RAND_MAX + 1));
    mutation(&population[mutation_i]);
  }

  // Fitness
  return ga_fitness_sort(population, population_n, fitness, fitness_data);
}

#endif