/*
Jason Lawrence
Princeton University

acls.cc

Several variants of the Alternating Constrained Least Squares (ACLS)
algorithm described in:

Lawrence, J., Ben-Artzi, A., DeCoro, C., Matusik, W., Pfister, H.,
Ramamoorthi, R., Rusinkiewicz, S. Inverse Shade Trees for
Non-Parametric Material Representation and Editing. ACM Transactions
on Graphics (Proceedings of ACM SIGGRAPH 2006).

This code calls several functions provided by the Numerical Algorithms
Group (NAG) library set and Intel's Math Kernel Library.

History:
   4/2006 - Created (jason)
*/

#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <signal.h>

// NAG library headers
#include <nag.h>
#include <nag_types.h>
#include <nage04.h>

// Intel's Math Kernel Library CBLAS headers
#include <mkl_cblas.h>

#include "acls.h"
#include "timer.h"

extern float *H_solidangles;
extern bool coltex;

void show_brdf (double *H, int th, int td, int pd);

// Allocate memory for ACLS workspace object.
acls_ws_t * create_acls_ws (int M,
			    int N,
			    int max_k)
{
  acls_ws_t *ws = (acls_ws_t *)malloc(sizeof(acls_ws_t));
  if (!ws) return 0;

  ws->max_k = max_k;

  ws->Wt   = (double *)malloc(sizeof(double)*M*(max_k+1));
  ws->Ht   = (double *)malloc(sizeof(double)*(N+2)*(max_k+1));
  ws->b    = (double *)malloc(sizeof(double)*(MAX(M,N)+2));
  ws->x    = (double *)malloc(sizeof(double)*max_k*(max_k+1));
  ws->kx   = (long *)malloc(sizeof(long)*(max_k+1)); bzero(ws->kx,sizeof(long)*(max_k+1));
  ws->bl   = (double *)malloc(sizeof(double)*(max_k+2));
  ws->bu   = (double *)malloc(sizeof(double)*(max_k+2));
  ws->cvec = (double *)malloc(sizeof(double)*(max_k+1));
  ws->sums = (double *)malloc(sizeof(double)*max_k);
  ws->objs = (double *)malloc(sizeof(double)*max_k);
  ws->inc  = (int *)malloc(sizeof(int)*M);
  ws->WH   = (float  *)malloc(sizeof(float)*M*N);
  ws->A    = (double *)malloc(sizeof(double)*(max_k+1));
  if (!ws->WH) {
    fprintf(stderr, "Out of memory in create_acls_ws.\n");
    exit(1);
  }

  return ws;
}

// Free memory associated with ACLS workspace object.
void free_acls_ws(acls_ws_t *ws)
{
  delete [] ws->Wt;
  delete [] ws->Ht;
  delete [] ws->b;
  delete [] ws->x;
  delete [] ws->kx;
  delete [] ws->bl;
  delete [] ws->bu;
  delete [] ws->cvec;
  delete [] ws->sums;
  delete [] ws->objs;
  delete [] ws->inc;
  delete [] ws->WH;
  delete [] ws->A;
}

// Compute error metric proposed in Equation 4.
double spunerrf(float *V,
		float *B,
		int M,
		int N,
		float *W,
		float *H,
		int K,
		float *WH,
		float lambda,
		float mu,
		int *inc)
{
  cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, M, N, K, 1.0, W, K, H, N, 0.0, WH, N);
  float *_V=V;
  float *_B=B;
  float *_WH=WH;

  float fit = 0.0;
  for(int i=0; i<M*N; i++)
    {
      float diff = (B ? *_B++ * (*_V++ - *_WH++) : (*_V++ - *_WH++));
      fit += diff*diff;
    }

  float sp=0.0;
  for (int i=0; i<M; i++)
    {
      for (int j=0; j<K; j++) {
	if (inc[i]==j) continue;
	float diff = 0 - W[i*K+j];
	sp += lambda*diff*diff;
      }
    }
  
  float un=0.0;
  for (int i=0; i<M; i++)
    {
      float sum=0.0;
      for (int j=0; j<K; j++)
	sum += W[i*K+j];
      float diff = 1.0-sum;
      un += mu*diff*diff;
    }

  return (fit+sp+un);
}

void cls_W (float *V,                       // Data matrix
	    float *B,                       // Confidence matrix
	    float *R,                       // Residual matrix (ignored)
	    int M, int N,                   // Problem resolution
	    float *W, float *H,             // Current factorization
	    int K,                          // Desired term count
	    acls_ws_t *ws,                  // ACLS Workspace Object
	    float lambda,                   // Sparsity parameter
	    float mu,                       // Unity norm parameter
	    int *inc,                       // The 'hidden' variables (strict association between observations in input set and basis vectors)
	    std::vector<double> *lbound,    // General value boundary constraints
	    std::vector<double> *ubound)
{
  static Nag_E04_Opt *options = 0;
  static NagError *fail = 0;
  static int nag_init_k = 0;
  static int nag_init_m = 0;
  static int nag_init_n = 0;

  if (options && (K!=nag_init_k || M!=nag_init_m || N!=nag_init_n))
    {
      e04xzc(options,"all",fail);
      delete options; options = 0;
      delete fail; fail = 0;
    }      

  if (!options)
    {
      options = new Nag_E04_Opt;
      fail    = new NagError;
      
      e04xxc(options);
      options->prob=Nag_LS1;
      options->list = false;
      options->print_level = Nag_NoPrint;
      options->print_deriv = Nag_D_NoPrint;
      
      bzero(fail,sizeof(NagError));
      fail->print = FALSE;

      nag_init_k = K;
      nag_init_m = M;
      nag_init_n = N;
    }

  double *Ht    = ws->Ht;
  double *b     = ws->b;
  double *x     = ws->x;
  long   *kx    = ws->kx;
  double *bl    = ws->bl;
  double *bu    = ws->bu;

  for (int j=0; j<K; j++){
    if (!lbound || !ubound) {
      bl[j]=0;   // strictly non-negative by default
      bu[j]=1e20;
    } else {
      bl[j]=(*lbound)[j]; // load arbitrary bound constraints
      bu[j]=(*ubound)[j];
    }
  }

  // Update W (one row at a time)
  for (int i=0; i<M; i++){

    // Loop over dimensions (we will take the one that produces the
    // lowest error)
    for (int k=0; k<K; k++){

      // Build target vector (weight by values in confidence matrix)
      for (int j=0; j<N; j++)
	b[j] = double(V[i*N+j] * B[i*N+j]);

      b[N]   =0.0;        // sparse term
      b[N+1] =sqrt(mu);   // unity norm term

      // Build vector of unknowns (the elements of this row of W in our case)
      for (int j=0; j<K; j++)
	x[k*K+j]=W[i*K+j];

      // Buid observations matrix (the matrix H weighted by the
      // confidence values in our case)
      for (int j=0; j<N; j++)
	for (int l=0; l<K; l++)
	  Ht[j*K+l] =double(H[l*N+j] * B[i*N+j]);

      // Add in parameter that controls sparsity of solution
      for (int l=0; l<K; l++)
	Ht[N*K+l] =(l==k ? 0.0 : sqrt(lambda));

      for (int l=0; l<K; l++)
	Ht[(N+1)*K+l] =sqrt(mu);


      // Call NAG routine to solve constrained least squares problem
      e04ncc(N+2,K,
	     0,(double *)0,0,   // no general constraints
	     bl,bu,(double *)0,b,Ht,K,kx,
	     x+k*K,&(ws->objs[k]),options,NAGCOMM_NULL,fail);

      options->used = TRUE;

    } // for (int k=0; k<K; k++){

    // Find solution with lowest error
    inc[i]=0;
    for (int j=1; j<K; j++)
      if (ws->objs[j] < ws->objs[inc[i]])
	inc[i]=j;

    // Store new mixing weights
    for (int j=0; j<K; j++)
      W[i*K+j] = float(x[inc[i]*K+j]);
  }
}

void cls_H_full (float *V, float *B,
		 int M, int N,
		 float *W, float *H,
		 int K)
{
  static Nag_E04_Opt options;
  static NagError fail;

  e04xxc(&options);
  options.prob=Nag_LS1;
  options.list = false;
  options.print_level = Nag_NoPrint;
  options.print_deriv = Nag_D_NoPrint;
  
  bzero(&fail,sizeof(NagError));
  fail.print = TRUE;

  static double objf;

  unsigned long sz = (unsigned long)(M*N)*(unsigned long)(K*N);

  int nclin = K;

  printf("Constraint matrix will be %.2fMB\n", double(sz)/double(ONE_MB)); fflush(stdout);

  static double *bl   = new double[K*N+K*3];
  static double *bu   = new double[K*N+K*3];

  static double *cvec = new double[K*N];
  static double *cm   = new double[sz];
  static double *b    = new double[M*N];
  static long   *kx   = new long[K*N];
  static double *x    = new double[K*N];
  static double *a    = new double[K*3*K*N];

  if ( !bl || !bu || !cvec || !cm || !b || !kx || !x || !a) {
    fprintf(stderr, "Out of memory in [globalfitH].\n");
    exit(1);
  }

  double *_bl = bl; double *_bu = bu;
  for (int i=0; i<K*N; i++){
    *_bl++=0;
    *_bu++=1e20;
  }
  for (int i=0; i<K*3; i++){
    *_bl++=0;
    *_bu++=3.14159;
  }

  bzero(cvec,sizeof(double)*K*N);
  bzero(kx,sizeof(long)*K*N);

  // Build linear constraint matrix
  bzero(a,sizeof(double)*K*3*K*N);

  if (coltex) {
    for (int j=0; j<K; j++) {
      for (int i=0; i<N; i++){
	a[ j*K*N + j*N + i] = H_solidangles[i];
      }
    }
  } else {
    for (int j=0; j<K; j++) {
      for (int c=0; c<3; c++) {
	for (int i=0; i<N/3; i++){
	  a[ (j*3+c)*K*N + j*N + i*3 + c ] = H_solidangles[i];//1.0;
	}
      }
    }
  }

  // Create b (target vector)
  float *_V=V; float *_B=B;
  double *_b=b;
  for (int j=0; j<M; j++){
    for (int i=0; i<N; i++)
      //      *_b++= (*_V++) * (*_B++);
      *_b++= *_V++;
  }

  // Create x (vector of unknowns)
  double *_x = x;
  for (int j=0; j<K; j++)
    for (int i=0; i<N; i++)
      *_x++ = H[j*N+i];


  // Build fit portion of constraint matrix
  bzero(cm,sizeof(double)*(M*N)*(K*N));
  for (int j=0; j<M; j++){
    for (int i=0; i<N; i++){
      int x = j*N+i;
      for (int k=0; k<K; k++)
	cm[x*(K*N)+k*N+i]=W[j*K+k];
    }
  }
  
  // Find optimal solution
  if (coltex) {
    e04ncc(M*N,N*K,
	   K,a,K*N,
	   bl,bu,cvec,b,cm,K*N,kx,
	   x,&objf,&options,NAGCOMM_NULL,&fail);
  } else {
    e04ncc(M*N,N*K,
	   K*3,a,K*N,
	   bl,bu,cvec,b,cm,K*N,kx,
	   x,&objf,&options,NAGCOMM_NULL,&fail);
  }
  options.used = TRUE;

  // Store solution
  _x = x;
  for (int j=0; j<K; j++)
    for (int i=0; i<N; i++)
      H[j*N+i] = float(*_x++);

#if 0
  for (int i=0; i<K; i++) {
    double rl=0;
    double gl=0;
    double bl=0;
    for (int j=0; j<N/3; j++) {
      rl += H_solidangles[j]*H[i*N+j*3+0];
      gl += H_solidangles[j]*H[i*N+j*3+1];
      bl += H_solidangles[j]*H[i*N+j*3+2];
    }
    printf("K=%d:\n",i);
    printf("\trl=%f\n",rl);
    printf("\tgl=%f\n",gl);
    printf("\tbl=%f\n",bl);fflush(stdout);
  }  
#endif

}

void cls_H (float *V, // Data matrix
	    float *B, // Confidence matrix
	    float *R, // Residual matrix
	    int M, int N, // Problem resolution
	    float *W, float *H,  // Current factorization
	    int K, // Desired term count
	    acls_ws_t *ws,  // ACLS Workspace Object
	    float lambda) // Sparsity index
{
  static Nag_E04_Opt *options = 0;
  static NagError *fail = 0;
  static int nag_init_k = 0;
  static int nag_init_m = 0;
  static int nag_init_n = 0;

  if (options && (K!=nag_init_k || M!=nag_init_m || N!=nag_init_n))
    {
      e04xzc(options,"all",fail);
      delete options; options = 0;
      delete fail; fail = 0;
    }      

  if (!options)
    {
      options = new Nag_E04_Opt;
      fail    = new NagError;
      
      e04xxc(options);
      //  options->optim_tol=.0001;
      options->prob=Nag_LS1;
      options->list = false;
      options->print_level = Nag_NoPrint;
      options->print_deriv = Nag_D_NoPrint;
      //  options->max_iter = 1000;
      
      bzero(fail,sizeof(NagError));
      fail->print = FALSE;

      nag_init_k = K;
      nag_init_m = M;
      nag_init_n = N;
    }


  double *Wt   = ws->Wt;
  double *b    = ws->b;
  double *x    = ws->x;
  long   *kx   = ws->kx;
  double *bl   = ws->bl;
  double *bu   = ws->bu;

  for (int j=0; j<K; j++){
    bl[j]=0;
    bu[j]=1e20;
  }
  bl[K]=1.0;
  bu[K]=1.0;    

  // Update H, one column at a time
  for (int i=0; i<N; i++){
    
    for (int j=0; j<M; j++)
      b[j]=double(V[j*N+i] * B[j*N+i]);
    
    for (int j=0; j<K; j++)
      x[j]=H[j*N+i];
    x[K]=1.0;
    
    for (int j=0; j<M; j++) {
      for (int k=0; k<K; k++)
	Wt[j*(K+1)+k] = double(W[j*K+k] * B[j*N+i]);
      Wt[j*(K+1)+K] = (R ? double(R[j*N+i]) : 0);
    }

    // Find local minima of constrained linear lsq problem
    e04ncc(M,K+1,0,(double *)0,0,
	   bl,bu,(double *)0,b,Wt,K+1,kx,
	   x,&ws->objf,options,NAGCOMM_NULL,fail);
    options->used = TRUE;

    // Store new basis reflectance values
    for (int j=0; j<K; j++)
      H[j*N+i] = float(x[j]);
  }
}

void cls_H (float *V, // Data matrix
	    float *B, // Confidence matrix
	    int M, int N, // Problem resolution
	    float *W, float *H,  // Current factorization
	    int K, // Desired term count
	    acls_ws_t *ws)  // ACLS Workspace Object
{
  static Nag_E04_Opt *options = 0;
  static NagError *fail = 0;
  static int nag_init_k = 0;
  static int nag_init_m = 0;
  static int nag_init_n = 0;

  if (options && (K!=nag_init_k || M!=nag_init_m || N!=nag_init_n))
    {
      e04xzc(options,"all",fail);
      delete options; options = 0;
      delete fail; fail = 0;
    }      

  if (!options)
    {
      options = new Nag_E04_Opt;
      fail    = new NagError;
      
      e04xxc(options);
      //  options->optim_tol=.0001;
      options->prob=Nag_LS1;
      options->list = false;
      options->print_level = Nag_NoPrint;
      options->print_deriv = Nag_D_NoPrint;
      //  options->max_iter = 1000;
      
      bzero(fail,sizeof(NagError));
      fail->print = FALSE;

      nag_init_k = K;
      nag_init_m = M;
      nag_init_n = N;
    }


  double *Wt   = ws->Wt;
  double *b    = ws->b;
  double *x    = ws->x;
  long   *kx   = ws->kx;
  double *bl   = ws->bl;
  double *bu   = ws->bu;

  for (int j=0; j<K; j++){
    bl[j]=0;
    bu[j]=1e20;
  }

  // Update H, one column at a time
  for (int i=0; i<N; i++){
    
    for (int j=0; j<M; j++)
      b[j]=double(V[j*N+i] * B[j*N+i]);
    
    for (int j=0; j<K; j++)
      x[j]=H[j*N+i];
    
    for (int j=0; j<M; j++) {
      for (int k=0; k<K; k++)
	Wt[j*K+k] = double(W[j*K+k] * B[j*N+i]);
      //	Wt[j*K+k] = double(W[j*K+k]);
    }

    // Find local minima of constrained linear lsq problem
    e04ncc(M,K,0,(double *)0,0,
	   bl,bu,(double *)0,b,Wt,K,kx,
	   x,&ws->objf,options,NAGCOMM_NULL,fail);
    options->used = TRUE;

    // Store new basis reflectance values
    for (int j=0; j<K; j++)
      H[j*N+i] = float(x[j]);
  }
}

float acls (float *V,                                            // Data matrix
	    float *B,                                            // Confidence matrix
	    float *R,                                            // Residual matrix
	    int M, int N,                                        // Dimensions
	    float *W, float *H,                                  // Factorization
	    int K,                                               // Desired term count
	    acls_ws_t *ws,                                       // ACLS Workspace Object
	    float lambda,                                        // Sparsity parameter
	    float mu,                                            // Unity norm parameter
	    void (*cb)(int, float *, float *, int, int, int),    // Visualization callback
	    int min_iters,                                       // Iteration bounds
	    int max_iters,
	    std::vector<double> *t,                              // Vector of time stamps 
	    std::vector<double> *ssds,                           // Vector or error stamps
	    std::vector<double> *lbound,                         // Value bounds for W
	    std::vector<double> *ubound)
{
  if (!ws || K > ws->max_k)
    {
      fprintf(stderr, "Invalid K value for workspace in [acls].\n");
      exit(1);
    }

  int *inc = ws->inc; bzero(inc,sizeof(int)*M);

  //  double old_ssd = sperrf(V,B,M,N,W,H,K,ws->WH,lambda,inc);
  double old_ssd = spunerrf(V,B,M,N,W,H,K,ws->WH,lambda,mu,inc);

  printf("ACLS: initial SSD %.6e\n", old_ssd); fflush(stdout);

  if (t) t->push_back(0.0);
  if (ssds) ssds->push_back(old_ssd);

  double tel=0.0;
  Timer timer;
  timer.start();
  
  // Iterate
  for (int it=0;(max_iters==0 || it<max_iters);it++)
    {
      // Update W
      cls_W(V,B,R,M,N,W,H,K,ws,lambda,mu,inc,lbound,ubound);

      float tssd = spunerrf(V, B, M, N, W, H, K, ws->WH, lambda, mu, inc);
      printf("\r%d cls_W: %.6e ", it,tssd); 
      printf("\n");
      fflush(stdout);

      // Update H
#if 1
      // Update routine considers H one column at a time
      cls_H(V,B,R,M,N,W,H,K,ws,lambda);
#else
      // Update routine considers all H simultaneously (necessary for energy conservation constraints).
      // Takes much longer!
      cls_H_full(V,B,M,N,W,H,K);
#endif

      tssd = spunerrf(V, B, M, N, W, H, K, ws->WH, lambda, mu, inc);
      printf("\r%d cls_H: %.6e ", it,tssd);
      printf("\n");
      fflush(stdout);

      if (cb) cb(it,W,H,M,N,K);

      if ( it % 1 == 0 )
	{
	  double secs = timer.stop(); tel+=secs;

	  float ssd = spunerrf(V, B, M, N, W, H, K, ws->WH, lambda, mu, inc);
	  float ossd = sperrf(V, B, M, N, W, H, K, ws->WH, 0.0, inc);

	  if (it > 0)
	    {
	      double secs_iter = secs / 1.0;
	      printf("\tsecs per iter ~ %.2f      ", secs_iter);
	    }	  
	  fflush(stdout);

	  if (t) t->push_back(tel);
	  if (ssds) ssds->push_back(ssd);

	  if ( (it > min_iters) && (ssd > old_ssd) )
	    {
	      fprintf(stderr, "\n\tError increased!");
	      old_ssd = ssd;
	    }

	  if ( (it > min_iters) && (fabs(1.0-ssd/old_ssd) < .01) )
	    {
	      fprintf(stdout, "\n\tACLS converged.");
	      old_ssd = ssd;
	      break;
	    }

	  old_ssd = ssd;
	  
	  timer.start();
	}
    }

  printf("\n"); fflush(stdout);

  return old_ssd;

}