/* 2 3 {rank,cr}=llsq(\\"type\\",xr,yr), linear least squares

Determines n coefficients c such that y ~= c0*f0(x) + c1*f1(x) + ...
where f0,f1,... are the n basis functions (e.g. 1, x, x*x, x*x*x, ... for type "poly").

Solves the equation U'c = y for c using SVD, where U is the n-by-m matrix created
using the basis functions evaluated at the m x values:

      f0(x0)  f0(x1)  ...
  U = f1(x0)  f1(x1)  ...
      ...

To implement another function type using different basis functions:

1) write a function like rf_llsq_poly() below to initialize the U matrix,
   and add the type name and function address to the struct llsq_function table.

2) add a function to rf/feval.c to evaluate the function type.

3) update the documentation to describe the function type.
*/

//----------------------------------------
//
// initialize U matrix using polynomial basis functions 1, x, x*x, x*x*x, ...
//
static int rf_llsq_poly( double **u, double *x, int m, int n)
{
 // row 0 of U is all 1's
 //
  for( int j = 0; j < m; ++j)
    u[0][j] = 1;

 // row 1 of U is the x values
 //
  if( n > 1)
    for( int j = 0; j < m; ++j)
      u[1][j] = x[j];

 // remaining rows of U are powers of the x values
 //
  for( int i = 2; i < n; ++i)
    for( int j = 0; j < m; ++j)
      u[i][j] = x[j]*u[i-1][j];

  return 0; // success
}

//----------------------------------------
//
// placeholder for type "data" - this routine is not invoked
//
static int rf_llsq_data( double **u, double *x, int m, int n)
{
  return 0;
}

//----------------------------------------

typedef int (*LLSQ_F)( double **u, double *x, int m, int n);

static struct llsq_function
{
  const char *type; LLSQ_F f;
}
llsq_functions[] =
{
  { "data",	rf_llsq_data },
  { "poly",	rf_llsq_poly },
//
// add other function types here...
//
  {      0,	0 }
};

//----------------------------------------

#define LLSQ_P dp

// fill in vector with cell values
//
static void rf_llsq_init( Data *d)
{
  *d->LLSQ_P++ = eval_cell( d->cp, &d->c);
}

// fill in coefficient range with result values
//
static void rf_llsq_coeff( Data *d)
{
  if( d->cp->n)
  {
    fprintf( out, "llsq: can't assign cell ");
    print_cell( &d->c, &cell_00);
    putc( '\n', out);
  }
  else
  { 
    d->cp->val = *d->LLSQ_P++;
    d->cp->state = EVALUATED;
  }
}

//----------------------------------------

double rf_llsq(const Node *e, const Cell *c)
{
  Data d;
  int m, n, rank = 0;
  double **u, **v, *s, *x, *y;
  LLSQ_F f;
  struct llsq_function *lf;

  Node *lval = Left(e);
  Node *cr = lval ? lval->next : NULL;

  Node *type = Right(e);
  Node *xr = type->next;
  Node *yr = xr->next;

  if( type->type != STRING)
  {
    fprintf( out, "llsq: first argument must be a string\n");
    return 0;
  }

  f = NULL;

  for( lf = llsq_functions; lf->type; ++lf)
    if( strcmp( lf->type, type->u.str) == 0) { f = lf->f; break; }

  if( f == NULL)
  {
    fprintf( out, "llsq: unknown function type \"%s\"\n", type->u.str);
    return 0;
  }

  if( check_range( cr, "llsq: second return value")
   || check_range( xr, "llsq: second argument")
   || check_range( yr, "llsq: third argument") )
    return 0;

  m = Rsize( yr);
  n = Rsize( cr);

  if( f == rf_llsq_data) // special case, no basis functions
  {
    if( Rsize( xr) != n*m)
    {
      fprintf( out, "llsq: incompatible range sizes\n");
      return 0;
    }
  }
  else
  {
    if( Rsize( xr) != m)
    {
      fprintf( out, "llsq: x,y ranges must be the same size\n");
      return 0;
    }
  }

  u = matalloc( n, m);
  v = matalloc( n, n);
  s = malloc( n*sizeof(double));
  y = malloc( m*sizeof(double));

 // x is used initially to store the m x values,
 // then after the svd it is used to store the n result coefficients,
 // so it has to be big enough to hold either of those
 //
  x = malloc( ((m>n)?m:n)*sizeof(double));

  if( !u || !v || !s || !y || !x)
  {
    fprintf( out, "llsq: memory allocation failed\n");
    exit(1);
  }

 // initialize y values
 //
  d.LLSQ_P = y;
  ss_traverse_range_adjust( yr, c, rf_llsq_init, &d, 0);

  if( debug)
    vecprint( "llsq y", "%8g", y, m);

 // initialize U matrix
 //
  if( f == rf_llsq_data) // special case, no basis functions
  {
    int dir_save = dir;

   // ZZ - svd requires the transpose of the data matrix
   //
    if( dir == BYROWS) dir = BYCOLS; else dir = BYROWS;

    d.LLSQ_P = u[0];
    ss_traverse_range_adjust( xr, c, rf_llsq_init, &d, 0);

    dir = dir_save;
  }
  else
  {
   // initialize x values
   //
    d.LLSQ_P = x;
    ss_traverse_range_adjust( xr, c, rf_llsq_init, &d, 0);

    if( debug)
      vecprint( "llsq x", "%8g", x, m);

    if( f( u, x, m, n)) goto cleanup;
  }

  if( debug)
    matprint( "llsq u, before svd", "%8g", u, n, m);

 // perform SVD, orthogonalize the rows of U
 //
  rank = svd( u, s, v, m, n, DBL_EPSILON);

  if( debug)
  {
    fprintf( out, "llsq rank = %d\n", rank);
    matprint( "llsq u, after svd", "%8g", u, n, m);
    vecprint( "llsq s", "%8g", s, n);
    matprint( "llsq v", "%8g", v, n, n);
  }

 // solve for LLSQ coefficients, store in x vector
 //
 // x = sum {  vj  (uj' b) / sj  }

  for( int i = 0; i < n; ++i)
    x[i] = 0;

  for( int j = 0; j < rank; ++j)
  {
    double t = 0;
    for( int i = 0; i < m; ++i)
      t += u[j][i] * y[i];
    t /= s[j];
    for( int i = 0; i < n; ++i)
      x[i] += v[j][i] * t;
  }

  if( debug)
    vecprint( "llsq c", "%8g", x, n);

 // assign coefficient result values
 //
  d.LLSQ_P = x;
  ss_traverse_range_adjust( cr, c, rf_llsq_coeff, &d, 0);

cleanup:
 
  matfree( u);
  matfree( v);
  free( s);
  free( y);
  free( x);

  if( lval)
    lval_helper( "llsq", lval, c, rank);

  return rank;
}