// genetic_ann.c - Example training of an XOR neural network by genetic algorithm

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

#define ANN_IMPLEMENTATION
#include "../include/ann.h"

#define POPULATION_N 100
#define RATE_SELECTION 0.2
#define RATE_MUTATION 0.2
#define MAGNITUDE_MUTATION 0.01
#define ERROR 0.00001

typedef ann_t *ga_individual_t;

#define GENETIC_IMPLEMENTATION
#include "../include/genetic.h"

double fitness(ann_t *const x, void const *c) {
  f64 *i = *((f64 **)c);
  f64 *t = *((f64 **)c + 1);
  f64 o[4];

  f64 e = 0, sum = 0;
  for (u64 b = 0; b < 4; b++) {
    ann_propagation_forward(x, &i[2 * b], &o[b]);
    e = ann_error_total(&o[b], &t[b], x->layer_neuron_n[x->layer_n - 1]);
    sum += e;
  }

  return sum;
}

void crossover(ann_t **y, ann_t *x0, ann_t *x1) {
  if (!y || !x0 || !x1 || !*y)
    return;
  if (x0->weight_n == 0 || x1->weight_n == 0)
    return;
  ann_t *child = (ann_t *)*y;
  if (child->weight_n == 0)
    return;
  uint32_t i = rand() % x0->weight_n;
  u64 max_copy = child->weight_n < x0->weight_n ? child->weight_n : x0->weight_n;
  max_copy = max_copy < x1->weight_n ? max_copy : x1->weight_n;
  u64 safe_i = i >= max_copy ? max_copy - 1 : (u64)i;

  memcpy(child->weight, x0->weight, safe_i * sizeof(f64));
  memcpy(child->weight + safe_i, x1->weight + safe_i, (max_copy - safe_i) * sizeof(f64));
  if (child->weight_n > max_copy) {
    memset(child->weight + max_copy, 0, (child->weight_n - max_copy) * sizeof(f64));
  }
}

void mutate(ann_t **x) {
  if (!x || !*x || (*x)->weight_n == 0)
    return;
  uint32_t i_mutation = rand() % (*x)->weight_n;
  f64 r =
      (2 * (MAGNITUDE_MUTATION * (f64)rand() / RAND_MAX)) - MAGNITUDE_MUTATION;
  (*x)->weight[i_mutation] += r;
}

int main(void) {
  srand(time(0));

  f64 input[12] = {0, 0, 0, 1, 1, 0, 1, 1};
  f64 output;
  f64 target[4] = {0, 1, 1, 0};

  ann_t *networks[POPULATION_N] = {0};

  for (u64 i = 0; i < POPULATION_N; i++) {
    ann_t *ann = malloc(sizeof(ann_t));
    if (!ann) {
      fprintf(stderr, "Failed to allocate ann_t\n");
      for (u64 k = 0; k < i; k++) { ann_free(networks[k]); free(networks[k]); }
      exit(1);
    }
    if (!ann_init(ann, 3, (u64[]){2, 2, 1})) {
      free(ann);
      fprintf(stderr, "Failed to init ann_t\n");
      for (u64 k = 0; k < i; k++) { ann_free(networks[k]); free(networks[k]); }
      exit(1);
    }
    ann_set_activation(ann, SIGMOID, IDENTITY);
    ann_random(ann);
    networks[i] = ann;
  }

 f64 e = ERROR + 1;
  uint32_t n = 0;
  uint32_t max_iterations = 1000000;
  uint32_t print_counter = 0;
  f64 *fitness_data[2] = {input, target};
  while (e > ERROR && n < max_iterations) {
   e = ga_generation(networks, POPULATION_N, RATE_SELECTION, RATE_MUTATION,
                         fitness, fitness_data, crossover, mutate);

    n++;
    print_counter++;

    if (print_counter > 1000) {
      printf("\rERROR: %10.5lf", e);
      print_counter = 0;
    }
  }

  printf("\33[2K\r");
  printf("ERROR: %lf", e);

  for (int i = 0; i < 1; i++) {
    for (int j = 0; j < 4; j++) {
      ann_propagation_forward(networks[i], (f64 const *const)&input[2 * j],
                              &output);
      ann_print_neuron(networks[i], &input[2 * j], &output);
    }
  }

  printf("\nITERATIONS: %u\n", n);

  for (u64 i = 0; i < POPULATION_N; i++) {
    ann_free(networks[i]);
    free(networks[i]);
  }
}