/* mainopt.c    by  M. Hazen
   Exercises a PSO search.  */

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include "particleswarm.h"
#define PI 3.14159265

float spherefunc(float* pos);
float rosenbrock(float* pos);
float rastrigin(float* pos);

int main() {
  float* opt;
  float mins[2] = {-2.0, -2.0};
  float maxs[2] = {2.5, 2.5};

  printf("Starting PSO on Sphere\n");
  opt = optimize(spherefunc, mins, maxs);
  printf("Optimal value is %f at [%f, %f]\n",
          spherefunc(opt), opt[0], opt[1]);
  free(opt);

  mins[0] = 0;
  mins[1] = 0;
  maxs[0] = 5;
  maxs[1] = 5;
  printf("Starting PSO on Rosenbrock\n");
  opt = optimize(rosenbrock, mins, maxs);
  printf("Optimal value is %f at [%f, %f]\n",
          rosenbrock(opt), opt[0], opt[1]);
  free(opt);
  for(int r=0; r<100; r++) {
  mins[0] = -5;
  mins[1] = -5;
  printf("Starting PSO on Rastrigin\n");
  opt = optimize(rastrigin, mins, maxs);
  printf("Optimal value is %f at [%f, %f]\n",
          rastrigin(opt), opt[0], opt[1]);
  free(opt);
  }

  return 0;
}

// spherefunc min at 0,0
float spherefunc(float* pos) {
  return pos[0]*pos[0] + pos[1]*pos[1];
}

// rosenbrock min at 1,1
float rosenbrock(float* pos) {
  return((1-pos[0])*(1-pos[0])+
          100*(pos[1]-pos[0]*pos[0])*(pos[1]-pos[0]*pos[0]));
}

// rastrigin min at 0,0
float rastrigin(float* pos) {
  float tmp = 20;
  tmp = tmp + pos[0]*pos[0] - 10 * cos(2*PI*pos[0]);
  tmp = tmp + pos[1]*pos[1] - 10 * cos(2*PI*pos[1]);
  return tmp;
}