#ifndef NE_H
#define NE_H

#include <math.h>
#include <stdlib.h>

// From stretchy_buffer.h version 1.03 (https://github.com/nothings/stb)
#ifndef STB_STRETCHY_BUFFER_H_INCLUDED
#define STB_STRETCHY_BUFFER_H_INCLUDED

#ifndef NO_STRETCHY_BUFFER_SHORT_NAMES
#define sb_free   stb_sb_free
#define sb_push   stb_sb_push
#define sb_count  stb_sb_count
#define sb_add    stb_sb_add
#define sb_last   stb_sb_last
#endif

#define stb_sb_free(a)         ((a) ? free(stb__sbraw(a)),0 : 0)
#define stb_sb_push(a,v)       (stb__sbmaybegrow(a,1), (a)[stb__sbn(a)++] = (v))
#define stb_sb_count(a)        ((a) ? stb__sbn(a) : 0)
#define stb_sb_add(a,n)        (stb__sbmaybegrow(a,n), stb__sbn(a)+=(n), &(a)[stb__sbn(a)-(n)])
#define stb_sb_last(a)         ((a)[stb__sbn(a)-1])

#define stb__sbraw(a) ((int *) (a) - 2)
#define stb__sbm(a)   stb__sbraw(a)[0]
#define stb__sbn(a)   stb__sbraw(a)[1]

#define stb__sbneedgrow(a,n)  ((a)==0 || stb__sbn(a)+(n) >= stb__sbm(a))
#define stb__sbmaybegrow(a,n) (stb__sbneedgrow(a,(n)) ? stb__sbgrow(a,n) : 0)
#define stb__sbgrow(a,n)      (*((void **)&(a)) = stb__sbgrowf((a), (n), sizeof(*(a))))

static void * stb__sbgrowf(void *arr, int increment, int itemsize)
{
   int dbl_cur = arr ? 2*stb__sbm(arr) : 0;
   int min_needed = stb_sb_count(arr) + increment;
   int m = dbl_cur > min_needed ? dbl_cur : min_needed;
   int *p = (int *) realloc(arr ? stb__sbraw(arr) : 0, itemsize * m + sizeof(int)*2);
   if (p) {
      if (!arr)
         p[1] = 0;
      p[0] = m;
      return p+2;
   } else {
      #ifdef STRETCHY_BUFFER_OUT_OF_MEMORY
      STRETCHY_BUFFER_OUT_OF_MEMORY ;
      #endif
      return (void *) (2*sizeof(int)); // try to force a NULL pointer exception later
   }
}
#endif // STB_STRETCHY_BUFFER_H_INCLUDED

#ifndef NE_REAL
	#define NE_REAL float
#endif

typedef struct {
	NE_REAL gradient;
	NE_REAL bias;
	NE_REAL* weights;
} ne_Neuron;

typedef ne_Neuron* ne_Layer;

typedef struct {
	NE_REAL fitness;
	ne_Layer* layers;
} ne_Network;

typedef ne_Network* ne_Population;

double ne_RandomUnit();

double ne_RandomClamped();

ne_Population ne_RandomNetworks(
		int* topology,
		int members,
		unsigned int rand_seed );

void ne_CopyNetwork(
		ne_Network* src,
		ne_Network* dest );

void ne_FreeNetwork( ne_Network* net );

void ne_FreePopulation( ne_Population* population );

void ne_Iterate(
		ne_Population* population,
		unsigned int rand_seed,
		double crossover_rate,
		double mutation_rate,
		int mutate_gradients );

void ne_GetOutputs(
		ne_Network* net,
		NE_REAL* inputs,
		NE_REAL* outputs );

#ifdef NE_IMPLEMENTATION

/** \brief Returns a random `double` from 0.0 to 1.0.
 *
 * @return A random `double` from 0.0 to 1.0.
 */
double ne_RandomUnit(){
	return (double)rand() / (double)RAND_MAX;
}

/** \brief Returns a random `double` from -1.0 to 1.0.
 *
 * @return A random `double` from -1.0 to 1.0.
 */
double ne_RandomClamped(){
	return ( (double)rand() / (double)RAND_MAX - 0.5 ) * 2.0;
}

/** \brief Returns a stetchy buffer of randomized neural networks.
 *
 * Returns a stretchy buffer of networks with randomized weights
 * centered at 0.0 and a standard deviation of
 * `sqrt( 1.0 / input_weights )`. Biases are initialized to 0.0. As used
 * with SELU.
 *
 * `members` must be a positive even number. If it is not, the program
 * will have undefined behavior.
 *
 * Example:
 * `ne_Population population = ne_RandomNetworks( topology, 20, seed );`
 *
 * @param  topology  Stretchy buffer of network layer sizes from input
 *                   to output. The first value is the number of inputs.
 * @param  members   Number of neural networks to create.
 * @param  rand_seed A seed for srand.
 * @return Stretchy buffer of randomized neural networks.
 */
ne_Population ne_RandomNetworks(
		int* topology,
		int members,
		unsigned int rand_seed ){

	static const ne_Network NewNetwork = { 0.0, NULL };
	static const ne_Neuron NewNeuron = { 0.0, 0.0, NULL };

	srand( rand_seed );

	ne_Population population = NULL;

	// ne_Iterate through the population.
	for( int m = 0; m < members; m++ ){
		// Add a network.
		stb_sb_push( population, NewNetwork );
		ne_Network* net = &stb_sb_last( population );
		// Set the first layer's number of weights per neuron to the number of inputs.
		int input_weights = topology[0];
		// Iterate through the layers.
		for( int i = 1; i < stb_sb_count( topology ); i++ ){
			int t = topology[i];
			//assert( t >= 1 );
			// Add a layer.
			stb_sb_push( net->layers, NULL );
			ne_Layer* layer = &stb_sb_last( net->layers );
			// Set the standard deviation for the layer's weights.
			NE_REAL stddev = sqrt( 1.0 / input_weights );
			// Iterate through the neurons.
			for( int n = 0; n < t; n++ ){
				// Add a neuron.
				stb_sb_push( *layer, NewNeuron );
				ne_Neuron* neuron = &stb_sb_last( *layer );
				// Add random weights.
				for( int w = 0; w < input_weights; w++ ){
					NE_REAL weight = ne_RandomClamped() * stddev;
					stb_sb_push( neuron->weights, weight );
				}
			}
			input_weights = t;
		}
	}

	return population;
}

/** \brief Deep copy `ne_Network` `src` into empty `ne_Network` `dest`.
 *
 * @param src  Pointer to the source neural network.
 * @param dest Pointer to the destination neural network.
 */
void ne_CopyNetwork(
		ne_Network* src,
		ne_Network* dest ){

	//assert( dest->layers == NULL );
	dest->fitness = src->fitness;
	dest->layers = stb_sb_add( dest->layers, stb_sb_count( src->layers ) );
	for( int x = 0; x < stb_sb_count( src->layers ); x++ ){
		ne_Layer* src_layer = &src->layers[x];
		ne_Layer* dest_layer = &dest->layers[x];
		*dest_layer = NULL;
		*dest_layer = stb_sb_add( *dest_layer, stb_sb_count( *src_layer ) );
		for( int y = 0; y < stb_sb_count( *src_layer ); y++ ){
			ne_Neuron* src_neuron = &(*src_layer)[y];
			ne_Neuron* dest_neuron = &(*dest_layer)[y];
			dest_neuron->gradient = src_neuron->gradient;
			dest_neuron->bias = src_neuron->bias;
			dest_neuron->weights = NULL;
			dest_neuron->weights = stb_sb_add(
				dest_neuron->weights,
				stb_sb_count( src_neuron->weights )
			);
			for( int i = 0; i < stb_sb_count( src_neuron->weights ); i++ ){
				dest_neuron->weights[i] = src_neuron->weights[i];
			}
		}
	}
}

/** \brief Deep free `ne_Network` `net`.
 *
 * @param net Pointer to the neural network to be freed.
 */
void ne_FreeNetwork( ne_Network* net ){
	static const ne_Network NewNetwork = { 0.0, NULL };
	for( int x = 0; x < stb_sb_count( net->layers ); x++ ){
		ne_Layer* layer = &net->layers[x];
		for( int y = 0; y < stb_sb_count( *layer ); y++ ){
			stb_sb_free( (*layer)[y].weights );
		}
		stb_sb_free( *layer );
	}
	stb_sb_free( net->layers );
	*net = NewNetwork;
}

/** \brief Deep free a population.
 *
 * @param net Pointer to the `ne_Population` to be freed.
 */
void ne_FreePopulation( ne_Population* population ){
	for( int i = 0; i < stb_sb_count( *population ); i++ ){
		ne_FreeNetwork( &(*population)[i] );
	}
	stb_sb_free( *population );
	*population = NULL;
}

/** \brief Breeds a new generation of neural networks.
 *
 * This is the main evolution function. Feed it your population and get
 * a slightly evolved population out the other end. Do this many times
 * to optimize for the fitness function.
 *
 * The input and output populations will always have the same count,
 * with the exception that passing in a population of negative or
 * non-even count results in undefined behavior.
 *
 * Make sure you set the `fitness` of each network before calling this
 * function.
 *
 * Example:
 * `ne_Iterate( &population, seed, 0.7, 0.001 );`
 *
 * @param population       A stretchy buffer of neural networks with a
 *                         homogenous topology.
 * @param rand_seed        A seed for `srand`.
 * @param crossover_rate   The probability that 2 chosen networks will
 *                         swap weights.
 * @param mutation_rate    A weight's probability of mutating.
 * @param mutate_gradients Whether or not to allow `gradient` values to
 *                         mutate.
 */
void ne_Iterate(
		ne_Population* population,
		unsigned int rand_seed,
		double crossover_rate,
		double mutation_rate,
		int mutate_gradients ){

	static const ne_Network NewNetwork = { 0.0, NULL };

	// Get the number of network inputs.
	int num_inputs = stb_sb_count( (*population)[0].layers[0][0].weights );

	// Recover the total number of weights from the first network.
	int
		num_weights = 0,
		input_weights = num_inputs;
	for( int i = 0; i < stb_sb_count( (*population)[0].layers ); i++ ){
		int layer_size = stb_sb_count( (*population)[0].layers[i] );
		// The bias counts as an extra weight.
		num_weights += layer_size * ( input_weights + 1 );
		input_weights = layer_size;
	}

	double total_fitness = 0.0;

	// Add up the fitnesses.
	for( int i = 0; i < stb_sb_count( *population ); i++ ){
		NE_REAL* fit = &(*population)[i].fitness;
		// Set negative or NaN fitnesses to 0, lest they cause a crash.
		if( *fit < 0.0 || isnan( *fit ) ){
			*fit = 0.0;
		}
		total_fitness += (double)( *fit );
	}

	srand( rand_seed );

	ne_Population new_population = NULL;

	// ne_Iterate through the population, producing 2 new members for every 2 old members.
	for( int half_n = 0; half_n < stb_sb_count( *population ) / 2; half_n++ ){
		// Create 2 networks.
		ne_Network
			net1 = NewNetwork,
			net2 = NewNetwork;
		// Pick 2 mates using roulette wheel selection.
		double
			slot = 0.0,
			ball1 = ne_RandomUnit() * total_fitness,
			ball2 = ne_RandomUnit() * total_fitness;
		for( int i = 0; i < stb_sb_count( *population ); i++ ){
			ne_Network* p = &(*population)[i];
			if( ball1 >= slot && ball1 <= slot + p->fitness ){
				ne_CopyNetwork( p, &net1 );
			}
			if( ball2 >= slot && ball2 <= slot + p->fitness ){
				ne_CopyNetwork( p, &net2 );
			}
			if( stb_sb_count( net1.layers ) > 0 && stb_sb_count( net2.layers ) > 0 ){
				break;
			}
			slot += fmax( 0.0, (double)p->fitness );
		}

		// Average old fitness for easy access if anything uses it.
		NE_REAL old_fitness = ( net1.fitness + net2.fitness ) * 0.5;
		net1.fitness = old_fitness;
		net2.fitness = old_fitness;

		// Swap weights for crossover_rate of couples.
		if( ne_RandomUnit() <= crossover_rate ){
			// Index weights by order of access as they are not contiguous in memory.
			int
				swap_offset = ne_RandomUnit() * num_weights,
				idx = 0;
			for( int l = 0; l < stb_sb_count( net1.layers ); l++ ){
				if( idx >= swap_offset ){
					break;
				}
				ne_Layer* layer1 = &net1.layers[l];
				ne_Layer* layer2 = &net2.layers[l];
				for( int n = 0; n < stb_sb_count( *layer1 ); n++ ){
					if( idx >= swap_offset ){
						break;
					}
					ne_Neuron* neuron1 = &(*layer1)[n];
					ne_Neuron* neuron2 = &(*layer2)[n];
					// Swap gradients as if they are part of the biases.
					// It is mathematically questionable to group both
					// values as one gene, but the entire neuron depends
					// on the gradient if a backpropagation step is
					// used, so some kind of grouping may be desirable.
					NE_REAL weight_tmp = neuron1->gradient;
					neuron1->gradient = neuron2->gradient;
					neuron2->gradient = weight_tmp;
					// Swap biases.
					weight_tmp = neuron1->bias;
					neuron1->bias = neuron2->bias;
					neuron2->bias = weight_tmp;
					idx++;
					for( int w = 0; w < stb_sb_count( neuron1->weights ); w++ ){
						if( idx >= swap_offset ){
							break;
						}
						// Swap weights.
						weight_tmp = neuron1->weights[w];
						neuron1->weights[w] = neuron2->weights[w];
						neuron2->weights[w] = weight_tmp;
						idx++;
					}
				}
			}
		}

		// Starting number of weights per neuron.
		input_weights = num_inputs;

		// Mutate the weights.
		for( int l = 0; l < stb_sb_count( net1.layers ); l++ ){
			ne_Layer* layer1 = &net1.layers[l];
			ne_Layer* layer2 = &net2.layers[l];
			// Set the standard deviation for the layers' weights.
			NE_REAL stddev = sqrt( 1.0 / input_weights );
			for( int n = 0; n < stb_sb_count( *layer1 ); n++ ){
				ne_Neuron* neuron1 = &(*layer1)[n];
				ne_Neuron* neuron2 = &(*layer2)[n];
				// If a particle hits a real, re-randomize it.
				if( mutate_gradients ){
					// TODO: Investigate different standard deviations.
					if( ne_RandomUnit() <= mutation_rate ){
						neuron1->gradient = ne_RandomClamped() * stddev;
					}
					if( ne_RandomUnit() <= mutation_rate ){
						neuron2->gradient = ne_RandomClamped() * stddev;
					}
				}
				if( ne_RandomUnit() <= mutation_rate ){
					neuron1->bias = ne_RandomClamped() * stddev;
				}
				if( ne_RandomUnit() <= mutation_rate ){
					neuron2->bias = ne_RandomClamped() * stddev;
				}
				for( int w = 0; w < stb_sb_count( neuron1->weights ); w++ ){
					if( ne_RandomUnit() <= mutation_rate ){
						neuron1->weights[w] = ne_RandomClamped() * stddev;
					}
					if( ne_RandomUnit() <= mutation_rate ){
						neuron2->weights[w] = ne_RandomClamped() * stddev;
					}
				}
			}
			input_weights = stb_sb_count( *layer1 );
		}

		// Add the networks to the new population.
		stb_sb_push( new_population, net1 );
		stb_sb_push( new_population, net2 );
	}

	for( int i = 0; i < stb_sb_count( *population ); i++ ){
		// Deep free the old data.
		ne_FreeNetwork( &(*population)[i] );
		// Shallow copy the memory address.
		(*population)[i] = new_population[i];
	}
}

/** \brief Runs data through a neural network.
 *
 * This is the main neural network function. It ingests an array of
 * values and outputs an array of other values that hopefully do what
 * you want them to do.
 *
 * The quality of the outputs depends on the quality of the neural
 * network. If you `ne_Iterate` the population the right number of times
 * with a smartly constructed fitness function and a good enough network
 * topology, then run the fittest member through `ne_GetOutputs`, you
 * should get high-quality outputs from your inputs.
 *
 * If either your `inputs` or `outputs` array is not large enough for
 * the network, you could get access violations, segmentation faults, or
 * undefined behavior. The library has no way to check the size of these
 * arrays.
 *
 * Stretchy buffers can be used instead of arrays for dynamic sizes of
 * `inputs` and `outputs`. (Think scripting and dynamic topologies.)
 * However, this function never checks the counts of these buffers.
 *
 * Example:
 * `ne_GetOutputs( &net, inputs, outputs );`
 *
 * @param net     Pointer to a neural network.
 * @param inputs  Array of `NE_REAL` numbers to be fed into the
 *                network's input layer. All numbers are valid. The size
 *                of the array must be at least the size of the
 *                network's input layer.
 * @param outputs Array of `NE_REAL` numbers to be filled by the
 *                network. The size of the array must be at least the
 *                size of the network's output layer.
 */
void ne_GetOutputs(
		ne_Network* net,
		NE_REAL* inputs,
		NE_REAL* outputs ){

	// SELU constants.
	static const NE_REAL a = 1.6732632423543772848170429916717;
	static const NE_REAL s = 1.0507009873554804934193349852946;

	// Get the number of network inputs.
	int num_inputs = stb_sb_count( net->layers[0][0].weights );

	// Create temporary input and output buffers.
	NE_REAL* in_tmp = NULL;
	NE_REAL* out_tmp = NULL;
	// The inputs array is assumed to be large enough to contain the
	// network's inputs. This keeps the API syntax simple.
	// Copy inputs into in_tmp.
	in_tmp = stb_sb_add( in_tmp, num_inputs );
	for( int i = 0; i < num_inputs; i++ ){
		in_tmp[i] = inputs[i];
	}

	for( int x = 0; x < stb_sb_count( net->layers ); x++ ){
		ne_Layer* layer = &net->layers[x];
		for( int y = 0; y < stb_sb_count( *layer ); y++ ){
			ne_Neuron* neuron = &(*layer)[y];
			// "treat the threshold as a weight that is always
			//  multiplied by an input of -1. This is usually referred
			//  to as the bias."
			NE_REAL activation = -1.0 * neuron->bias;
			for( int i = 0; i < stb_sb_count( neuron->weights ); i++ ){
				activation += in_tmp[i] * neuron->weights[i];
			}
			// SELU
			if( activation < 0.0 ){
				activation = a * exp( activation ) - a;
			}
			activation *= s;
			stb_sb_push( out_tmp, activation );
		}
		stb_sb_free( in_tmp );
		// Shallow copy the memory address.
		in_tmp = out_tmp;
		out_tmp = NULL;
	}

	// The outputs array is assumed to be large enough to contain the
	// network's outputs.
	for( int i = 0; i < stb_sb_count( in_tmp ); i++ ){
		outputs[i] = in_tmp[i];
	}
	stb_sb_free( in_tmp );
}

#endif // #ifdef NE_IMPLEMENTATION

#endif // #ifndef NE_H