#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    file[] = {"failures.dat"};
    const char    sfile[] = {"failures_poisson_reg_save"};
    long          rec1 = 0;
    long          nvar = 6;
    long          nrec = 48;
    long          dblk = 48;
    double       *data = 0;
    long          nxvar = 0;
    long         *xvar = 0;
    long          yvar = 5;
    long          iwts = -1;
    long          ioff = -1;
    char          link = 'l';
    double        a = -1.0;
    double        dev = 0.0;
    long          df = 0;
    double       *b = 0;
    double       *se = 0;
    double       *cov = 0;
    double       *model = 0;
    double        scale = 0.0;
    double        tol = 0.0;
    double        eps = 0.0;
    long          maxit = 0;
    double        fv = 0.0;
    double        h = 0.0;
    double        res = 0.0;
    double        pred = 0.0;
    double        se_pred = 0.0;
    int           info = 0;
    FILE         *fp = 0;
    long          i, j, p;

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

    /*
      Read data values.
    */
    if ((fp = fopen(file,"r")) == 0) {
        printf("\n Data file named %s 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);

    /*
      Compute poisson regression.
    */
    nagdmc_poisson_reg(rec1,nvar,nrec,dblk,data,0,0,0,nxvar,xvar,yvar,iwts,ioff,
                       link,a,&dev,&df,b,se,cov,model,scale,tol,eps,maxit,&info);

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

    /*
      Example of saving a poisson 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(cov);
        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);
        free(cov);
        return 2;
    }

    /*
      Print information on model parameters.
    */
    printf("\n Parameter\tCo-efficient\tStandard Error\n");
    for (i=0; i<p; ++i)
        printf(" %-8li\t%-8.4f\t%-8.4f\n",i,b[i],se[i]);

    /*
      Extract and print information from the fitted model.
    */
    printf("\n %-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);
            free(cov);
            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);
        free(cov);
        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);
    free(cov);
    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;
}