// mccalibrate - Calibration post-processor for linmctool.
// Copyright (c) 2010,2011 pabr@pabr.org
// See http://www.pabr.org/linmctool/
//
// Compile with: gcc --std=gnu99 -Wall mccalibrate.c -lm -o mccalibrate

#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>
#include <math.h>

#if 1
  typedef double num;
# define SQRT sqrt
#else
  typedef float num;
# define SQRT sqrtf
#endif

void fatal(char *msg) {
  if ( errno ) perror(msg); else fprintf(stderr, "%s\n", msg);
  exit(1);
}

/***** Device calibration data *****/

struct dev {
  int initialized;
  num amin[3], amax[3];  // Range for -1..+1 g
  num gmin[3], gmax[3];  // Range for -1..+1 rad/s
  num mmin[3], mmax[3];  // Range for local magnetic field (arbitrary scaling)
};

/***** Data storage *****/

void store_calibration(const struct dev *d, const char *addr) {
  char name[22];  sprintf(name, "%s.cal", addr);
  char bak[22];  sprintf(bak, "%s.bak", addr);
  rename(name, bak);
  FILE *f = fopen(name, "wc");
  if ( ! f ) fatal(name);
  for ( int i=0; i<3; ++i )
    fprintf(f, "A %f %f\n", d->amin[i], d->amax[i]);	  
  for ( int i=0; i<3; ++i )
    fprintf(f, "G %f %f\n", d->gmin[i], d->gmax[i]);	  
  for ( int i=0; i<3; ++i )
    fprintf(f, "M %f %f\n", d->mmin[i], d->mmax[i]);	  
  fclose(f);
  fprintf(stderr, "Wrote %s\n", name);
}

int read_stored_calibration(struct dev *d, const char *addr) {
  char name[22];  sprintf(name, "%s.cal", addr);
  FILE *f = fopen(name, "r");
  if ( ! f ) {
    fprintf(stderr, "%s not found, using default calibration\n", name);
    return -1;
  }
  for ( int i=0; i<3; ++i )
    if ( fscanf(f, "A %lf %lf\n", &d->amin[i], &d->amax[i]) != 2 ) fatal(name);
  for ( int i=0; i<3; ++i )
    if ( fscanf(f, "G %lf %lf\n", &d->gmin[i], &d->gmax[i]) != 2 ) fatal(name);
  for ( int i=0; i<3; ++i )
    if ( fscanf(f, "M %lf %lf\n", &d->mmin[i], &d->mmax[i]) != 2 ) fatal(name);
  fclose(f);
  fprintf(stderr, "Read calibration from %s\n", name);
  return 0;
}

void init_calibration_wiimote(struct dev *d, const char *addr) {
  d->amin[0] = d->amin[1] = d->amin[2] = -24;
  d->amax[0] = d->amax[1] = d->amax[2] =  24;
}

void init_calibration_sixaxis(struct dev *d, const char *addr) {
  d->amin[0] = -110; d->amin[1] = -110; d->amin[2] = -110;
  d->amax[0] =  110; d->amax[1] =  110; d->amax[2] =  110;
  d->gmin[2] = -50;
  d->gmax[2] =  50;
}

void init_calibration_psmove(struct dev *d, const char *addr) {
  d->amin[0] = -4300; d->amin[1] = -4300; d->amin[2] = -4300;
  d->amax[0] =  4300; d->amax[1] =  4300; d->amax[2] =  4300;
  d->gmin[0] = d->gmin[1] = d->gmin[2] = -640;
  d->gmax[0] = d->gmax[1] = d->gmax[2] =  640;
  d->mmin[0] = -150; d->mmin[1] = -150; d->mmin[2] = -150;
  d->mmax[0] =  150; d->mmax[1] =  150; d->mmax[2] =  150;
  read_stored_calibration(d, addr);
}

/***** Cost function *****/

num eval_algebraic(const num s[][3], int ns, num m[3], num M[3]) {
  num c[3] = { (m[0]+M[0])/2, (m[1]+M[1])/2, (m[2]+M[2])/2 };
  num a[3] = { (M[0]-m[0])/2, (M[1]-m[1])/2, (M[2]-m[2])/2 };
  num aa = cbrt(a[0]*a[1]*a[2]);
  num aa4 = aa*aa*aa*aa;
  num k[3] = { 1/(a[0]*a[0]), 1/(a[1]*a[1]), 1/(a[2]*a[2]) };
  num err = 0;
  for ( int i=0; i<ns; ++i ) {
    num u[3] = { s[i][0]-c[0], s[i][1]-c[1], s[i][2]-c[2] };
    num e = u[0]*u[0]*k[0] + u[1]*u[1]*k[1] + u[2]*u[2]*k[2] - 1;
    err += e*e * aa4;
  }
  return err;
}

/***** Calibration *****/

num fit_ellipsoid_randomwalk(const num s[][3], int ns, num min[3], num max[3]) {
  num best = eval_algebraic(s, ns, min, max);
  num oldbest;
  do {
    oldbest = best;
    for ( int n=0; n<100000; ++n ) {
      num k = 10.0 / n;
      num m[3], M[3];
      memcpy(m, min, sizeof(m));
      memcpy(M, max, sizeof(M));
      for ( int i=0; i<3; ++i ) {
	num r = M[i]-m[i];
	m[i] += k * r * (rand()-RAND_MAX/2) / RAND_MAX;
	M[i] += k * r * (rand()-RAND_MAX/2) / RAND_MAX;
	if ( m[i] > M[i] ) { num tmp=m[i]; m[i]=M[i]; M[i]=tmp; }
      }
      num err = eval_algebraic(s, ns, m, M);
      if ( err < best ) {
	memcpy(min, m, sizeof(m));
	memcpy(max, M, sizeof(M));
	best = err;
      }
    }
  } while ( best != oldbest );
  best /= ns;
  best = sqrt(sqrt(best));
#if 0
  fprintf(stderr, "randomwalk best %f   %f %f  %f %f  %f %f\n",
	  best, min[0],max[0], min[1],max[1], min[2],max[2]);
#endif
  return best;
}

num fit_ellipsoid(const num s[][3], int ns, num min[3], num max[3], int pct[3]) {
  num err = fit_ellipsoid_randomwalk(s, ns, min, max);
  // Check sample coverage
  num smin[3], smax[3];
  for ( int i=0; i<3; ++i ) {
    smin[i] = smax[i] = s[0][i];
    for ( int j=1; j<ns; ++j ) {
      if ( s[j][i] < smin[i] ) smin[i] = s[j][i];
      if ( s[j][i] > smax[i] ) smax[i] = s[j][i];
    }
  }
  for ( int i=0; i<3; ++i ) pct[i] = 100 * (smax[i]-smin[i]) / (max[i]-min[i]);
  return err;
}

void show_change(struct dev *p, struct dev *n) {
  void aux(char *name, num pmin[3], num pmax[3], num nmin[3], num nmax[3]) {
    for ( int i=0; i<3; ++i )
      fprintf(stderr, "%s%d   min %+6.0f ppm   max %+6.0f ppm\n", name, i,
	      1e6 * (nmin[i]-pmin[i]) / (pmax[i]-pmin[i]),
	      1e6 * (nmax[i]-pmax[i]) / (pmax[i]-pmin[i]));
  }
  aux("A", p->amin, p->amax, n->amin, n->amax);
  aux("G", p->gmin, p->gmax, n->gmin, n->gmax);
  aux("M", p->mmin, p->mmax, n->mmin, n->mmax);
}

/***** main *****/

#define MAXSAMPLES 256
num acc_samples[MAXSAMPLES][3];
num mag_samples[MAXSAMPLES][3];
int nsamples = 0;

void calibrate_AM(struct dev *d, const char *addr) {
  struct dev prevcal = *d;
  int acc_pct[3];
  num acc_err = fit_ellipsoid(acc_samples, nsamples, d->amin,d->amax, acc_pct);
  int mag_pct[3];
  num mag_err = fit_ellipsoid(mag_samples, nsamples, d->mmin,d->mmax, mag_pct);

  fprintf(stderr, "ACC: %d samples, coverage %d%% %d%% %d%%, err %f\n",
	  nsamples, acc_pct[0],acc_pct[1],acc_pct[2], acc_err);
  fprintf(stderr, "MAG: %d samples, coverage %d%% %d%% %d%%, err %f\n",
	  nsamples, mag_pct[0],mag_pct[1],mag_pct[2], mag_err);
  show_change(&prevcal, d);

  store_calibration(d, addr);
#if 0
  char plotname[26]; sprintf(plotname, "%s.calplot", addr);
  FILE *f = fopen(plotname, "wc");
  if ( f ) {
    fprintf(f, "set parametric\n");
    fprintf(f, "set trange [0:6.29]\n");
    void plotfit(num v[][3], num min[3], num max[3], int ix, int iy) {
      num x0 = (max[ix]+min[ix])/2, xr = (max[ix]-min[ix])/2;
      num y0 = (max[iy]+min[iy])/2, yr = (max[iy]-min[iy])/2;
      fprintf(f, "plot %f+%f*cos(t),%f+%f*sin(t)", x0,xr,y0,yr);
      for ( int i=0; i<nsamples; ++i )
	fprintf(f, ", %f,%f with linespoints", v[i][ix], v[i][iy]);
      fprintf(f, "\n");
      fprintf(f, "pause -1\n");
    }
    plotfit(acc_samples, d->amin, d->amax, 0, 1);
    plotfit(acc_samples, d->amin, d->amax, 0, 2);
    plotfit(acc_samples, d->amin, d->amax, 1, 2);
    plotfit(mag_samples, d->mmin, d->mmax, 0, 1);
    plotfit(mag_samples, d->mmin, d->mmax, 0, 2);
    plotfit(mag_samples, d->mmin, d->mmax, 1, 2);
    fclose(f);
    fprintf(stderr, "Wrote %s\n", plotname);
  }
#endif
}

void add_sample(struct dev *d, const char *addr, num acc[3], num mag[3]) {
  if ( nsamples == MAXSAMPLES ) exit(0);
  memcpy(acc_samples[nsamples], acc, sizeof(acc_samples[nsamples]));
  memcpy(mag_samples[nsamples], mag, sizeof(mag_samples[nsamples]));
  ++nsamples;

  if ( nsamples % 6 ) {
    fprintf(stderr, " Got %d samples, please continue.\n", nsamples);
  } else {
    fprintf(stderr, " Computing...\n");
    calibrate_AM(d, addr);
  }
}

void usage() {
  fprintf(stderr, "Usage: mccalibrate [--static] [--batch]\n");
  exit(1);
}

int main(int argc, char *argv[]) {
  int batch_mode = 0;
  int do_static_calibration = 0;

  for ( int i=1; i<argc; ++i )
    if      ( ! strcmp(argv[i], "--static") ) do_static_calibration = 1;
    else if ( ! strcmp(argv[i], "--batch") ) batch_mode = 1;
    else usage();

  int tty;
  if ( do_static_calibration && !batch_mode ) {
    tty = open("/dev/tty", O_RDONLY|O_NONBLOCK);
    if ( tty < 0 ) fatal("Need interactive terminal /dev/tty");
  } else tty = -1;

  #define MAXDEVS 4
  struct dev devs[MAXDEVS];
  memset(devs, 0, sizeof(devs));

  num acc_avg[3], mag_avg[3];
  int navg = -1;  // Number of samples in accumulators, or -1.

  char line[256];
  while ( fgets(line, sizeof(line), stdin) ) {
    int index, seq;
    char addr[18];
    int a[3], g[3], m[3];
    if ( sscanf(line, "%d %17s PSMOVE seq=%d aX=%d aY=%d aZ=%d"
		" gX=%d gY=%d gZ=%d mX=%d mY=%d mZ=%d",
		&index, addr, &seq,
		&a[0],&a[1],&a[2],
		&g[0],&g[1],&g[2],
		&m[0],&m[1],&m[2]) != 12 ) fatal(line);
    if ( index<0 || index>=MAXDEVS ) continue;
    struct dev *dev = &devs[index];
    if ( ! dev->initialized ) {
      init_calibration_psmove(dev, addr); // TBD
      dev->initialized = 1;
    }
    
    if ( do_static_calibration ) {
      static int msg_done = 0;
      if ( ! msg_done ) {
	fprintf(stderr, "Put PS MOVE in at least 6 orientations.\n"
		"Press RETURN to sample when it is stationary.\n");
	msg_done = 1;
      }

      char key;
      if ( navg<0 && (batch_mode||read(tty,&key,sizeof(key)) == sizeof(key)) ) {
	// Initiate averaging.
	memset(acc_avg, 0, sizeof(acc_avg));
	memset(mag_avg, 0, sizeof(mag_avg));
	navg = 0;
	fprintf(stderr, "Please wait 1s...");
      }
      
      if ( navg >= 0 ) {
	// Averaging in progress.
	write(5, line, strlen(line));  // Log for batch mode.
	for ( int j=0; j<3; ++j ) { acc_avg[j]+=a[j]; mag_avg[j]+=m[j]; }
	++navg;
	if ( navg == 176 ) {
	  // Averaging complete.
	  for ( int j=0; j<3; ++j ) { acc_avg[j]/=navg; mag_avg[j]/=navg; }
	  add_sample(dev, addr, acc_avg, mag_avg);
	  navg = -1;
	}
      }

    } else {
      // Flow-through mode
      num A[3], G[3], M[3];
      for ( int i=0; i<3; ++i ) {
	A[i] = (2*a[i]-(dev->amax[i]+dev->amin[i]))
	  * 9.81 / (dev->amax[i]-dev->amin[i]);
	G[i] = (2*g[i]-(dev->gmax[i]+dev->gmin[i]))
	  / (dev->gmax[i]-dev->gmin[i]);
	M[i] = (2*m[i]-(dev->mmax[i]+dev->mmin[i]))
	  / (dev->mmax[i]-dev->mmin[i]);
      }
      fprintf(stdout, "%d %s PSMOVE seq=%-2d"
	      "  aX=%-9.5lf aY=%-9.5lf aZ=%-9.5lf"
	      "  gX=%-9.5lf gY=%-9.5lf gZ=%-9.5lf"
	      "  mX=%-8.4lf mY=%-8.4lf mZ=%-8.4lf\n",
	      index, addr, seq,
	      A[0],A[1],A[2],
	      G[0],G[1],G[2],
	      M[0],M[1],M[2]);
      fflush(stdout);
    }

  }

  return 0;
}