#include "empca.h"
#include "amdmex.h"
#include "svd.h"

int verbose = 0 ;

/*  One observation per column. (Left) eigenvector matrix and
    eigenvalue vector are pre-allocated. Output is one eigenvector per
    column. Either or both eval and evec_r can be NULL if their values
    are not needed. */

/*  NOTE: assumes the input matrix is zero mean! */

void empca(dym *obs, int iters, dym *evec, dyv *eval, dym *evec_r) {
  int i, j ;

  bool eval_is_null = (eval == NULL) ;
  bool evec_r_is_null = (evec_r == NULL) ;

  int obs_dim = dym_rows(obs) ;
  int num_obs = dym_cols(obs) ;
  int k = dym_cols(evec) ;

  dym *evec_t = mk_dym(k, obs_dim) ;
  dym *evec_inv = mk_dym(k, k) ;
  dym *x = mk_dym(k, num_obs) ;
 
  dym *x_t = mk_dym(num_obs, k);

  /*  Stuff we need for orthogonalization */

  dym *x_evec = mk_dym(k, k) ;
  dyv *ssq_evec_r ;

  /* If we haven't been given dyvs and dyms for eigenvalues or right eigenvectors,
     allocate them locally, and free them before we go home.
  */

  if (eval_is_null)
    eval = mk_dyv(k) ;
  if (evec_r_is_null)
    evec_r = mk_dym(k, num_obs) ;

  /* Make sure our matrices have the right shape */

  if (dym_rows(obs) != dym_rows(evec)
      || dym_cols(evec) != dyv_size(eval)
      || k > dym_rows(evec)) {
    fprintf(stderr, "Error: size of evec matrix, eval vector, or number of components requested is not compatible with dimensions of obs matrix passed to empca.\n") ;
    my_breakpoint() ;
  }
  
  /* ///////////////////////////////////////// */
  /*  Randomize evec to start */

  for (i = 0; i < obs_dim; ++i)
    for (j = 0; j < k; ++j)
      dym_set(evec, i, j, gen_gauss()) ;

  if (verbose)
    fprintf_dym(stdout,"evec_initial",evec,"\n");

  /* ///////////////////////////////////////// */
  /*  Iterate */

  for (i = 0; i < iters; ++i) {
    if (verbose)
      printf("iteration %d\n", i) ;

    /*  x = inv(evec'*evec)*evec'*obs ; */

    dym_transpose(evec, evec_t) ;

    dym_mult(evec_t, evec, evec_inv) ;
    invert_dym(evec_inv, evec_inv) ;

    /* fprintf_dym(stdout,"evec_inv",evec_inv,"\n"); */

    dym_mult(evec_t, obs, x) ;
    dym_mult(evec_inv, x, x) ;

    /* fprintf_dym(stdout,"x",x,"\n"); */

    /*  evec = obs*x'*inv(x*x') ; */

    dym_transpose(x, x_t) ;
    dym_mult(x, x_t, evec_inv) ;
    invert_dym(evec_inv, evec_inv) ;

    dym_mult(x_t, evec_inv, x_t) ;
    dym_mult(obs, x_t, evec) ;

    if (verbose)
      fprintf_dym(stdout,"evec",evec,"\n");
  }

  /* ////////////////////////////// */
  /*  Now we've got to orthogonalize */

  /*  evec = orth(evec); */
  /*  [xevec,eval] = truepca(evec'*data); */
  /*  evec = evec*xevec; */

  /* fprintf_dym(stdout,"before_orth",evec,"\n"); */ 

  orth(evec,evec) ;

  /* fprintf_dym(stdout,"evec_orth",evec,"\n"); */

  dym_transpose(evec, evec_t) ;
  dym_mult(evec_t, obs, x) ;
  
  mk_pca(x, x_evec, eval) ;

  /* fprintf_dym(stdout,"x_evec",x_evec,"\n"); */

  dym_mult(evec, x_evec, evec) ;

  if (verbose)
    fprintf_dym(stdout,"final evec",evec,"\n");

  /* From here, we jump through a few hoops to get the eigenvalues and right (vs left)
     eigenvectors into the right shape.
  */
  dym_transpose(evec, evec_t) ;
  dym_mult(evec_t, obs, evec_r) ;

  /* Scale so evec_r's are normalized */

  ssq_evec_r = mk_sumsq_of_dym_cols(evec_r, NULL) ;

  for (i = 0; i < k; ++i)
    dyv_set(eval, i, sqrt(dyv_ref(ssq_evec_r, i))) ;
  
  dyv_reciprocal(eval, ssq_evec_r) ;
  dym_scale_rows(evec_r, ssq_evec_r, evec_r) ; 

  free_dyv(ssq_evec_r) ;

  /*  Clean up at the end */

  if (eval_is_null)
    free_dyv(eval) ;
  if (evec_r_is_null)
    free_dym(evec_r) ;

  free_dym(evec_t) ;
  free_dym(evec_inv) ;
  free_dym(x) ;
  free_dym(x_t) ;
  free_dym(x_evec) ;
}

/* ///////////////////////////////////////////////////////////////////// */
/* Make each row zero mean, returning a dyv containing the original
   mean observation so that future data can be transformed using the
   same transform.
*/

dyv *mk_rows_zero_mean(dym *obs) {
  int obs_dim = dym_rows(obs) ;
  int num_obs = dym_cols(obs) ;

  dyv *mean_vec = mk_dyv(obs_dim) ;

  int i, j ;

  for (i = 0; i < obs_dim; ++i) {
    double mean_val = 0 ;
    for (j = 0; j < num_obs; ++j)
      mean_val += dym_ref(obs, i, j) ;
    mean_val /= num_obs ;
    
    dyv_set(mean_vec, i, mean_val) ;

    for (j = 0; j < num_obs; ++j)
      dym_set(obs, i, j, dym_ref(obs, i, j) - mean_val) ;
  }
  return mean_vec ;
}

/* ////////////////////////////////////////////////////////////////////
   Make each row unit mean, returning a dyv containing the original
   stddev so that future data can be transformed to the same scale by
   dividing by the returned vector.

   NOTE: this doesn't assume zero-mean data, but should probably be
   called after mk_rows_zero_mean, so that the mean vector is
   recovered.
*/

dyv *mk_rows_unit_variance(dym *obs)
{
  int obs_dim = dym_rows(obs) ;
  int num_obs = dym_cols(obs) ;

  dyv *stddev_vec = mk_dyv(obs_dim) ;

  int i, j ;

  /* Compute mean of column squares */

  for (i = 0; i < obs_dim; ++i) {
    double mean_val = 0 ;
    double var_val = 0 ;
    double stddev ;

    for (j = 0; j < num_obs; ++j) {
      double val = dym_ref(obs, i, j) ;
      mean_val += val ;
      var_val += val * val ;
    }

    mean_val /= num_obs ;
    var_val /= num_obs ;

    /* variance is E[x^2] - E[x]^2 */

    var_val -= mean_val * mean_val ;

    /* To get unit variance, we actually divide by the stddev */

    stddev = sqrt(var_val) ;
    dyv_set(stddev_vec, i, stddev) ;

    if (stddev)
      for (j = 0; j < num_obs; ++j)
	dym_set(obs, i, j, dym_ref(obs, i, j) / stddev) ;
  }

  return stddev_vec ;
}

/* //////////////////////////////////////// */
/* Helper routines for empca */

void orth(dym *A, dym *A_orth) {
  /*
    [U,S,V] = svd(A,0);
    [m,n] = size(A);
    if m > 1, s = diag(S);
    elseif m == 1, s = S(1);
    else s = 0;
    end
    tol = max(m,n) * max(s) * eps;
    r = sum(s > tol);
    Q = U(:,1:r);
  */

  int num_rows = dym_rows(A) ;
  int num_cols = dym_cols(A) ;
  int k = (num_rows > num_cols) ? num_cols : num_rows ;

  dym *u = mk_dym(num_rows, k) ;
  dyv *s = mk_dyv(k) ;
  dym *v = mk_dym(k, num_cols) ;
  
  dym_svd(A, u, s, v) ;

  /*  Should figure out what real rank is, but... */

  copy_dym(u, A_orth) ;

  /*  And clean up */

  free_dym(u) ;
  free_dyv(s) ;
  free_dym(v) ;

}

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

void mk_pca(dym *data, dym *evecs, dyv *evals) {
  /*
    [xevec,eval] = truepca(evec'*data);

    [d,N]  = size(dataset);

    mm = mean(dataset')';
    dataset = dataset - mm*ones(1,N);

    cc = cov(dataset',1);
    [cvv,cdd] = eig(cc);
    [zz,ii] = sort(diag(cdd));
    ii = flipud(ii);
    evects = cvv(:,ii);
    cdd = diag(cdd);
    evals = cdd(ii);
  */

  int dim = dym_rows(data) ;
  int num_obs = dym_cols(data) ;

  dym *data_t = mk_dym(num_obs, dim) ;
  dym *cov ;
  dyv *mean;

  /* fprintf_dym(stdout,"pca data",data,"\n"); */

  dym_transpose(data, data_t) ;

  /* We need to divide by n, not n-1, so do a round of */
  /* post-processing after computing the covariance. */

  mean = mk_mean_of_dym_rows(data_t,NULL);
  cov = mk_covariance_of_dym_rows(data_t,mean,NULL) ;
  dym_scalar_mult(cov, (num_obs-1.0)/num_obs, cov) ;

  /* fprintf_dym(stdout,"pca cov",cov,"\n"); */

  sorted_eigens_of_spd_dym(cov, evecs, evals) ;

  dym_transpose(evecs, evecs) ;

  /* fprintf_dym(stdout,"pca evecs",evecs,"\n"); */
  /* fprintf_dyv(stdout,"pca evals",evals,"\n"); */
  
  free_dym(data_t) ;
  free_dym(cov) ;
  free_dyv(mean) ;
}
