#include <nagdmc.h>

/*
  handle_info() prints to screen information based on the value of the info
  parameter.
*/
int
handle_info(const char func[], int info);

/*
  imputed_values() is the function described in 'Explanatory Code' in 
  nagdmc_impute_em.pdf.
*/
void 
imputed_values(long nvar, double data[], long nrepl, long indexes[]);

int
main(void) {
    const char    file[] = {"continuous.dat"};
    long          rec1 = 0;
    long          nvar = 8;
    long          nrec = 200;
    long          dblk = 200;
    double       *data = 0;
    long          nxvar = 0;
    long         *xvar = 0;
    double       *wt = 0;
    double        mval = -999.0;
    long          maxit = 20;
    long          it;
    double        tol = 1.0e-3;
    long          nrepl = 0;
    long          nempty = 0;
    double       *mean = 0;
    double       *cov = 0;
    double        sumwts = 0.0;
    int           info = 0;
    long         *indexes = 0;
    double       *xbar = 0;
    double       *ss = 0;
    FILE         *fp = 0;
    long          i, j, p;
    
#define XVAR(I) (nxvar == 0 ? I : xvar[I])

    /*
      Number of parameters.
    */
    if (nxvar == 0)
        p = nvar;
    else
        p = nxvar;

    /*
      Memory allocation.
    */
    if (!(data = (double *)malloc(dblk*nvar * sizeof(double))) ||
        !(mean = (double *)malloc(p * sizeof(double))) ||
        !(cov = (double *)malloc((p*(p+1)/2) * sizeof(double))) ||
        !(xbar = (double *)malloc(p * sizeof(double))) ||
        !(ss = (double *)malloc((p*(p+1)/2) * sizeof(double)))) {
        if (data) free(data);
        if (mean) free(mean);
        if (cov) free(cov);
        if (xbar) free(xbar);
        printf(" Memory allocation failure.\n\n");
        return 2;
    }
    
    /*
      Read data values.
    */
    if ((fp = fopen(file,"r")) == 0) {
        printf(" Data file named %s was not found.\n\n",file);
        return 2;
    }
    
    for (i=0; i<dblk; ++i) {
        for (j=0; j<nvar; ++j) 
            fscanf(fp,"%lf ",&data[i*nvar+j]);
    }

    fclose(fp);
    
    /*
      Imputation by multivariate normal E-M.
    */
    indexes = nagdmc_impute_em(rec1,nvar,nrec,dblk,data,nxvar,xvar,wt,mval,
                               tol,maxit,&it,&nrepl,&nempty,mean,cov,&sumwts,
                               &info);

    if (handle_info("nagdmc_impute_em",info)) {
        free(data);
        free(mean);
        free(cov);
        free(xbar);
        free(ss);
        nagdmc_free_impute(indexes);
        return 2;
    }

    /*
      Explanatory code from function document.
    */
    imputed_values(nvar,data,nrepl,indexes);

    /*
      Return allocated memory to operating system.
    */
    free(data);
    free(mean);
    free(cov);
    free(xbar);
    free(ss);
    nagdmc_free_impute(indexes);
    
    return 0;
}

void
imputed_values(long nvar, double data[], long nrepl, long indexes[]) {
    long i, j, k;

#define MISSING_ROW(I) indexes[I]
#define MISSING_COL(I) indexes[I+nrepl]
#define DATA(I,J) data[(I)*nvar+J]

    printf("\n\tRow \tCol \tValue\n");
    for (i=0; i<nrepl; ++i) {
        j = MISSING_ROW(i);
        k = MISSING_COL(i);
        printf("\t%-4li\t%-4li\t%-8.4f\n",j,k,DATA(j,k));
    }
}

int
handle_info(const char func[], int info) {
    if (info == -999)
    {
        printf(" Invalid licence, please contact NAG.\n\n");
        return 2;
    }
    else if (info > 0)
    {
        printf(" Error code %i from %s.\n\n",info,func);
        return 1;
    }
    else if (info < 0)
        printf (" Information code %i from %s.\n\n",info,func);

    return 0;
}