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

#include <stdio.h>

#include "pso.h"

void printarr(const double* arr, const int len){
	int i;
	printf("[ ");
	for(i=0;i<len;i++)
		printf("%f ",arr[i]);
	printf("]\n");
}

void init_swarm(const unsigned int particles, const unsigned int dimension, 
				const unsigned int neighborhood, const double alpha, 
				const double beta, const double gamma, const double delta, 
				const double epsilon,
				double (*evaluate)(const double *particle, const unsigned int dim), 
				struct swarm *s){
	unsigned int i;
	s->num_particles = particles;
	s->dimensions = dimension;
	s->neighborhood = neighborhood;
	s->alpha = alpha;
	s->beta = beta;
	s->gamma = gamma;
	s->delta = delta;
	s->epsilon = epsilon;
	
	s->particles = (double**)(malloc(particles*sizeof(double*)));
	s->velocities = (double**)(malloc(particles*sizeof(double*)));
	s->personal_bests = (double**)(malloc(particles*sizeof(double*)));

	for(i=0;i<particles;i++){
		s->particles[i] = (double*)(malloc(dimension*sizeof(double)));
		s->velocities[i] = (double*)(malloc(dimension*sizeof(double)));
		s->personal_bests[i] = (double*)(malloc(dimension*sizeof(double)));
	}
	s->global_best = (double*)(malloc(dimension*sizeof(double)));
	s->evaluate = evaluate;
}

void free_swarm(struct swarm *s){
	unsigned int i;
	for(i=0;i<s->num_particles;i++){
		free(s->particles[i]);
		free(s->velocities[i]);
		free(s->personal_bests[i]);
		s->particles[i] = s->velocities[i] = s->personal_bests[i] = NULL;
	}
	free(s->particles);
	free(s->velocities);
	free(s->personal_bests);
	free(s->global_best);
	s->particles = s->velocities = s->personal_bests = NULL;
	s->global_best = NULL;
	s->evaluate = NULL;
}

void randomize(const double min, const double max, struct swarm *s){
	double scale = max-min;
	int p,d;
	for(p=0;p<s->num_particles;p++){
		for(d=0;d<s->dimensions;d++){
			s->particles[p][d] = (RNDM()*scale)+min;
			s->velocities[p][d] = (RNDM()*scale)+min;
		}
	}
}

void update_bests(const char first, struct swarm *s){
	int p;
	char gbu = first;
	double fit,old_fit,gb_fit;
	if(!first) gb_fit = s->evaluate(s->global_best,s->dimensions);
	for(p=0;p<s->num_particles;p++){
		fit = s->evaluate(s->particles[p],s->dimensions);
		if(gbu || fit > gb_fit){
			memcpy(s->global_best,s->particles[p],s->dimensions*sizeof(double));
			gb_fit = fit;
			gbu = 0;
		}
		if(first) memcpy(s->personal_bests[p],s->particles[p],s->dimensions*sizeof(double));
		else {
			old_fit = s->evaluate(s->personal_bests[p],s->dimensions);
			if(fit > old_fit) memcpy(s->personal_bests[p], s->particles[p],s->dimensions*sizeof(double));
		}
	}
}

const double* neighborhood_best(const unsigned int particle, struct swarm *s){
	double new_fit, old_fit;
	int i,p;
	const double* rv = s->particles[particle];
	old_fit = s->evaluate(rv,s->dimensions);
	for(i=0;i<s->neighborhood;i++){
		p = (int)(RNDM()*(s->num_particles));
		new_fit = s->evaluate(s->particles[p],s->dimensions);
		if(new_fit > old_fit){
			rv = s->particles[p];
			old_fit = new_fit;
		}
	}
	return rv;
}

void update_velocities(struct swarm *s){
	int p,d;
	double A,B,C,D;
	const double* best_neighbor;
	for(p=0;p<s->num_particles;p++){
		best_neighbor = neighborhood_best(p,s);
		for(d=0;d<s->dimensions;d++){
			A = s->alpha;
			B = RNDM()*s->beta ;
			C = RNDM()*s->gamma;
			D = RNDM()*s->delta;
			s->velocities[p][d] = 
				A*s->velocities[p][d] +
				B*(s->personal_bests[p][d] - s->particles[p][d]) +
				C*(best_neighbor[d] - s->particles[p][d]) +
				D*(s->global_best[d] - s->particles[p][d]);
		}
	}
}

void update_particles(struct swarm *s){
	int p,d;
	for(p=0;p<s->num_particles;p++){
		for(d=0;d<s->dimensions;d++){
			s->particles[p][d] = (s->epsilon)*(s->velocities[p][d]);
		}
	}
}

void run_one_step(const char first, struct swarm *s){
	update_bests(first,s);
	update_velocities(s);
	update_particles(s);
}

double distance(const double* vec1, const double* vec2, const int len){
	int i;
	double sum = 0.0;
	for(i=0;i<len;i++)
		sum += pow(vec1[i]-vec2[i],2);
	return sqrt(sum);
}

void run(const char use_iterations, const unsigned int num_iterations,
		 const char use_ideal, const double ideal_fitness,
		 const char use_convergence, const double distance_epsilon, 
		 struct swarm *s){
	double sum;
	int p, i = 0;
	do {
		run_one_step((i>0)?0:1,s);
		if(use_iterations && i >= num_iterations){
			printf("Out of iterations!\n");
			break;
		}
		if(use_ideal){
			if(s->evaluate(s->global_best,s->dimensions) >= ideal_fitness){
				printf("Hit the ideal fitness\n");
				break;
			}
		}
		if(use_convergence){
			sum = 0.0;
			for(p=0;p<s->num_particles;p++)
				sum += distance(s->global_best,s->particles[p],s->dimensions);
			sum = sum/((double)s->num_particles);
			if(sum<=distance_epsilon){
				printf("Particles have converged: %f\n",sum);
				break;
			}
		}
		printf("Iteration %d:\n Best Particle:",i); printarr(s->global_best,s->dimensions);
		i++;
	} while(1);
}