#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>
#include <sys/select.h>
#include <termios.h>
#include <ctype.h>

// adjustable parameters
#define POPULATION_SIZE     200        // total population size
#define NUM_SURVIVORS       4          // number of survivors per generation (must be less than POPULATION_SIZE)
#define GENOME_LEN          28         // number of loci in each genome (each locus is one character)
#define MUTATION_RATE       0.01       // probability that a locus is mutated per generation
#define FITNESS_THRESHOLD   28         // program stops when this fitness threshold is met or exceeded
#define DISPLAY_INTERVAL    1          // display results every n generations, where n == DISPLAY_INTERVAL
#define PAUSE_TIME          50000      // pause time in microseconds
#define GENOMES_TO_DISPLAY  1          // in STEP_MODE, number of (the fittest) genomes to display at each step
#define STEP_MODE           1          // if set, program will pause every n generations and display the genomes
#define ENABLE_LATCHING     0          // if set, enable latching
#define ENABLE_SELECTION    1          // if set, enable selection; otherwise, select survivors randomly

#define CHARSET_SIZE        27         // 26 uppercase letters plus the space character
#define NUM_LINES           300        // number of lines in the program's display "window"
#define CHARS_PER_LINE      120        // for each line in the "window"

// A genome consists of an array of loci plus the associated fitness value for each locus
typedef struct {
  char   locus[GENOME_LEN];
  int fitness;
} genome_t;

genome_t target;                        // target phrase
genome_t genome_array[POPULATION_SIZE]; // genomes of all members of the population

char weasel_to_ascii(char c) {
  switch(c) {
    case  0: return ' ';
    case  1: return 'A';
    case  2: return 'B';
    case  3: return 'C';
    case  4: return 'D';
    case  5: return 'E';
    case  6: return 'F';
    case  7: return 'G';
    case  8: return 'H';
    case  9: return 'I';
    case 10: return 'J';
    case 11: return 'K';
    case 12: return 'L';
    case 13: return 'M';
    case 14: return 'N';
    case 15: return 'O';
    case 16: return 'P';
    case 17: return 'Q';
    case 18: return 'R';
    case 19: return 'S';
    case 20: return 'T';
    case 21: return 'U';
    case 22: return 'V';
    case 23: return 'W';
    case 24: return 'X';
    case 25: return 'Y';
    case 26: return 'Z';
  }
}

char ascii_to_weasel(char c) {
  switch(c) {
    case ' ': return  0;
    case 'A': return  1;
    case 'B': return  2;
    case 'C': return  3;
    case 'D': return  4;
    case 'E': return  5;
    case 'F': return  6;
    case 'G': return  7;
    case 'H': return  8;
    case 'I': return  9;
    case 'J': return 10;
    case 'K': return 11;
    case 'L': return 12;
    case 'M': return 13;
    case 'N': return 14;
    case 'O': return 15;
    case 'P': return 16;
    case 'Q': return 17;
    case 'R': return 18;
    case 'S': return 19;
    case 'T': return 20;
    case 'U': return 21;
    case 'V': return 22;
    case 'W': return 23;
    case 'X': return 24;
    case 'Y': return 25;
    case 'Z': return 26;
  }
}

// Swap two genomes given their indices.
void swap_genomes(int g1, int g2) {
  genome_t temp_genome;
  if (g1 == g2) {
    return;
  } else {
    temp_genome = genome_array[g1];
    genome_array[g1] = genome_array[g2];
    genome_array[g2] = temp_genome;
  }
}

// Thoroughly and randomly scramble the population within the genome array.
void scramble_genomes() {
  for (int i=0; i < 4*POPULATION_SIZE; i++) {
    swap_genomes(rand()%POPULATION_SIZE, rand()%POPULATION_SIZE);
  }
}

// Calculate a genome's fitness by determining the number of loci at which
// it matches the target.
int fitness(genome_t *genome) {
  int matches = 0;

  // scan the entire genome, counting the number of loci that match the target
  for (int i=0; i < GENOME_LEN; i++) {
    if (genome->locus[i] == target.locus[i]) ++matches;
  }

  return matches;
}

// Display a genome as a character string.
void display_genome(genome_t *genome) {
  for (int i=0; i < GENOME_LEN; i++) {
    putchar(weasel_to_ascii(genome->locus[i]));
  }
}

// Convert an ascii string to a Weasel genome
void weaselize (char *ascii_str, genome_t *genome) {
  for (int i=0; i<GENOME_LEN; i++) {
    genome->locus[i] = ascii_to_weasel(ascii_str[i]);
  }
}

// move cursor to the home position
void cursor_home() {
  printf("\x1b[H");
}

// overwrite everything on the screen with spaces
void clear_screen() {
  cursor_home();
  for (int i=0; i < NUM_LINES; i++) {
    for (int j=0; j < CHARS_PER_LINE; j++) {
      putchar(' ');
    }
    printf("\n\r");
  }
}

// Mutate 'old' genome, placing results in 'new'
void mutate(genome_t *old, genome_t *new) {
  for (int i=0; i < GENOME_LEN; i++) {
    if ((double)rand()/(double)RAND_MAX < MUTATION_RATE && (!ENABLE_LATCHING || old->locus[i] != target.locus[i])) {
        new->locus[i] = (char)(rand()%CHARSET_SIZE); // mutate this locus
    } else {
      new->locus[i] = old->locus[i];        // don't mutate
    }
  }
  new->fitness = fitness(new);
}

// Initialize a genome with a random string
void init_genome_rand(genome_t *genome) {
  for (int i=0; i < GENOME_LEN; i++) {
    genome->locus[i] = (char)(rand()%CHARSET_SIZE);
  }
  genome->fitness = fitness(genome);
}

// Comparison function for use by qsort() library routine
int compare_fitness (const void * elem1, const void * elem2) {
    if (((genome_t *)elem2)->fitness  > ((genome_t *)elem1)->fitness) return 1;
    if (((genome_t *)elem2)->fitness == ((genome_t *)elem1)->fitness) return 0;
    return -1;
}

//
// From http://stackoverflow.com/questions/448944/c-non-blocking-keyboard-input
//

struct termios orig_termios;

void reset_terminal_mode()
{
    tcsetattr(0, TCSANOW, &orig_termios);
}

void set_conio_terminal_mode()
{
    struct termios new_termios;

    /* take two copies - one for now, one for later */
    tcgetattr(0, &orig_termios);
    memcpy(&new_termios, &orig_termios, sizeof(new_termios));

    /* register cleanup handler, and set the new terminal mode */
    atexit(reset_terminal_mode);
    cfmakeraw(&new_termios);
    tcsetattr(0, TCSANOW, &new_termios);
}

int kbhit()
{
    struct timeval tv = { 0L, 0L };
    fd_set fds;
    FD_ZERO(&fds);
    FD_SET(0, &fds);
    return select(1, &fds, NULL, NULL, &tv);
}

int getch()
{
    int r;
    unsigned char c;
    if ((r = read(0, &c, sizeof(c))) < 0) {
        return r;
    } else {
        return c;
    }
}

int main() {

  long long generation = 0;               // generation count
  int i;
  int selection_enabled = ENABLE_SELECTION;
  char ascii_target[GENOME_LEN+1];
  char c;

  set_conio_terminal_mode();

  printf("\n\n\n\n\r");

  // set cursor to home position, clear screen, and rehome
  clear_screen();
  cursor_home();

  // seed the random number generator with the current time
  srand(time(NULL));

  // set up ascii version of initial fitness target
  ascii_target[0]  = 'M';
  ascii_target[1]  = 'E';
  ascii_target[2]  = 'T';
  ascii_target[3]  = 'H';
  ascii_target[4]  = 'I';
  ascii_target[5]  = 'N';
  ascii_target[6]  = 'K';
  ascii_target[7]  = 'S';
  ascii_target[8]  = ' ';
  ascii_target[9]  = 'I';
  ascii_target[10] = 'T';
  ascii_target[11] = ' ';
  ascii_target[12] = 'I';
  ascii_target[13] = 'S';
  ascii_target[14] = ' ';
  ascii_target[15] = 'L';
  ascii_target[16] = 'I';
  ascii_target[17] = 'K';
  ascii_target[18] = 'E';
  ascii_target[19] = ' ';
  ascii_target[20] = 'A';
  ascii_target[21] = ' ';
  ascii_target[22] = 'W';
  ascii_target[23] = 'E';
  ascii_target[24] = 'A';
  ascii_target[25] = 'S';
  ascii_target[26] = 'E';
  ascii_target[27] = 'L';

  // set up weasel version of the fitness target
  weaselize(ascii_target, &target);

  // randomly initialize the genomes of the entire population
  for (i=0; i < POPULATION_SIZE; i++) {
    init_genome_rand(&genome_array[i]);
  }
  printf("\n\n\rInitial Organism 0 Genome:\n\n\r");
  display_genome(&genome_array[0]);
  printf("\n\n\n\r");

  while (genome_array[0].fitness < FITNESS_THRESHOLD) {

    // preserve the survivors, but convert the rest into mutated copies of the survivors
    for (i=0; i < POPULATION_SIZE - NUM_SURVIVORS; i++) {
      mutate(&genome_array[i%NUM_SURVIVORS], &genome_array[i+NUM_SURVIVORS]);
    }

    // now mutate the survivors themselves
    for (i=0; i<NUM_SURVIVORS; i++) {
      mutate(&genome_array[i], &genome_array[i]);
    }

    if (selection_enabled) {
      // sort the genomes into descending order by fitness
      qsort(genome_array, POPULATION_SIZE, sizeof(genome_t), compare_fitness);
    } else {
      // randomly scramble the genome array
      scramble_genomes();
    }

    ++generation;

    // in step mode, pause and display all genomes periodically
    if (STEP_MODE && generation % DISPLAY_INTERVAL == 0) {
      clear_screen();
      cursor_home();
      printf("\n\rTarget Phrase: ");
      printf("%s", ascii_target);
      printf("\n\n\rGeneration: %lld\n\r", generation);
      printf("\n\rSelection: %s\n\n\r", selection_enabled ? "ON" : "OFF");
      for (i = 0; i < GENOMES_TO_DISPLAY; i++) {
        printf("Organism %d  Fitness %d\n\n\r", i, genome_array[i].fitness);
        display_genome(&genome_array[i]);
        printf("\r\n\n\n");
      }
      usleep(PAUSE_TIME); // pause briefly before continuing
    }

    if (!STEP_MODE && generation % DISPLAY_INTERVAL == 0) {
      printf("generation = %lld  fitness = %d \n\r", generation, genome_array[0].fitness);
    }

    if (kbhit()) { // keystroke detected
      switch(getch()) {
        case 'q': // quit the program
          exit(0);

        case 'p': // pause until next keystroke
          printf("Paused: Hit any key to continue\r\n");
          while(!kbhit()){
            ;
          }
          (void)getch(); // consume character
          break;

        case 's': // toggle selection on and off
          selection_enabled = !selection_enabled;
          break;

        case 't': // set new target phrase
          printf("Enter new target phrase (spaces and letters only): ");
          fflush(stdout);
          i = 0;
          do {
            while (!kbhit()) {      // wait for keystroke
              ;
            }
            c = toupper(getch());
            if (c == '\r') break;
            if (isalpha(c) || c == ' ') {
              putchar(c);
              fflush(stdout);
              ascii_target[i++] = c;   // add character to target phrase
            }
          } while (ascii_target[i-1] != '\r' && i-1 < GENOME_LEN);
          for(; i<GENOME_LEN; i++) {
            ascii_target[i] = ' ';  // pad remainder of target with spaces
          }
          weaselize(ascii_target, &target);
          // need to recompute fitness values because the target has changed
          for (i=0; i<POPULATION_SIZE; i++) {
            genome_array[i].fitness = fitness(&genome_array[i]);
          }
          break;
      }
    }
  }

  printf("\n\n\rWinning generation = %lld  fitness = %d \n\r", generation, genome_array[0].fitness);
  printf("\n\rWinning genome:\n\n\r");
  display_genome(&genome_array[0]);
  printf("\n\n\r");
  tcsetattr(0, TCSANOW, &orig_termios);
}