#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    datafile[] = {"hald.dat"};
    const char    sfile[] = {"hald_stepwise_reg_save"};
    long          rec1 = 0;
    long          nvar = 5;
    long          nrec = 13;
    long          dblk = 13;
    double       *data = 0;
    long          nxvar = 0;
    long         *xvar = 0;
    long          yvar = 4;
    long          iwts = -1;
    int           xtype[] = {0,0,0,0};
    double        Fin = 2.0;
    double        Fout = 2.0;
    double        eps = 1.0e-6;
    double       *b = 0;
    double       *se = 0;
    double        rms = 0.0;
    double        R2 = 0.0;
    long          df = 0;
    double       *model = 0;
    int           printlevel = 1;
    const char   *file = 0;
    int           info = 0;
    long          p;
    double        fv = 0.0;
    double        res = 0.0;
    double        h = 0.0;
    double        pred = 0.0;
    double        se_pred = 0.0;
    FILE         *fp = 0;
    long          i, j, l;


    /*
      Maximum number of parameters in the model.
    */
    if (nxvar > 0)
        p = nxvar + 1;
    else
        p = nvar - (iwts>=0);
    
    /*
      Memory allocation.
    */
    l = (3 * p * (p + 1)) / 2 + nvar + 14;
    if (!(b = (double *)malloc(p * sizeof(double))) ||
        !(se = (double *)malloc(p * sizeof(double))) ||
        !(model = (double *)malloc(l * sizeof(double))) ||
        !(data = (double *)malloc(dblk*nvar * sizeof(double)))) {
        printf (" Memory allocation failure.\n\n");
        return 2;
    }

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

    /*
      Compute stepwise regression.
    */
    nagdmc_stepwise_reg(rec1,nvar,nrec,dblk,data,nxvar,xvar,yvar,iwts,xtype,
                        Fin,Fout,eps,b,se,&R2,&rms,&df,model,printlevel,file,
                        &info);
    
    if (handle_info("nagdmc_stepwise_reg",info)) {
        free(data);
        free(model);
        free(b);
        free(se);
        return 2;
    }

    /*
      Example of saving a linear regression model to file and re-loading it.
    */  
    nagdmc_save_reg(sfile,model,&info);
    if (handle_info("nagdmc_save_reg",info)) {
        free(data);
        free(b);
        free(se);
        free(model);
        return 2;
    }

    free(model); model = 0;

    model = nagdmc_load_reg(sfile,&info);

    if (handle_info("nagdmc_load_reg",info)) {
        free(data);
        free(b);
        free(se);
        return 2;
    }

    /*
      Extract and print information from the fitted model.
    */
    printf(" %-8s  %-12s  %-8s  %-8s\n","Record","Fitted value", "Residual",
           "Leverage");

    for (i=0; i<nrec; ++i) {
        j = rec1 + i;
        
        nagdmc_extr_reg(model,&data[j*nvar],0,&fv,&res,&h,&info);

        if (handle_info("nagdmc_extr_reg",info)) {
            free(data);
            free(b);
            free(se);
            nagdmc_free_reg(model);
            return 2;
        }

        printf(" %-8li  %-12.4f  %-8.4f  %-8.4f\n",j,fv,res,h);
    }

    /*
      Predict a value for the last record in the data.
    */
    printf("\n Predictions\n\n");
    printf(" %-8s  %-15s  %-18s\n","Record","Predicted value",
           "S.E. Predicted Value");

    j = dblk-1;

    nagdmc_predict_reg(model,&data[j*nvar],0,0,&pred,&se_pred,1,&info);

    if (handle_info("nagdmc_predict_reg",info)) {
        free(data);
        free(b);
        free(se);
        nagdmc_free_reg(model);
        return 2;
    }
    
    printf(" %-8li  %-15.4f  %-18.4f\n",j,pred,se_pred);
    
    /*
      Return allocated memory to operating system.
    */
    free(data);
    free(b);
    free(se);
    nagdmc_free_reg(model);

    return 0;
}

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;
}