// mctrack - Motion-tracking demo for linmctool.
// Copyright (c) 2010,2011 pabr@pabr.org
// See http://www.pabr.org/linmctool/
//
// Compile with: gcc --std=gnu99 -Wall mctrack-20110304.c -lX11 -lm -o mctrack

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

#include <X11/Xlib.h>
#include <X11/Xutil.h>

/***** Utilities *****/

void fatal(const char *msg) { perror(msg); exit(1); }

Display *display;
int screen;
Window window;
GC gc;
char *colors[] = { "white", "red", "green", "blue" };
unsigned long pixels[4];

void xopen(int width, int height) {
  display = XOpenDisplay(getenv("DISPLAY"));
  if ( ! display ) fatal("XOpenDisplay");
  screen = DefaultScreen(display);
  XSetWindowAttributes xswa;
  xswa.background_pixel = BlackPixel(display, screen);
  window = XCreateWindow(display, DefaultRootWindow(display), 
			 0, 0, width, height, 1,
			 CopyFromParent, InputOutput,
			 CopyFromParent, CWBackPixel, &xswa);
  if ( ! window ) fatal("XCreateWindow");
  XMapWindow(display, window);
  gc = XCreateGC(display, window, 0, NULL);
  if ( ! gc ) fatal("XCreateGC");
  Colormap cmap = DefaultColormap(display, screen);
  for ( int i=0; i<sizeof(colors)/sizeof(colors[0]); ++i ) {
    XColor c;
    if ( ! XAllocNamedColor(display, cmap, colors[i], &c, &c) )
      fatal("AllocNamedColor");
    pixels[i] = c.pixel;
  }
}

/**********************************************************************/
// Math

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

void normalize(const num u[3], num res[3]) {
  num k = 1.0 / SQRT(u[0]*u[0]+u[1]*u[1]+u[2]*u[2]);
  for ( int i=0; i<3; ++i ) res[i] = u[i] * k;
}

void cross(const num u[3], const num v[3], num res[3]) {
  res[0] = u[1]*v[2] - u[2]*v[1];
  res[1] = u[2]*v[0] - u[0]*v[2];
  res[2] = u[0]*v[1] - u[1]*v[0];
}

/**********************************************************************/
// Tracking

struct tracker {
  int initialized;
  int seqnum;
  num R[3][3];  // Rotates device coordinates to global coordinates.
                // R[0] is global x-axis (EAST) expressed in local coordinates.
                // R[1] is global y-axis (NORTH) expressed in local coordinates.
                // R[2] is global z-axis (UP) expressed in local coordinates.
  num p[3];     // Position of device origin in global frame.
  num v[3];     // Velocity of device origin in global frame.
  num A0[3];    // Lowpass-filtered acceleration in local coordinates.
  num Gz0;      // Lowpass-filtered Z gyro.
};

// Compute absolute orientation from accelerometers and magnetometers.
// Assume accelerometers point UP (-gravity only, device not accelerating).
// Assume magnetic field points SOUTH (with unknown elevation).

void absolute_orientation(const num A[3], const num M[3], num R[3][3]) {
  // A = -gravity = UP direction expressed in local coordinates => R[2].
  normalize(A, R[2]);
  // UP ^ MagneticField = EAST direction in local coordinates => R[0].
  cross(R[2], M, R[0]);
  normalize(R[0], R[0]);
  // UP ^ EAST = global NORTH direction in local coordinates => R[1].
  cross(R[2], R[0], R[1]);
}

// Track 6DoF position+orientation with accelerometers, gyros and magnetometers.

void update_tracker_AGM(struct tracker *t, num dt,
			const num A[3], const num G[3], const num M[3]){
  if ( ! t->initialized ) {
    memcpy(t->A0, A, sizeof(t->A0));
    absolute_orientation(A, M, t->R);
    t->initialized = 1;
  }

  num OmegaX[3][3] = { {     0, -G[2],  G[1] },
		       {  G[2],     0, -G[0] },
		       { -G[1],  G[0],     0 } };
  num oldR[3][3];  memcpy(oldR, t->R, sizeof(oldR));

  // Integrate rotation rate: R += R * Omega_X * dt
  for ( int i=0; i<3; ++i )
    for ( int j=0; j<3; ++j )
      for ( int k=0; k<3; ++k )
	t->R[i][j] += oldR[i][k] * OmegaX[k][j] * dt;

  // Bias R toward (noisy) absolute orientation: R = R0*k + R*(1-k)
  num kR = 0.01;
  num R0[3][3];
  absolute_orientation(t->A0, M, R0);
  for ( int i=0; i<3; ++i )
    for ( int j=0; j<3; ++j ) t->R[i][j] = R0[i][j]*kR + t->R[i][j]*(1-kR);

  // Make R orthonormal.
  normalize(t->R[2], t->R[2]);
  cross(t->R[2], t->R[0], t->R[1]);
  normalize(t->R[1], t->R[1]);
  cross(t->R[1], t->R[2], t->R[0]);

  // Rotate the lowpass-filtered acceleration: A0 -= (Omega^A0)*dt
  num dA0[3];  cross(G, t->A0, dA0);
  for ( int i=0; i<3; ++i ) t->A0[i] -= dA0[i] * dt;
  
  // Lowpass-filter the acceleration: A0 = A*k + A0*(1-k)
  num kA = 0.001;
  for ( int i=0; i<3; ++i ) t->A0[i] = A[i]*kA + t->A0[i]*(1-kA);

  // Integrate acceleration: v += R * (A-A0)*dt
  num a[3];
  for ( int i=0; i<3; ++i ) {
    a[i] = 0;
    for ( int k=0; k<3; ++k ) a[i] += t->R[i][k] * (A[k]-t->A0[k]);
  }
  for ( int i=0; i<3; ++i ) t->v[i] += a[i] * dt;

  // Integrate velocity: p += v*dt
  for ( int i=0; i<3; ++i ) t->p[i] += t->v[i]*dt;

  // Damp to prevent integration drift on position and velocity.
  for ( int i=0; i<3; ++i ) { t->v[i]*=0.90; t->p[i]*=0.990; }
}

// Track 3DoF orientation with accelerometers and one gyro.

void update_tracker_tilt_gz(struct tracker *t, num dt, num A[3], num Gz) {
  if ( ! t->initialized ) {
    num M[3] = { 0, -1, 0 };  // Pretend Y points NORTH initially.
    absolute_orientation(A, M, t->R);
    t->Gz0 = Gz;
    t->initialized = 1;
  }

  // Lowpass-filter the gyro: Gz0 = Gz*k + Gz0*(1-k)
  num kG = 0.01;
  for ( int i=0; i<3; ++i ) t->Gz0 = Gz*kG + t->Gz0*(1-kG);
  
  // R[1] is the global y-axis (NORTH) in local coordinates.
  // Rotate it and use it to simulate the local magnetic field M.
  // R[1] -= (Omega^R[1])*dt
  // SimulatedM = - R'[1]
  num G[3] = { 0, 0, Gz-t->Gz0 };
  num dR1[3];  cross(G, t->R[1], dR1);
  num M[3];  for ( int i=0; i<3; ++i ) M[i] = - (t->R[1][i] -= dR1[i]*dt);

  // Bias NORTH toward front of device.
  num R_rate = 0.3;
  M[1] -= R_rate * dt;

  absolute_orientation(A, M, t->R);
}

// Estimate 3DoF orientation with accelerometers and one IR beacon.
// Result is incorrect if device is accelerating.

struct blob { int x, y, s; };  // In sensor coordinates.

void update_tracker_tilt_ir(struct tracker *t, num dt,
			    num A[3], struct blob blobs[4]) {
  t->initialized = 1;
  struct blob *b = NULL;
  for ( int i=0; i<4; ++i )
    if ( blobs[i].s!=15 && (!b||blobs[i].s>b->s) ) b = &blobs[i];
  if ( b ) {
    // Assume IR beacon points NORTH.  Use that to simulate compass.
    num focal_h = 1382;
    num focal_v = 1000;
    num M[3] = { -(b->x-512)/focal_h, -1, -(b->y-430)/focal_v };
    absolute_orientation(A, M, t->R);
  } else {
    // Beacon not found. Pretend Y-axis of device points NORTH.
    num M[3] = { 0, -1, 0 };
    absolute_orientation(A, M, t->R);
  }
}

// Estimate 2DoF tilt with accelerometers only.
// Result is incorrect if device is accelerating.

void update_tracker_tilt(struct tracker *t, num dt, num A[3]) {
  t->initialized = 1;
  // Pretend Y-axis of device points NORTH (with unknown elevation).
  num M[3] = { 0, -1, 0 };
  absolute_orientation(A, M, t->R);
}

/**********************************************************************/

int main(int argc, char *argv[]) {
  int ww = 640, wh = 480;
  xopen(ww, wh);

  void draw_3dline(int screenx0, struct tracker *t,
		   num X0,num Y0,num Z0, num X1,num Y1,num Z1) {
    void proj(num X, num Y, num Z, int screen[2]) {
      num x = t->p[0] + t->R[0][0]*X + t->R[0][1]*Y + t->R[0][2]*Z;
      num y = t->p[1] + t->R[1][0]*X + t->R[1][1]*Y + t->R[1][2]*Z;
      num z = t->p[2] + t->R[2][0]*X + t->R[2][1]*Y + t->R[2][2]*Z;
#ifdef STICKYTAPE_DEMO  // Magnetic north is toward the left of the LCD monitor.
      if ( screenx0 == ww/2 ) { num tmp = y; y=x; x=-tmp; }
#endif
      y += 0.15;
      screen[0] = screenx0 + (wh/2)*x/y;
      screen[1] = wh/2 - (wh/2)*z/y;
    }
    int s0[2], s1[2];
    proj(X0,Y0,Z0, s0);
    proj(X1,Y1,Z1, s1);
    XDrawLine(display, window, gc, s0[0],s0[1], s1[0],s1[1]);
    ++s0[0]; ++s1[0];
    XDrawLine(display, window, gc, s0[0],s0[1], s1[0],s1[1]);
    ++s0[1]; ++s1[1];
    XDrawLine(display, window, gc, s0[0],s0[1], s1[0],s1[1]);
    --s0[0]; --s1[0];
    XDrawLine(display, window, gc, s0[0],s0[1], s1[0],s1[1]);
  }

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

  void draw() {
    XClearWindow(display, window);
    XSetForeground(display, gc, pixels[0]);
    XDrawString(display, window, gc, ww/2, 16,"UP", 2);
    XDrawString(display, window, gc, 16, wh/2,"WEST", 4);
    int ndevs = 0;
    for ( int i=0; i<MAXDEVS; ++i ) if ( devs[i].initialized ) ++ndevs;
    for ( int i=0; i<MAXDEVS; ++i ) {
      if ( ! devs[i].initialized ) continue;
      XSetForeground(display, gc, pixels[i]);
      int sx0 = (ww/2+i*ww)/ndevs;
      XDrawPoint(display, window, gc, sx0, wh/2);
      num W=0.02, H=0.01, D=0.1;
      draw_3dline(sx0, &devs[i], -W,0,-H, -W,D,-H);
      draw_3dline(sx0, &devs[i], -W,0,+H, -W,D,+H);
      draw_3dline(sx0, &devs[i], +W,0,-H, +W,D,-H);
      draw_3dline(sx0, &devs[i], +W,0,+H, +W,D,+H);
      draw_3dline(sx0, &devs[i], -W,0,-H, -W,0,+H);
      draw_3dline(sx0, &devs[i], -W,0,+H, +W,0,+H);
      draw_3dline(sx0, &devs[i], +W,0,+H, +W,0,-H);
      draw_3dline(sx0, &devs[i], +W,0,-H, -W,0,-H);
      draw_3dline(sx0, &devs[i], -W,D,-H, -W,D,+H);
      draw_3dline(sx0, &devs[i], -W,D,+H, +W,D,+H);
      draw_3dline(sx0, &devs[i], +W,D,+H, +W,D,-H);
      draw_3dline(sx0, &devs[i], +W,D,-H, -W,D,-H);
      draw_3dline(sx0, &devs[i], 0,0,0, D,0,0);
      draw_3dline(sx0, &devs[i], 0,0,0, 0,0,D);
      draw_3dline(sx0, &devs[i], 0,0,0,
		  -devs[i].A0[0]*D/9.81,
		  -devs[i].A0[1]*D/9.81,
		  -devs[i].A0[2]*D/9.81);
    }
    XSync(display, False);
  }

  void draw_vector(int di, int vi, num v[3], num range) {
    XSetForeground(display, gc, pixels[0]);
    int sx0 = ww/2+vi*100, sy0, r = 100;
    sy0 = wh * 3 / 4;
    XDrawLine(display, window, gc, sx0,sy0,sx0+r*v[0]/range,sy0-r*v[1]/range);
    sy0 = wh * 7 / 8;
    XDrawLine(display, window, gc, sx0,sy0,sx0+r*v[0]/range,sy0-r*v[2]/range);
    XSync(display, False);
  }
  void draw_intvector(int di, int vi, int v[3], num range) {
    XSetForeground(display, gc, pixels[0]);
    int sx0 = ww/2+(vi-0.5)*300, sy0, r = 100;
    sy0 = wh * 5 / 8;
    //    XDrawLine(display, window, gc, sx0,sy0,sx0+r*v[0]/range,sy0-r*v[1]/range);
    XDrawPoint(display, window, gc, sx0, sy0);
    XDrawPoint(display, window, gc, sx0+r*v[0]/range, sy0-r*v[1]/range);
    sy0 = wh * 6 / 8;
    //    XDrawLine(display, window, gc, sx0,sy0,sx0+r*v[0]/range,sy0-r*v[2]/range);
    XDrawPoint(display, window, gc, sx0, sy0);
    XDrawPoint(display, window, gc, sx0+r*v[0]/range, sy0-r*v[2]/range);
    XSync(display, False);
  }
  struct timeval tdisp;
  gettimeofday(&tdisp, NULL);

  char line[256];
  while ( fgets(line, sizeof(line), stdin) ) {
    int index;
    if ( sscanf(line, "%d", &index) != 1 ) fatal(line);
    if ( index<0 || index>=MAXDEVS ) continue;
    struct tracker *dev = &devs[index];
    int seqnum;
    num a[3], g[3], m[3];
    struct blob ir[4];
    char addr[18];
    if ( sscanf(line, "%*d %17s PSMOVE seq=%d aX=%lf aY=%lf aZ=%lf"
		" gX=%lf gY=%lf gZ=%lf mX=%lf mY=%lf mZ=%lf", addr, &seqnum,
		&a[0],&a[1],&a[2],
		&g[0],&g[1],&g[2],
		&m[0],&m[1],&m[2]) == 11 ) {
      // PS Move
#if 0
      draw_intvector(index, 0, a, 4500);
      draw_intvector(index, 1, m, 200);
#endif
      num dt = 1.0/175;
      if ( ! dev->initialized ) dev->seqnum = seqnum;
      else dt *= (seqnum-dev->seqnum) & 31;
      dev->seqnum = seqnum;
      update_tracker_AGM(dev, dt, a, g, m);
#if 0
      draw_vector(index, 0, A, 9.81);
      draw_vector(index, 1, M, 1.0);
#endif
    } else if ( sscanf(line, "%*d %17s %*s aX=%lf aY=%lf aZ=%lf gZ=%lf",
		       addr, &a[0],&a[1],&a[2], &g[2]) == 5 ) {
      // SIXAXIS/DS3
      a[2] = -a[2];  // Az points down on this device.
      num dt = 1.0 / 88;
#ifdef STICKYTAPE_DEMO  // SIXAXIS is upside-down.
      a[0] = - a[0];
      a[2] = - a[2];
      g[2] = - g[2];
#endif
      update_tracker_tilt_gz(dev, dt, a, g[2]);
    } else if ( sscanf(line, "%*d %17s WIIMOTE aX=%lf aY=%lf aZ=%lf"
		       " ir0x=%d ir0y=%d ir0s=%d ir1x=%d ir1y=%d ir1s=%d"
		       " ir2x=%d ir2y=%d ir2s=%d ir3x=%d ir3y=%d ir3s=%d",
		       addr,
		       &a[0], &a[1], &a[2],
		       &ir[0].x, &ir[0].y, &ir[0].s,
		       &ir[1].x, &ir[1].y, &ir[1].s,
		       &ir[2].x, &ir[2].y, &ir[2].s,
		       &ir[3].x, &ir[3].y, &ir[3].s) == 1+3+12 ) {
      // Wiimote
      num dt = 1.0 / 98;
      update_tracker_tilt_ir(dev, dt, a, ir);
    } else fatal(line);
    
    struct timeval now;
    gettimeofday(&now, NULL);
    int display_fps = 20;
    int age = (now.tv_sec-tdisp.tv_sec)*1000000 + (now.tv_usec-tdisp.tv_usec);
    if ( age >= 1000000/display_fps ) { draw(); tdisp=now; }
  }

  return 0;
}