#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);

int
main(void) {
    const char    data_file[] = {"iris.dat"};
    const char    save_file[] = {"knnc_save"};
    long          rec1 = 0;
    long          nvar = 5;
    long          nrec = 130;
    long          dblk = 150;
    double       *data = 0;
    long          nxvar = 4;
    long          xvar[] = {0,1,2,3};
    long          iproot = 0;
    long          yvar = 4;
    long          ng = 3;
    long          nig[] = {44,43,43};
    double        prior[] = {0.34, 0.33, 0.33};
    double        rho = 0.5;
    long          uc = -555;
    long          norm = 2;
    long          k = 3;
    long         *res = 0;
    long         *nn = 0;
    double       *dist = 0;
    int           info = 0;
    long          i, j;
    FILE         *fp = 0; 

    /*
      Allocate memory for data and read data values.
    */
    if (!(data = (double *)malloc(dblk*nvar * sizeof(double)))) {
        printf (" Memory allocation failure.\n\n");
        return 2;
    }

    if ((fp = fopen(data_file,"r")) == 0) {
        printf("\n Data file named %s not found. \n\n",data_file);
        return 2;
    }
    
    for (i=0; i<dblk; ++i) {
        for (j=0; j<nvar; ++j) 
            fscanf(fp,"%lf ",&data[i*nvar+j]);
        
        data[i*nvar+yvar] += 1.0;
    }
    
    fclose(fp);

    /*
      Compute binary search tree using training data.
    */
    nagdmc_kdtree(rec1,nvar,nrec,dblk,data,nxvar,xvar,yvar,ng,nig,&iproot,&info);

    if (handle_info("nagdmc_kdtree",info)) {
        free(data);
        return 2;
    }

    /*
      The exmaple of saving and re-loading a tree lattice now follows.
    */

    /*
      Save computed binary tree.
    */
    nagdmc_save_kdtree(iproot,save_file,&info);

    if (handle_info("nagdmc_save_kdtree",info)) {
        free(data);
        return 2;
    }
    
    /*
      Return the memory used in the binary tree to the operating system.
    */
    nagdmc_free_kdtree(iproot);

    iproot = 0;

    /*
      Load the saved binary tree into memory.
    */
    nagdmc_load_kdtree(save_file,&iproot,&info);

    if (handle_info("nagdmc_load_kdtree",info)) {
        free(data);
        return 2;
    }

    /*
      Memory allocation for 20 k-nearest neighbours classifications.
    */
    if (!(res = (long *)malloc(20 * sizeof(long)))) {
        printf(" Memory allocation failure.\n\n");
        free(data);
        return 2;
    }
    
    if (!(nn = (long *)malloc(20*k * sizeof(long)))) {
        printf (" Memory allocation failure.\n\n");
        free(data);
        free(res);
        return 2;
    }
    
    if (!(dist = (double *)malloc(20*k * sizeof(double)))) {
        printf (" Memory allocation failure.\n\n");
        free(data);
        free(res);
        free(nn);
        return 2;
    }

    /*
      Calculate k-nearest neighbour classifications.
    */
    nagdmc_knnc(130,nvar,20,dblk,data,iproot,prior,rho,uc,norm,k,res,nn,
                dist,&info);
    
    if (handle_info("nagdmc_knnc",info)) {
        free(data);
        free(res);
        free(nn);
        free(dist);
        return 2;
    }

    /*
      Print classification results.
    */
    printf("\n %li-Nearest Neighbour model:\n\n",k);
    printf(" Record No.\tPrediction");
    for (j=0; j<k; ++j)
        printf("\tNN-%li Index (Dist)",j+1);
    
    for (i=0; i<20; ++i) {
        printf("\n %-10li\t%-10li",rec1+i,res[i]);

        for (j=0; j<k; ++j)
            printf("\t%-10li (%1.1e)",nn[i*k+j],dist[i*k+j]);
    }
    printf("\n\n");

    /*
      Return allocated memory to the operating system.
    */
    free(data);
    free(res);
    free(nn);
    free(dist);
    if (iproot)
        nagdmc_free_kdtree(iproot);

    return 0;
}

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

    return 0;
}