/* nsfo.c (new utility for Sigma/Phi/Omega computation) 08-12jan99 jw */

#undef SHOWPARAMS
#undef DEBUG

#include <errno.h>
#include "std.h"
#include <stdio.h>
#include <string.h>
#include <math.h>

#include "bsfunc.h"
#include "bscode.h"

#define pow10(y)	(pow(10, (y)))

#define PROG_NAME	"nsfo"
#define PROG_VERSION	"1.00"
#define AUTHOR_NAME	"Jiri Wackermann"
#define SYSERRMSG	strerror( errno )
#define LINEMAX		2048
#define OUT_SIGMA	0x001
#define OUT_PHI		0x002
#define OUT_OMEGA	0x004
#define LOG_SIGMA	0x010
#define LOG_PHI		0x020
#define LOG_OMEGA	0x040
#define OUT_MATRIX	0x100
#define OUT_TRIANG	0x200
#define OUT_DEFAULT	( OUT_SIGMA | OUT_PHI | OUT_OMEGA )

/* function prototypes */
extern void error( int, ... )		__attribute__((noreturn));
static void print_help( void )		__attribute__((noreturn));
static void print_version( void )	__attribute__((noreturn));
static void read_args( int, char ** );
static FILE *open_file( char *, char *, FILE * );
static int  read_data( FILE *, int, int, double[][], char ** );
static void select_data( int, int, int, double[][], double[][], char s[] );
static void print_matrix( FILE *, char *, int, int, int, double[][] );
static void compute_sf( int, int, double[][], double, double[] );
static double entropy( int, double[] );

static void v_copy( int, double[], double[] );
static void v_copy_select( int, double[], double[], char[] );
static void v_sub( int, double[], double[] ); 
static void v_divide( int, double[], double );
static double v_sum_line( int, double[] );
static double v_sumsq_line( int, double[] );
static void center_rows( int, int, double[][] );
static void center_columns( int, int, double[][] );
static void covariance( int, int, double[][], double[][], int );
static void cov_to_cor( int, double[][], int );
static double m_trace( int, double[][] );
static void jacobi( int, double[][], double[], double[][], double );

/* global data */

static char *helpinfo = 
  "usage:\n"
  "  %s [options] [infile [outfile]]\n"
  "options:\n"
  "  -1        compute Omega only\n"
  "  -2        compute Sigma and Phi only\n"
  "  -3        compute all three descriptors\n"
  "  -c flags  center in space `s' and/or in time `t'\n"
  "  -C        print cov/cor matrix, full\n"
  "  -f freq   sampling frequency in samples/time unit\n"
  "  -h n      skip first n lines\n"
  "  -l        logarithmise Omega\n"
  "  -L        logarithmise all descriptors\n"
  "  -N n,k    data matrix size n samples by k channels\n"
  "  -n        suppress terminating LF\n"
  "  -o fmt    numeric format string (default: `%%le' or $OFMT)\n"
  "  -p n      precision 10^(-n)\n"
  "  -R        use correlation matrix\n"
  "  -s mask   select channels by mask\n"
  "  -T        print cov/cor matrix, triangular\n"
  "  -v        print run-time messages to stderr\n"
  "  -x k      exclude channel k\n"
  "  -X        exclude all channels in sequence (implies -1)\n"
  "  -!        print version info and exit\n"
  ;

char *progname = NULL;
static char *ifname	= "-";		/* input file name, default = stdin  */
static char *ofname	= "-";		/* output file name,default = stdout */
static char *outfmt	= "%le";	/* numeric output format string */
static int   n_chin	= 0;		/* number of channels on input */
static int   n_chan	= 0;		/* number of channels to process */
static int   n_samp	= 0;		/* number of samples */
static double f_samp	= 0.;		/* sampling frequency */
static int   center_t	= 0;		/* center in time */
static int   center_s	= 0;		/* center in space */
static int   ln_skip	= 0;		/* lines to skip */
static int   no_lf	= 0;		/* suppress terminating '\n' */
static int   out_mask	= OUT_DEFAULT;	/* output select bitmask */
static int   trans_cor	= 0;		/* transform to correlation matrix */
static char *select_p	= 0;		/* ptr to channels selector string */
static int   verbose	= 0;		/* print run-time info */
static double epsilon	= 1e-6;		/* precision for jacobi() */
static int   excl_ch	= 0;		/* excluded channel index */
static char *mask_ch	= 0;		/* channel selection bitstring */

/* code starts here */

void main( int argc, char **argv ) {
  FILE   *ifile, *ofile;
  double (*data_p)[][], (*src_p)[][], (*cov_p)[][], sfo[3];
  char   *mat, *msg;

  read_args( argc, argv );

		#ifdef SHOWPARAMS
		fprintf( stderr, "ifname = `%s'\n", ifname );
		fprintf( stderr, "ofname = `%s'\n", ofname );
		fprintf( stderr, "outfmt = `%s'\n", outfmt );
		fprintf( stderr, "n_chin = %d\n", n_chin );
		fprintf( stderr, "n_chan = %d\n", n_chan );
		fprintf( stderr, "n_samp = %d\n", n_samp );
		fprintf( stderr, "f_samp = %lf\n", f_samp );
		fprintf( stderr, "center_t = %d\n", center_t );
		fprintf( stderr, "center_s = %d\n", center_s );
		fprintf( stderr, "out_mask = %06x\n", out_mask );
		fprintf( stderr, "select_p = `%s'\n", select_p );
		fprintf( stderr, "trans_cor = %d\n", trans_cor );
		fprintf( stderr, "verbose = %d\n", verbose );
		fprintf( stderr, "epsilon = %lf\n", epsilon );
		fprintf( stderr, "excl_ch = %d\n", excl_ch );
		#endif

	/* open input/output files */
  ifile = open_file( ifname, "rt", stdin );
  ofile = open_file( ofname, "wt", stdout );

	/* allocate space for data buffer */
  src_p = calloc( n_samp * n_chin, sizeof( double ));
  if ( !src_p )
    error( 1, "data buffer", SYSERRMSG, 0 );
  if ( n_chan == n_chin )
    data_p = src_p;
  else {
    data_p = calloc( n_samp * n_chan, sizeof( double ));
    if ( !src_p )
      error( 1, "data buffer", SYSERRMSG, 0 );
  }

	/* allocate space for covariance matrix */
  if ( out_mask & OUT_OMEGA ) {
    cov_p = calloc( n_chan * n_chan, sizeof( double ));
    if ( !cov_p )
      error( 1, SYSERRMSG, 0 );
  }

  if ( ln_skip ) {
    if ( !read_data( ifile, ln_skip, 0, NULL, &msg ))
      error( 1, msg, 0 );
  }
	/* read & prepare data */
  if ( verbose ) {
    fprintf( stderr, "%s: %s -> %s\n", progname, ifname, ofname );
    fprintf( stderr, "  reading data: %dx%d\n", n_samp, n_chan );
    mat = trans_cor ? "correlation" : "covariance";
  }
  if ( !read_data( ifile, n_samp, n_chin, *src_p, &msg ))
    error( 1, msg, 0 );
  fclose( ifile );

  if ( n_chan != n_chin )
    select_data( n_samp, n_chin, n_chan, *src_p, *data_p, mask_ch );

  if ( center_t || center_s ) {
    if ( verbose ) fprintf( stderr, "  centering data: " );
    if ( center_t ) {
      if ( verbose ) fprintf( stderr, "t" );
      center_rows( n_samp, n_chan, *data_p );
    }
    if ( center_s ) {
      if ( verbose ) fprintf( stderr, "s" );
      center_columns( n_samp, n_chan, *data_p );
    }
    if ( verbose ) fputc( '\n', stderr );
  }

	/* compute Sigma/Phi */
  if ( out_mask & ( OUT_SIGMA | OUT_PHI )) {
    if ( verbose )
      fprintf( stderr, "  computing integral descriptors\n" );
    compute_sf( n_samp, n_chan, *data_p, f_samp, sfo );
  }

	/* compute Omega */
  if ( out_mask & OUT_OMEGA ) {
    double T, lambda[n_chan];

    if ( verbose )
      fprintf( stderr, "  computing %s matrix: %dx%d\n", mat, n_chan, n_chan );
    covariance( n_samp, n_chan, *data_p, *cov_p, 0 );
    if ( trans_cor )
      cov_to_cor( n_chan, *cov_p, 0 );
    if ( out_mask & OUT_MATRIX ) {
      if ( verbose )
        fprintf( stderr, "  writing %s matrix\n", mat );
      print_matrix( ofile, outfmt, ( out_mask & OUT_TRIANG ), n_chan, n_chan,
                    *cov_p );
    }
    T = m_trace( n_chan, *cov_p );
    if ( verbose )
      fprintf( stderr, "  computing eigenvalues\n" );
    jacobi( n_chan, *cov_p, lambda, NULL, T * epsilon );
    T = v_sum_line( n_chan, lambda );
    v_divide( n_chan, lambda, T );
    sfo[2] = entropy( n_chan, lambda );
  }

	/* log transform */
  if ( out_mask & LOG_SIGMA )
    sfo[0] = log10( sfo[0] );
  if ( out_mask & LOG_PHI )
    sfo[1] = log10( sfo[1] );
  if ( out_mask & LOG_OMEGA )
    sfo[2] = sfo[2] / log( 10. );
  else
    sfo[2] = exp( sfo[2] );

	/* write results */
  if ( out_mask & OUT_SIGMA ) {
    fputc( ' ', ofile );
    fprintf( ofile, outfmt, sfo[0] );
  }
  if ( out_mask & OUT_PHI ) {
    fputc( ' ', ofile );
    fprintf( ofile, outfmt, sfo[1] );
  }
  if ( out_mask & OUT_OMEGA ) {
    fputc( ' ', ofile );
    fprintf( ofile, outfmt, sfo[2] );
  }
  if ( !no_lf )
    fputc( '\n', ofile );

	/* cleanup */
  fclose( ofile );
  free( cov_p );
  free( data_p );
  exit( 0 );
}

void print_help() {
  fprintf( stderr, helpinfo, progname );
  exit( 0 );
}

void print_version() {
  fprintf( stderr, "This is `%s' (version %s) by %s\n",
           PROG_NAME, PROG_VERSION, AUTHOR_NAME );
  exit( 0 );
}

void read_args( int argc, char **argv ) {
  char *p;
  int   o;

  if ( p = getenv( "OFMT" ))
    outfmt = p;
  progname = argv[0];
  opterr = 0;

  while (( o = getopt( argc, argv, "123c:Cf:h:lLnN:o:p:Rs:Tx:vx:X!?" )) != -1 )
    switch ( o ) {
      case '1':
        out_mask &= ~OUT_DEFAULT;
        out_mask |=  OUT_OMEGA;
        break;
      case '2':
        out_mask &= ~OUT_DEFAULT;
        out_mask |=  OUT_SIGMA | OUT_PHI;
        break;
      case '3':
        out_mask &= ~OUT_DEFAULT;
        out_mask |=  OUT_SIGMA | OUT_PHI | OUT_OMEGA;
        break;
      case 'c':
        if ( index( optarg, 's' )) center_s = 1;
        if ( index( optarg, 't' )) center_t = 1;
        break;
      case 'C':
        out_mask |= OUT_MATRIX;
        break;
      case 'f':
        f_samp = strtod( optarg, NULL );
        break;
      case 'h':
        ln_skip = atoi( optarg );
        break;
      case 'l':
        out_mask |= LOG_OMEGA;
        break;
      case 'L':
        out_mask |= LOG_SIGMA | LOG_PHI | LOG_OMEGA;
        break;
      case 'n':
         no_lf = 1;
         break;
      case 'N':
        if ( sscanf( optarg, " %d,%d", &n_samp, &n_chin ) == 2 ) {
          if ( n_samp <= 1 ) error( 1, "invalid number of samples", 0 );
          if ( n_chin <= 1 ) error( 1, "invalid number of channels", 0 );
        } else
          error( 1, "invalid data size", optarg, 0 );
        break;
      case 'o':
        outfmt = optarg;
        break;
      case 'p': {
        int n = atoi( optarg );
        if ( n > 0 )
          epsilon = pow10( (double)( -n ));
        else
          error( 1, "invalid precision", 0 );
      } break;
      case 'R':
        trans_cor = 1;
        break;
      case 's':
        select_p = optarg;
        break;
      case 'T':
        out_mask = OUT_MATRIX | OUT_TRIANG;
        break;
      case 'v':
        verbose = 1;
        break;
      case 'x':
        excl_ch = atoi( optarg );
        break;
      case 'X':
        excl_ch = -1;
        break;
      case '!':
        print_version();
        break;
      case '?':
        if ( o == optopt )
          print_help();
        else
          error( 1, "illegal option", &optopt, 0 );
    }

  if ( optind < argc )
    ifname = argv[optind];
  if ( optind+1 < argc )
    ofname = argv[optind+1];

  if ( n_samp == 0 || n_chin == 0 )
    error( 1, "data size unknown", 0 );
  if (( out_mask & OUT_PHI ) && f_samp == 0. )
    error( 1, "sampling frequency unknown", 0 );
  if ( excl_ch && !( 1 <= excl_ch && excl_ch <= n_chin ))
    error( 1, "excluded channel out of range", 0 );
  if ( excl_ch && select_p )
    error( 1, "conflicting options: -sx\n", 0 );

  if ( select_p ) {
    mask_ch = bs_alloc( n_chin );
    if ( bs_decode( select_p, mask_ch, n_chin, 0 ) < 0 )
      error( 1, "decode error", select_p, 0 );
    n_chan = bs_bits( mask_ch, n_chin );
		#ifdef DEBUG
		fprintf( stderr, "\t\tn_chan=%d\n", n_chan );
		#endif
  } else
    n_chan = n_chin;
}

FILE *open_file( char *fn, char *mod, FILE *def ) {
  FILE *f = ( fn[0] == '-' ) ? def  : fopen( fn, mod );
  if ( !f )
    error( 1, fn, SYSERRMSG, 0 );
  return f;
}

int read_data( FILE *f, int m, int n, double data[m][n], char **msgp ) {
  int     i, j, ok = 0;
  char    line[LINEMAX] = "";
  char   *lp, *msg = NULL;

	/* read matrix data */
  for ( i = 0; i < m; i++ ) {
    if ( lp = fgets( line, sizeof line, f )) {
      if ( data )
        for ( j = 0; j < n; j++ )
          if ( *lp != '\n' )
            data[i][j] = strtod( lp, &lp );
          else {
            msg = "unexpected eol"; goto done;
          }
    } else {
      msg = "unexpected eof"; goto done;
    }
  }
  ok = 1;

done:
  if ( msgp ) *msgp = msg;
  return ok;
}

void select_data( int m, int n1, int n2, double a[m][n1], double b[m][n2],
                  char *s ) {
  int i;

  for ( i = 0; i < m; i++ )
    v_copy_select( n1, b[i], a[i], s );
}

void print_matrix( FILE *f, char *out, int tri, int m, int n, double mat[m][n] ) {
  int i, j;

  for ( i = 0; i < m; i++ ) {
    int l = tri ? i+1 : n;
    for ( j = 0; j < l; j++ ) {
      if ( j ) fputc( ' ', f );
      fprintf( f, out, mat[i][j] );
    }
    fputc( '\n', f );
  }
}

void compute_sf( int N, int K, double data[N][K], double f, double res[2] ) {
  int    n;
  double m[2], diff[K];

  for ( n = 0; n < N; n++ ) {
    m[0] += v_sumsq_line( K, data[n] );
    if ( n == 0 ) continue;
    v_copy( K, diff, data[n] );
    v_sub(  K, diff, data[n-1] );
    m[1] += v_sumsq_line( K, diff );
  }
  m[0] /= (double)( N );
  m[1] /= (double)( N-1 );
  res[0] = sqrt( m[0] / (double)( K ));
  res[1] = m[0] > 0. ? f * sqrt( m[1] / m[0] ) / ( M_PI * 2. ) : 0.;
}

double entropy( int n, double p[n] ) {
  int i;
  double e = 0;

  for ( i = 0; i < n; i++ )
    if ( p[i] > 0. )
      e -= p[i] * log( p[i] );
  return e;
}

/*-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-*/

void v_zero( int n, double v[] ) {
  while ( n-- )
    *v++ = 0.;
}

void v_copy( int n, double v[], double u[] ) {
  bcopy( u, v, n * sizeof( double )); 
}

void v_copy_select( int n, double v[], double u[], char bs[] ) {
  int   i, j, m = 1;
  char *p = bs;

  for ( i = j = 0; i < n; i++ ) {
    if ( *p & m ) v[j++] = u[i];
    m <<= 1;
    if ( m == 0x100 ) {
      p++; m = 1;
    }
  }
}

void v_add( int n, double v[], double u[] ) {
  while ( n-- > 0 )
    *v++ += *u++;
}

void v_sub( int n, double v[], double u[] ) {
  while ( n-- > 0 )
    *v++ -= *u++;
}

void v_divide( int n, double v[], double k ) {
  while ( n-- )
    *v++ /= k;
}

double v_sum_line( int n, double v[] ) {
  double s = 0.;
  while ( n-- )
    s += *v++;
  return s;
}

double v_sumsq_line( int n, double v[] ) {
  double s = 0.;
  while ( n-- ) {
    s += *v * *v; v++;
  }
  return s;
}

void v_sub_line( int n, double v[], double k ) {
  while ( n-- )
    *v++ -= k;
}

void m_zero( int m, int n, double mat[m][n] ) {
  bzero( mat, m * n * sizeof( double ));
} 

void m_unit( int n, double mat[n][n] ) {
  int i;

  m_zero( n, n, mat );
  for ( i = 0; i < n; i++ )
    mat[i][i] = 1.;
}

double m_trace( int n, double mat[n][n] ) {
  int i;
  double t = 0.;

  for ( i = 0; i < n; i++ )
    t += mat[i][i];
  return t;
}

void m_getdiag( int n, double mat[n][n], double v[n] ) {
  int i;

  for ( i = 0; i < n; i++ )
    v[i] = mat[i][i];
}

void m_divide( int m, int n, double mat[n][n], double k ) {
  int     i;

  for ( i = 0; i < m; i++ )
    v_divide( n, mat[i], k );
}

void add_m_prod( int n, double mat[n][n], double u[n] ) {
  int i, j;

  for   ( i = 0; i <  n; i++ )
    for ( j = 0; j <= i; j++ )
      mat[i][j] += u[i] * u[j];
}

void sum_rows( int m, int n, double mat[m][n], double s[n] ) {
  int i;

  v_zero( n, s );
  for ( i = 0; i < m; i++ )
    v_add( n, s, mat[i] );
}

void sub_rows( int m, int n, double mat[m][n], double v[n] ) {
  int i;

  for ( i = 0; i < m; i++ )
    v_sub( n, mat[i], v );
}

void center_rows( int m, int n, double mat[m][n] ) {
  double mean[n];

  sum_rows( m, n, mat, mean );
  v_divide( n, mean, (double)( m ));
  sub_rows( m, n, mat, mean );
}

void center_columns( int m, int n, double mat[m][n] ) {
  int    i;

  for ( i = 0; i < m; i++ )
    v_sub_line( n, mat[i], v_sum_line( n, mat[i] ) / (double)( n ));
}

void makesym( int n, double mat[n][n] ) {
  int i, j;

  for   ( i = 1; i < n; i++ )
    for ( j = 0; j < i; j++ )
      mat[j][i] = mat[i][j];
}

void covariance( int m, int n, double mat[m][n], double cov[n][n], int tri ) {
  int i;

  m_zero( n, n, cov );
  for ( i = 0; i < m; i++ )
    add_m_prod( n, cov, mat[i] );
  if ( !tri )
    makesym( n, cov );
  m_divide( n, n, cov, (double)( m ));
}

void cov_to_cor( int n, double mat[n][n], int tri ) {
  int  i, j;
  double  d;

  for   ( i = 1; i < n; i++ )
    for ( j = 0; j < i; j++ )
      if (( d = mat[i][i] * mat[j][j]) != 0. )
        mat[i][j] /= sqrt( d );
  for ( i = 0; i < n; i++ )
    mat[i][i] = 1.;
  if ( !tri )
    makesym( n, mat );
}

/* --- the following routines are called by jacobi() --- */

double max_non_diag( int n, double mat[n][n], int *ip, int *jp ) {
  double  a, m = -1.;
  int     i, j, I, J;

  for   ( i = 1; i < n; i++ )
    for ( j = 0; j < i; j++ )
      if ( m < ( a = fabs( mat[i][j] ))) {
        m = a; I = i; J = j;
      }
  *ip = I;
  *jp = J;
  return m;
}

void rotate_rows( int n, double mat[n][n], int i, int j, double a, double b ) {
  double x, y;
  int k;

  for ( k = 0; k < n; k++ ) {
    x = mat[i][k] *  a + mat[j][k]  * b;
    y = mat[i][k] * -b + mat[j][k]  * a;
    mat[i][k] = x;
    mat[j][k] = y;
  }
}

void rotate_cols( int n, double mat[n][n], int i, int j, double a, double b ) {
  double x, y;
  int k;

  for ( k = 0; k < n; k++ ) {
    x = mat[k][i] *  a + mat[k][j] * b;
    y = mat[k][i] * -b + mat[k][j] * a;
    mat[k][i] = x;
    mat[k][j] = y;
  }
}

void jacobi( int n, double mat[n][n], double l[n], double v[n][n], double eps ) {
  double m, d, r, c, s;
  int    i, j;

  do {
    m = max_non_diag( n, mat, &i, &j );
    d = mat[i][i] - mat[j][j];
    r = sqrt( d * d + 4. * mat[i][j] * mat[i][j] );
    c = sqrt( ( 1. + d / r ) / 2. );
    s = sqrt( ( 1. - d / r ) / 2. );
    if ( mat[i][j] < 0. ) s = -s;
    rotate_cols( n, mat, i, j, c, s );
    rotate_rows( n, mat, i, j, c, s );
    if ( v )
      rotate_rows( n, v, i, j, c, s );
  } while ( m > eps );

  m_getdiag( n, mat, l );
}

/* end of file */