// genetic_pushdown_btc.c - Example training of a pushdown automaton by genetic algorithm

#include "../include/escape.h"
#include <float.h>
#include <math.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <tgmath.h>
#include <time.h>

typedef double stack_value_t;
#define STACK_IMPLEMENTATION
#include "../include/stack.h"

typedef stack_value_t pushdown_value_t;
#define PUSHDOWN_IMPLEMENTATION
#include "../include/pushdown.h"

typedef struct {
  size_t size;
  pushdown_instruction_t *op;
} program_t;

typedef program_t ga_individual_t;
#define GENETIC_IMPLEMENTATION
#include "../include/genetic.h"

#define POPULATION_N 128
#define RATE_SELECTION 0.25
#define RATE_MUTATION 0.2
#define ERROR 1000

#define generate() ((pushdown_instruction_t)(rand() % END))

void crossover(program_t *y, program_t x0, program_t x1) {
  if (!y || x0.size == 0 || x1.size == 0) {
    return;
  }

  uint32_t x0_crossover = rand() % x0.size;
  uint32_t x1_crossover = rand() % x1.size;

  uint32_t size = x0_crossover + (x1.size - x1_crossover);

  pushdown_instruction_t *instructions =
      calloc(size + 1, sizeof(pushdown_instruction_t));
  if (!instructions) {
    free(y->op);
    y->op = NULL;
    y->size = 0;
    return;
  }
  memcpy(instructions, x0.op, x0_crossover * sizeof(pushdown_instruction_t));
  memcpy(instructions + x0_crossover, x1.op + x1_crossover,
         (x1.size - x1_crossover) *  sizeof(pushdown_instruction_t));
  instructions[size] = END;

  free(y->op);
  y->op = instructions;
  y->size = size;
}

void mutation(program_t *x) {
  if (!x || x->size == 0) {
    return;
  }
  uint32_t i = rand() % x->size;
  x->op[i] = generate();
}

double BTC[] = {
#include "./genetic_pushdown_btc.txt"
};

double fitness(program_t individual, void const *data) {
  (void)data;
  if (!individual.op)
    return DBL_MAX;

  double error = 0;
  const uint32_t K = 16;
  size_t btc_n = sizeof(BTC) / sizeof(*BTC);
  if (btc_n < K)
    return 0;

  for (size_t i = 0; i < btc_n - K; i++) {
    stack_t input = {0};
    stack_t output = {0};

    stack_from_array(&input, &BTC[i], K);
    stack_init(&output);
    pushdown_run(&output, &input, individual.op);
    double e = fabs(stack_peek(&output) - BTC[i + K]) + individual.size;
    error += e;

    stack_free(&input);
    stack_free(&output);
  }

  return error;
}

// void kill(int signal) { (void)signal; } // removed: shadows POSIX kill(2)

int main(void) {
  program_t population[POPULATION_N] = {0};

  srand(time(0));

  // Initialize population
  for (uint64_t i = 0; i < POPULATION_N; i++) {
    uint64_t n = rand() % 10 + 10;
    pushdown_instruction_t *program =
        calloc(n + 1, sizeof(pushdown_instruction_t));
    if (!program) {
      fprintf(stderr, "Failed to allocate program\n");
      for (uint64_t k = 0; k < i; k++) free(population[k].op);
      return 1;
    }
    program[n] = END;

    for (uint64_t j = 0; j < n; j++) {
      program[j] = generate();
    }

    population[i].op = program;
    population[i].size = n;
  }

  uint32_t n = 0;

  double error = ERROR + 1, error_last = 0;
  double rate_mutation = RATE_MUTATION;
  double rate_selection = RATE_SELECTION;

  // Conduct optimization
  uint32_t max_iterations = 1000000;
  while (error > ERROR && n < max_iterations) {
   error_last = error;
    error = ga_generation(population, POPULATION_N, rate_selection,
                           rate_mutation, fitness, NULL, crossover, mutation);

    if (error == error_last) {
      rate_mutation = fmin(1.1 * rate_mutation, 0.95);
    } else {
      rate_mutation = RATE_MUTATION;
    }

    printf("%u\t%zu\t%lf\t%lf\n", n, population[0].size, error, error / POPULATION_N);
    n++;
  }

  clock_t t1 = clock();
  (void)t1;

  pushdown_instruction_t *op = population[0].op;
  if (!op) {
    fprintf(stderr, "Best individual has no program\n");
    return 1;
  }
  while (*op != END) {
    printf("%d\n", *op);
    op++;
  }

  for (uint64_t i = 0; i < POPULATION_N; i++) {
    free(population[i].op);
  }

  return 0;
}