#include <float.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include "genalg.h"


genalg_t * ga_create(
  size_t pop_count, 
  double (*fitness)(void*),
  void * (*create_child)(void*),
  void (*org_free)(void*)
) {
  genalg_t *ga = malloc(sizeof(genalg_t));
  if (ga == NULL) {
    fprintf(stderr, "failed to allocate genalg_t structure\n");
    return NULL;
  }
  ga->population = malloc(pop_count * sizeof(void*));
  if (ga->population == NULL) {
    fprintf(stderr, "failed to allocate population buffer\n");
    free(ga);
    return NULL;
  }
  ga->pop_count = pop_count;
  ga->fitness = fitness;
  ga->create_child = create_child;
  ga->org_free = org_free;
  return ga;
}


void ga_free(genalg_t *ga) {
  for (size_t i=0; i<ga->pop_count; i++) {
    ga->org_free(ga->population[i]);
  }
  free(ga->population);
  free(ga);
}


void * ga_tournament_select(genalg_t *ga, int tournament_sz) {
  double best_fitness = -DBL_MAX;
  size_t best_index = 0;

  for (int i=0; i<tournament_sz; i++) {
    size_t idx = rand() % ga->pop_count;
    double fitness = ga->fitness(ga->population[idx]);
    if (fitness > best_fitness) {
      best_fitness = fitness;
      best_index = idx;
    }
  }

  return ga->population[best_index];
}


struct genalg_stats_t ga_population_statistics(genalg_t *ga) {
  struct genalg_stats_t stats = { 0, 0, DBL_MAX, -DBL_MAX, NULL };
  // compute mean
  for (size_t i=0; i<ga->pop_count; i++) {
    stats.mean += ga->fitness(ga->population[i]);
  }
  stats.mean /= ga->pop_count;
  // compute other stats
  for (size_t i=0; i<ga->pop_count; i++) {
    double fitness = ga->fitness(ga->population[i]);
    double diff = fitness - stats.mean;
    stats.variance += diff*diff;
    stats.min = fmin(stats.min, fitness);
    stats.max = fmax(stats.max, fitness);
    if (stats.max == fitness) {
      stats.best = ga->population[i];
    }
  }

  return stats;
}


// helper: qsort from fittest to least fit
double (*fitness_sort_cmp)(void*) = NULL;
int fitness_sort(const void *a, const void *b) {
  if (fitness_sort_cmp == NULL) {
    return 0;
  }

  double fa = fitness_sort_cmp((void*)a);
  double fb = fitness_sort_cmp((void*)b);
  if (fa < fb) { return 1; }
  if (fa > fb) { return -1; }
  return 0;
}


int ga_replace_population(genalg_t *ga, int tournament_sz, int keep) {
  // create new population buffer
  void **new_population = malloc(ga->pop_count * sizeof(void*));
  if (new_population == NULL) {
    fprintf(stderr, "failed to allocate new population buffer\n");
    return 1;
  }

  // copy [keep] most fit from previous generation
  if (keep) {
    fitness_sort_cmp = ga->fitness;
    qsort(ga->population, ga->pop_count, sizeof(void*), fitness_sort);
    memcpy(new_population, ga->population, keep * sizeof(void*));
  }

  // fill new_population with children
  for (size_t i=keep; i<ga->pop_count; i++) {
    void * parent = ga_tournament_select(ga, tournament_sz);
    new_population[i] = ga->create_child(parent);
    if (new_population[i] == NULL) {
      for (size_t j=0; j<i; j++) {
        ga->org_free(new_population[j]);
      }
      free(new_population);
      return 1;
    }
  }

  // delete un-kept population
  for (size_t i=keep; i<ga->pop_count; i++) {
    ga->org_free(ga->population[i]);
  }
  free(ga->population);
  ga->population = new_population;

  return 0;
}