// version 1.0.0 14Feb2018
#include <stdlib.h>
#include <stdio.h>
#include "stplugin.h"
#include "mycalcsw.h"

// Note: Matrices are long vectors with row-major storage
//    The i,j element of an r x c matrix is 
//    the (i-1)*r + (j-1) element of the of the vector
//    under C-style zero-base indexing
//
//    Define preprocesor macros to facilitate readability

#define  M(i, j)   *(*mat + (i)*c + j)
#define  B(j)      *(bmat+j)
#define  E(j)      *(emat+j)
#define  V(i, j)   *(vmat + (i)*c + j)
#define  C(i, j)   *(cmat + (i)*c + j)

ST_retcode InitCmat(ST_double **mat, ST_int r, ST_int c)
{
    ST_int  i, j ;
    char    msg[80] ;

    *mat = (ST_double *) malloc((size_t) r*c*sizeof(ST_double)) ;
    if (*mat == NULL ) {
        sprintf(msg, "Insufficient memory\n") ;
        SF_error(msg) ;
        return( (ST_retcode) 909) ;
    }

    for(i=0; i<r; i++) {
        for(j=0; j<c; j++) {
            M(i, j) = (ST_double) 0 ;
        }
    }

    return((ST_int) 0) ;
}

ST_retcode MyAve(ST_double *bmat, ST_int first, ST_int last, ST_int nVars, 
    ST_int *nObs)
{
    ST_int     i, j ;
    ST_double  value ;
    ST_retcode rc ;
    char       msg[80] ;

    rc    = (ST_retcode) 0 ;
    *nObs = (ST_int) 0 ;
    for(i=first-1; i<last; i++) {                  
        if (SF_ifobs(i+1)) {
            ++(*nObs) ;
            for(j=0; j<nVars; j++) {
                rc = SF_vdata(j+1, i+1, &value); 
                if(rc>0) {
                    sprintf(msg, "Problem accessing Stata data\n") ;
                    SF_error(msg) ;
                    return(rc) ;
                }
                if (SF_is_missing(value)) {
                    sprintf(msg, "missing values encountered\n") ;
                    SF_error(msg) ;
                    return( (ST_retcode) 416 ) ;
                }
                B(j) += value ; 
            }
        }
    }

    DivideByScalar(bmat, (ST_int) 1, nVars, (ST_double) *nObs) ; 
    return(rc) ;
}

ST_retcode MyV(ST_double *bmat, ST_double *vmat, ST_int first, ST_int last,
    ST_int nVars, ST_int nObs) 
{
    ST_int     i, j, j2, c ;
    ST_double  *emat, value  ;
    char       msg[80] ;
    ST_retcode rc ;

// used in macros for matrices
    c     = nVars ;
    emat  = NULL;

    rc = InitCmat(&emat, 1, nVars ) ; 
    if (rc>0) {
        return( rc ) ;
    }

    for(i=first-1; i<last; i++) {
        if (SF_ifobs(i+1)) {
            for(j=0; j<nVars; j++) {
                rc = SF_vdata(j+1, i+1, &value);     
                if(rc>0) {
                    free(emat) ;
                    sprintf(msg, "Problem accessing Stata data\n") ;
                    SF_error(msg) ;
                    return(rc) ;
                }
                if (SF_is_missing(value)) {
                    free(emat) ;
                    sprintf(msg, "missing values encountered\n") ;
                    SF_error(msg) ;
                    return( (ST_retcode) 416 ) ;
                }
                E(j) = value - B(j) ;
            }
            for(j=0; j<nVars; j++) {
// matrix is symmetric: only fill in lower diag in loop over observations
                for(j2=0; j2<=j; j2++) {
                    V(j,j2) += (E(j))*(E(j2)) ;
                }
            }
        }
    }

// Fill in above the diagonal
    for(j=0; j<nVars; j++) {
        for(j2=j+1; j2<nVars; j2++) {
            V(j,j2) = V(j2,j) ;
        }
    }

    DivideByScalar(vmat, nVars, nVars, (ST_double) (nObs*(nObs-1)) )  ;

    free(emat) ;
    return(rc) ;
}

ST_retcode CopyCtoStataMatrix(ST_double *cmat, char *smat, ST_int r, ST_int c)
{
    ST_int      i, j; 
    ST_retcode  rc ;
    char        msg[80] ;

    rc = (ST_retcode) 0 ;
    for(i=0; i<r; i++) {
        for(j=0; j<c; j++) {
            rc = SF_mat_store(smat, (i+1) , (j+1), C(i,j) ) ;
            if(rc>0) {
                sprintf(msg, "cannot access Stata matrix %s\n", smat) ;
                SF_error(msg) ;
                return(rc) ;
            }
        }
    }
    return(rc) ;
}


// Replace each element in r x c matrix mat with that element 
// divided by val
void DivideByScalar(ST_double *vmat, ST_int r, ST_int c, ST_double val) 
{
    ST_int  i, j ;

    for(i=0; i<r; i++) {
        for(j=0; j<c; j++) {
            V(i,j) /= val ;  
        }
    }
}

#undef M
#undef B
#undef E
#undef V