/*-----------------------------------------------------------------------------

Copyright (C) 2017

A. Ronald Gallant
Post Office Box 659
Chapel Hill NC 27514-0659
USA

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

-------------------------------------------------------------------------------

Function      varcov - compute mean and variance-covariance matrix for X

Syntax        #include "libscl.h"
              INTEGER varcov
                (const realmat& X, realmat& mu, realmat& V, REAL eps=EPS); 
              INTEGER varcov
                (const trealmat& X, realmat& mu, realmat& V, REAL eps=EPS); 

Prototype in  libscl.h

Description   Computes mean vector mu and variance-covariance matrix V for 
              an n by p input matrix X: on return mu is p by 1 and V is 
	      p by p. The divisor for mu is n and for V it is max(n-1,1).

Return value  Estimated rank of V.

Remark        If X is a trealmat, then results are for T(X). If X is null, 
              then mu and V are null and zero is returned. The computation
	      is broken up into subtotals within blocks to improve accuracy.

Functions     Library: (none)
called        libscl:  realmat

------------------------------------------------------------------------------*/

#include "libscl.h"

namespace {
  const INTEGER blocksize = 2048;
}


namespace scl {

  INTEGER varcov(const realmat& X, realmat& mu, realmat& V, REAL eps) 
  {
    if (X.size() == 0) {
      realmat null;
      mu = V = null;
      INTEGER rank = 0;
      //error("Error, varcov, X is a null matrix");
      return rank;
    }

    INTEGER n = X.nrow();
    INTEGER p = X.ncol();

    const INTEGER blocks = n/blocksize;
    const INTEGER remain = n - blocks*blocksize;

    REAL blkscale = 1.0/REAL(blocksize);
    REAL muscale = REAL(blocksize)/REAL(n);
    REAL sigscale = ( n - 1 > 0 ? REAL(blocksize)/REAL(n - 1) : muscale );

    mu.resize(p,1,0.0);

    if (remain > 0) {

      for (INTEGER j=1; j<=p; ++j) {
        for (INTEGER i=1; i<=remain; ++i) {
          mu[j] += X(i,j);
        }
      }

      for (INTEGER j=1; j<=p; ++j) mu[j] *= blkscale;
    }

    if (blocks > 0) {

      for (INTEGER b=1; b<=blocks; ++b) {

        realmat ave(p,1,0.0);

        for (INTEGER j=1; j<=p; ++j) {
          for (INTEGER i=1; i<=blocksize; ++i) {
            INTEGER idx = remain + blocksize*(b - 1) + i;
            ave[j] += X(idx,j);
          }
        }

        for (INTEGER j=1; j<=p; ++j) mu[j] += blkscale*ave[j];
      }
    }

    for (INTEGER j=1; j<=p; ++j) mu[j] *= muscale;


    V.resize(p,p,0.0);
       
    if (remain > 0) {

      for (INTEGER j=1; j<=p; ++j) {
        for (INTEGER k=j; k<=p; ++k) {
          for (INTEGER i=1; i<=remain; ++i) {
            V(j,k) += (X(i,j)-mu[j])*(X(i,k)-mu[k]);
          }
        }
      }

      for (INTEGER j=1; j<=p; ++j) {
        for (INTEGER k=j; k<=p; ++k) {
          V(j,k) *= blkscale;
        }
      }
    }

    if (blocks > 0) {

      for (INTEGER b=1; b<=blocks; ++b) {

        realmat ave(p,p,0.0);

        for (INTEGER j=1; j<=p; ++j) {
          for (INTEGER k=j; k<=p; ++k) {
            for (INTEGER i=1; i<=blocksize; ++i) {
              INTEGER idx = remain + blocksize*(b - 1) + i;
              ave(j,k) += (X(idx,j)-mu[j])*(X(idx,k)-mu[k]);
            }
          }
        }

        for (INTEGER j=1; j<=p; ++j) {
          for (INTEGER k=j; k<=p; ++k) {
            V(j,k) += blkscale*ave(j,k);
          }
        }

      }
    }
        
    for (INTEGER j=1; j<=p; ++j) {
      for (INTEGER k=j; k<=p; ++k) {
        V(j,k) *= sigscale;
        V(k,j) = V(j,k);
      }
    }

    realmat R;
    INTEGER rank = cholesky(V,R,eps);

    return rank;
  }

  INTEGER varcov(const trealmat& X, realmat& mu, realmat& V, REAL eps) 
  {
    if (X.size() == 0) {
      realmat null;
      mu = V = null;
      INTEGER rank = 0;
      //error("Error, varcov, X is a null matrix");
      return rank;
    }

    INTEGER n = X.nrow();
    INTEGER p = X.ncol();

    const INTEGER blocks = n/blocksize;
    const INTEGER remain = n - blocks*blocksize;

    REAL blkscale = 1.0/REAL(blocksize);
    REAL muscale = REAL(blocksize)/REAL(n);
    REAL sigscale = ( n - 1 > 0 ? REAL(blocksize)/REAL(n - 1) : muscale );

    mu.resize(p,1,0.0);

    if (remain > 0) {

      for (INTEGER j=1; j<=p; ++j) {
        for (INTEGER i=1; i<=remain; ++i) {
          mu[j] += X(i,j);
        }
      }

      for (INTEGER j=1; j<=p; ++j) mu[j] *= blkscale;
    }

    if (blocks > 0) {

      for (INTEGER b=1; b<=blocks; ++b) {

        realmat ave(p,1,0.0);

        for (INTEGER j=1; j<=p; ++j) {
          for (INTEGER i=1; i<=blocksize; ++i) {
            INTEGER idx = remain + blocksize*(b - 1) + i;
            ave[j] += X(idx,j);
          }
        }

        for (INTEGER j=1; j<=p; ++j) mu[j] += blkscale*ave[j];
      }
    }

    for (INTEGER j=1; j<=p; ++j) mu[j] *= muscale;


    V.resize(p,p,0.0);
       
    if (remain > 0) {

      for (INTEGER j=1; j<=p; ++j) {
        for (INTEGER k=j; k<=p; ++k) {
          for (INTEGER i=1; i<=remain; ++i) {
            V(j,k) += (X(i,j)-mu[j])*(X(i,k)-mu[k]);
          }
        }
      }

      for (INTEGER j=1; j<=p; ++j) {
        for (INTEGER k=j; k<=p; ++k) {
          V(j,k) *= blkscale;
        }
      }
    }

    if (blocks > 0) {

      for (INTEGER b=1; b<=blocks; ++b) {

        realmat ave(p,p,0.0);

        for (INTEGER j=1; j<=p; ++j) {
          for (INTEGER k=j; k<=p; ++k) {
            for (INTEGER i=1; i<=blocksize; ++i) {
              INTEGER idx = remain + blocksize*(b - 1) + i;
              ave(j,k) += (X(idx,j)-mu[j])*(X(idx,k)-mu[k]);
            }
          }
        }

        for (INTEGER j=1; j<=p; ++j) {
          for (INTEGER k=j; k<=p; ++k) {
            V(j,k) += blkscale*ave(j,k);
          }
        }

      }
    }
        
    for (INTEGER j=1; j<=p; ++j) {
      for (INTEGER k=j; k<=p; ++k) {
        V(j,k) *= sigscale;
        V(k,j) = V(j,k);
      }
    }

    realmat R;
    INTEGER rank = cholesky(V,R,eps);

    return rank;
  }

}
