#pragma once 

#include <algorithm>
#include <cstdlib>
#include <vector>

#include "sync.h"
#include "rand.h"

namespace genetic {

template <class T> struct Array;
template <class T> struct Stats;
template <class T> struct Strategy;
struct CellTracker;

template <class T> Stats<T> run(Strategy<T>);

template <class T> struct Strategy {
    // Number of worker threads that will be evaluating cell fitness
    int num_threads;

    int batch_size; // Number of cells a worker thread tries to work on in a row
                    // before accessing/locking the work queue again.
    int num_cells;  // Size of the population pool
    int num_generations; // Number of times (epochs) to run the algorithm
    bool test_all; // Sets whether or not every cell's fitness is evaluated every
                   // generation
    float test_chance; // Chance to test any given cell's fitness. Relevant only
                       // if test_all is false.
    bool enable_crossover; // Cells that score well in the evaluation stage
                           // produce children that replace low-scoring cells
    int crossover_parent_num;        // Number of unique high-scoring parents in a
                                     // crossover call.
    int crossover_parent_stride; // Number of parents to skip over when moving to
                                 // the next set of parents. A stride of 1 would
                                 // produce maximum overlap because the set of
                                 // parents would only change by one every
                                 // crossover.
    int crossover_children_num;  // Number of children to expect the user to
                                 // produce in the crossover function.
    bool enable_mutation;          // Cells may be mutated
                                   // before fitness evaluation
    float mutation_chance; // Chance for any given cell to be mutated cells during
                           // the mutation
    uint64_t rand_seed;
    bool higher_fitness_is_better; // Sets whether or not to consider higher
                                   // fitness values better or worse. Set this to
                                   // false if fitness is an error function.

    // User defined functions
    T (*make_default_cell)();
    void (*mutate)(T &cell_to_modify);
    void (*crossover)(const Array<T *> parents, const Array<T *> out_children);
    float (*fitness)(const T &cell);
};

template<class T> struct Stats {
    std::vector<T> best_cell;
    std::vector<float> best_cell_fitness;
};

struct CellTracker {
    float score;
    int cellid;
};

template <class T> struct Array {
    T *data;
    int len;

    T &operator[](int i) { return data[i]; }
};

template <class T> Array<T> make_array(int len) {
    return {
    .data = (T*)malloc(sizeof(T)*len),
    .len = len
     };
}

template <class T> Stats<T> run(Strategy<T> strat) {
    // Create cells
    Array<T> cells = make_array<T>(strat.num_cells);
    for (int i = 0; i < cells.len; i++) cells[i] = strat.make_default_cell();

    // Create cell trackers
    Array<CellTracker> trackers = make_array<CellTracker>(strat.num_cells);
    for (int i = 0; i < trackers.len; i++) trackers[i] = { .score=0, .cellid=i };

    // Init stat tracker
    Stats<T> stats;

    // Run the algorithm
    for (int gen = 0; gen < strat.num_generations; gen++) {
	// 1. mutate
	for (int i = 0; i < trackers.len; i++) {
	    if (abs(norm_rand(strat.rand_seed)) < strat.mutation_chance) {
		strat.mutate(cells[trackers[i].cellid]);
	    }
	}
	// 2. crossover
	if (strat.enable_crossover) {
	    int parent_end = strat.crossover_parent_num;
	    int child_begin = trackers.len-strat.crossover_children_num;
	    while (parent_end <= child_begin) {
		// Get pointers to all the parent cells
		Array<T*> parents = make_array<T*>(strat.crossover_parent_num);
		for (int i = parent_end-strat.crossover_parent_num; i < parent_end; i++) {
		    parents[i] = &cells[trackers[i].cellid];
		}

		// Get pointers to all the child cells (these will be overwritten)
		Array<T*> children = make_array<T*>(strat.crossover_children_num); 
		for (int i = child_begin; i < child_begin+strat.crossover_children_num; i++) {
		    children[i] = &cells[trackers[i].cellid];
		}
		strat.crossover(parents, children);
		parent_end += strat.crossover_parent_stride;
		child_begin -= strat.crossover_children_num;
	    }
	}
	// 3. evaluate
	if (strat.test_all) {
	    for (int i = 0; i < trackers.len; i++) {
		trackers[i].score = strat.fitness(cells[trackers[i].cellid]);
	    }
	} else {
	    for (int i = 0; i < trackers.len; i++) {
		if (abs(norm_rand(strat.rand_seed)) < strat.test_chance) {
		    trackers[i].score = strat.fitness(cells[trackers[i].cellid]);
		}
	    }
	}
	// 4. sort
	std::sort(&trackers[0], &trackers[trackers.len-1], [strat](CellTracker &a, CellTracker &b){ return strat.higher_fitness_is_better ? a.score < b.score : a.score > b.score; });

	printf("Gen: %d, Best Score: %f\n", gen, trackers[0].score);
	stats.best_cell.push_back(cells[trackers[0].cellid]);
	stats.best_cell_fitness.push_back(trackers[0].score);
    }
    return stats;
}

} // namespace genetic