/* Umbrella tests for a one-way layout
   Test whether treatment groups increase to a particular treatment group
     then decrease afterwards
   Two versions: Known peak and unknown peak */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
/* Prototypes for all functions in  statlib.c */
#include "statlib.h"

/* Compile with e.g.   */
/*    gcc -Wall -o Umbrella Umbrella.c statlib.c   */
/* with both statlib.h and statlib.c in the same directory as Umbrella.c */

/* Number of permutation simulations to estimate P-values */
/*   and the starting random-number seed. */
int nsims=10000;
unsigned long wseed=123456;

#define MAXK 10       /* Upper bound for number of treatment groups */
#define MAXN 50       /* Upper bound for number of observations per group */
#define MAXTOT (MAXK*MAXN)  /* Upper bound for total number of observations */

/* Two data sets */

/* Dataset I: */

/* Title and source */
const char *title_one="Five treatment groups with 10 observations each";
const char *source_one="(Simulated data)";
/* Descriptive names for treatment groups */
const char *treatnames_one[MAXK] =
   { "#1", "#2", "#3", "#4", "#5", "#6" };

/* Number of treatment groups and treatment-group sample sizes */
const int kk_one=5;
const int nn_one[MAXK] = { 10, 10, 10, 10, 10 };

/* Data with treatment groups as rows */
/* See comments in OneWayLayout.c for syntax of double arrays */
/*   and how to initialize them. */
const double xx_one[MAXK][MAXN] =
 { { 308,  308,   12,  206,  368,  264,   88,  240,  118,  264 },  
   { 126,  252,  262,  258,  136,  204,  170,  276,  314,  308 },  
   { 218,  364,  300,  322,   46,  180,  278,  454,   70,  188 },  
   { 358,  366,  394,  252,  430,  274,  216,  284,  350,  290 },  
   { 332,  358,  356,  178,  262,  158,  176,  278,  240,  304 } };

/* How to display the data in a printf() statement */
const char *dataformat_one="%3g";

/* Dataset II: */

/* Title and source */
const char *title_two="Fasting metabolic rates in deer by months";
const char *source_two="Hollander & Wolfe Table 6.8, p215";
const char *treatnames_two[MAXK] =
  { "Jan-Feb", "Mar-Apr", "May-Jun",
    "Jul-Aug", "Sep-Oct", "Nov-Dec" };

/* Number of treatment groups and treatment-group sample sizes */
const int kk_two=6;
const int nn_two[MAXK] = { 7, 3, 5, 4, 4, 3 };

/* Data with treatment groups as rows */
const double xx_two[MAXK][MAXN] =
 { { 36.0, 33.6, 26.9, 35.8, 30.1, 31.2, 35.3 },
   { 39.9, 29.1, 43.4 },
   { 44.6, 54.4, 48.2, 55.7, 50.0 },
   { 53.8, 53.9, 62.5, 46.6 },
   { 44.3, 34.1, 35.7, 35.6 },
   { 31.7, 22.1, 30.7 } };

/* How to display the data in a printf() statement */
const char *dataformat_two="%4.1f";



/* Carry out umbrella tests for one dataset */

void do_umbrella_tests (const char *title, const char *source,
  const char *treatnames[MAXK],  const double xx[MAXK][MAXN],
  const int nn[MAXK], const int kk, const int pp,
  const char *dataformat);


int main(int argc, char **argv)
 { printf (
   "\nUMBRELLA TESTS for a one-way layout\n"
   "Test whether treatment groups increase to a particular treatment group\n"
   "  then decrease afterwards\n"
   "Two versions: Known peak and unknown peak\n");
   /* Helpful messages on segmentation faults and some other fatal errors */
   /* This function is in statlib.c */
   set_newsignals();

   /* Seed random-number seed */
   /* Enter `ProgramName <Number>' to set starting RN seed=<Number> */
   if (argc>=2) wseed=atol(argv[1]);
   wseed=setrand(wseed);  /* If wseed=0, replace by true starting seed */
   printf ("\nStarting random-number seed: %lu\n",wseed);
   printf ("(Enter a number on the command line "
     "to specify a different seed.)\n");

   /* First dataset */
   printf ("\nFIRST DATA SET:\n");
   do_umbrella_tests(title_one, source_one, treatnames_one,
     xx_one, nn_one, kk_one, 4, dataformat_one);

   printf ("\n--------------------------------------------\n");
   printf ("--------------------------------------------\n");

   /* Second dataset */
   printf ("\nSECOND DATA SET:\n");
   /* Switch to 100,000 simulations for the second dataset */
   nsims=100000;
   do_umbrella_tests(title_two, source_two, treatnames_two,
     xx_two, nn_two, kk_two, 4, dataformat_two);

   /* Show total random numbers used */
   show_randnums_used();

   return 0; }



/* Functions to carry out known-p and unknown-p umbrella tests */

void do_one_umbrella(int pp, const double ww[MAXTOT], 
    const int nfirst[MAXK], const int nn[MAXK], int kk);

void do_max_umbrella(const double ww[MAXTOT], 
    const int nfirst[MAXK], const int nn[MAXK], int kk);



void do_umbrella_tests (const char *title, const char *source,
  const char *treatnames[MAXK],  const double xx[MAXK][MAXN],
  const int nn[MAXK], const int kk, const int pp,
  const char *dataformat)
 { int i,j,r, ntot;
   /* nfirst[i] is the offset in wwdat[] of the first value in row #i */
   int nfirst[MAXK];  
   /* Variables for midranks: */
   double wwdat[MAXTOT];      /* All observations */
   double mrank[MAXTOT];      /* The corresponding midranks */

   /* Show dataset description and source */
   printf ("%s\n%s\n", title,source);

   /* Count the total number of observations */
   for (ntot=i=0; i<kk; i++)  ntot+=nn[i];

   /* Display raw data */
   printf ("Number of treatments: %d    "
     "Total number of observations: %d\n", kk,ntot);
   for (i=0; i<kk; i++)
    { printf ("%s (n=%d): ", treatnames[i], nn[i]);
      /* `dataformat' is a function argument */
      for (j=0; j<nn[i]; j++)
       { printf ("  ");  printf (dataformat, xx[i][j]); }
      printf("\n"); }

   printf ("\nData means and rank averages for %d treatment groups:\n",kk);
   /* First, find the midranks */
   /* Copy observed data to a single one-dimensional array  wwdat[], */
   /*   keeping track of the initial offsets of each treatment group */
   /* Thus treatment #i is at  nfirst[i] <= r < nfirst[i]+nn[i] */
   for (r=i=0; i<kk; i++)
    { nfirst[i]=r;
      for (j=0; j<nn[i]; j++)  wwdat[r++]=xx[i][j]; }
   /* Put joint (Kruskal-Wallis) midranks at the corresponding */
   /*   locations in  mrank[] */
   /* makeranks() (in statlib.c) is intelligent enough to know */
   /*   that tgsize=NULL means that we don't want tie-group information */
   makeranks(mrank, wwdat,ntot, NULL,NULL);

   /* Now display the treatment-group averages */
   for (i=0; i<kk; i++)
    { int nf=nfirst[i], len=nn[i];
      double xsum, rsum;
      /* Sum over the i-th treatment group */
      xsum=rsum=0.0;
      if (len>0) /* This allows for empty treatment groups */
       { for (j=nf; j<nf+len; j++)
          { xsum += wwdat[j];  rsum += mrank[j]; }
         xsum=xsum/len;  rsum=rsum/len; }
      printf ("%s (n=%d):  DataAv: %6.2f   RankAv: %5.2f\n", 
        treatnames[i], len, xsum, rsum); }

   printf ("\nTHE UMBRELLA TEST for KNOWN p=%d (Mack-Wolfe):\n",pp);
   do_one_umbrella(pp, wwdat,nfirst,nn, kk);

   printf ("\nTHE UMBRELLA TEST with UNKNOWN p (Chen-Wolfe form):\n");
   printf ("\nNote that the one-sided P-values in the table above are NOT\n"
     "  multiple-comparison corrected.\n");

   /* Copy the observed data back to  wwdat[]  to be safe */
   for (r=i=0; i<kk; i++) for (j=0; j<nn[i]; j++)
     wwdat[r++]=xx[i][j];
   do_max_umbrella(wwdat, nfirst,nn,kk);
            }



/* Functions used in both do_one_umbrella() and do_max_umbrella() */

void make_mannwhitney_diffs(double uu[MAXK][MAXK],
   const double ww[MAXTOT], const int nfirst[MAXK],
   const int nn[MAXK], int kk);
void show_mwdiffs(double uu[MAXK][MAXK], int kk);
void show_example(int pp, double uu[MAXK][MAXK],
  const int nfirst[MAXK], int kk);
void show_umbrella_table(double uu[MAXK][MAXK], int kk,
    const int nn[MAXK], int pp);
double get_aap (int pp, double uu[MAXK][MAXK], int kk);
void show_pval (int nbig, int nsims);


/* The Umbrella Test for known p=pp (Mack-Brown) (see text) */

void do_one_umbrella(int pp, const double ww[MAXTOT], 
    const int nfirst[MAXK], const int nn[MAXK], int kk)
 { int i,r, nbig,ns, ntot;
   double uu[MAXK][MAXK], vv[MAXK][MAXK];
   double wtemp[MAXTOT], obs_aap;
   printf ("Single-umbrella test:\n");
   printf ("Is the data significant for a `peak' in "
     "treatment-group medians at p=%d?\n",pp);

   /* Put the Mann-Whitney differences of  ww  into uu[][] */
   make_mannwhitney_diffs(uu, ww,nfirst,nn,kk);

   /* Show the MW difference table */
   printf ("\nMann-Whitney differences between the treatment groups:\n");
   show_mwdiffs(uu,kk);

   /* Show an example calculation */
   show_example(pp,uu,nn,kk);

   /* Find the observed JT-like statistic Ap for p=pp */
   obs_aap=get_aap(pp,uu,kk);
   /* Show a table for Ap for all p */
   show_umbrella_table(uu,kk, nn,pp);

   printf ("\nSimulating the P-value of Ap=%g (p=%d) "
     "using %s permutations:\n", obs_aap, pp, commnum(nsims));

   /* Copy the data in  ww[]  to  wtemp[]  for permutations */
   for (ntot=i=0; i<kk; i++) ntot+=nn[i];
   for (r=0; r<ntot; r++)  wtemp[r]=ww[r];
   /* Now do the simulations. Since p is fixed, we use Ap as the */
   /*   test statistic */
   nbig=0;
   for (ns=0; ns<nsims; ns++)
    { double aap;
      shuffle_dub(wtemp,ntot);
      /* Put the Mann-Whitney differences of  wtemp[]  in vv[][] */
      make_mannwhitney_diffs(vv, wtemp,nfirst,nn,kk);
      /* Calculate Ap for the permuted data */
      aap=get_aap(pp,vv,kk);
      /* get_aap() returns an integer or half-integer, so no need */
      /*   to worry about roundoff error. */
      if (aap>=obs_aap) nbig++; }
   /* Display the estimated P-value and a CI for the P-value */
   show_pval(nbig,nsims); }


/* Subroutines for do_umbrella() and do_max_umbrella() */

/* For data for  kk  treatment groups stored in a single linear */
/*   array ww[], with the i-th treatment group starting at  nfirst[i]    */
/*   of length  nn[i],  put the corresponding Mann-Whitney differences */
/*   in the array  uu[][].  */

void make_mannwhitney_diffs(double uu[MAXK][MAXK],
   const double ww[MAXTOT], const int nfirst[MAXK],
   const int nn[MAXK], int kk)
 { int i,j, a,b;
   for (i=0; i<kk; i++) for (j=i+1; j<kk; j++)
    { int nfi=nfirst[i], ni=nn[i], nfj=nfirst[j], nj=nn[j];
      int ucount=0;  double ustat;
      for (a=0; a<ni; a++) for (b=0; b<nj; b++)
       { if (ww[nfi+a]<ww[nfj+b]) ucount+=2;
         else if (ww[nfi+a]==ww[nfj+b]) ucount++; }
      uu[i][j]=ustat=(double)ucount/2;
      uu[j][i]=ni*nj-ustat; }}


void show_mwdiffs(double uu[MAXK][MAXK], int kk)
 { int i,j;
   for (i=0; i<kk; i++)
    { printf (" %2d.",i+1);
      for (j=0; j<kk; j++)
       { double mw=uu[i][j];
         if (i==j) printf ("     -  ");
         else if (mw==floor(mw)) printf ("   %3d  ", (int)mw);
         else  printf ("   %5.1f", mw); }
      printf("\n"); }}


void show_example(int pp, double uu[MAXK][MAXK],
  const int nfirst[MAXK], int kk)
 { char *sgn="";  int i,j;  double ap=0.0;
   printf ("\nExample of an Ap statistic:\n");
   printf ("  Ap(%d) = ",pp);
   for (i=1; i<pp; i++) for (j=i+1; j<=pp; j++)
    { printf ("%sU(%d,%d)",sgn,i,j);  sgn="+"; }
   sgn=" + ";
   for (i=kk; i>pp; i--) for (j=i-1; j>=pp; j--)
    { printf ("%sU(%d,%d)",sgn,i,j);  sgn="+"; }
   printf("\n  =");  sgn="";
   for (i=1; i<pp; i++) for (j=i+1; j<=pp; j++)
    { double u=uu[i-1][j-1];  printf(" %s %g",sgn,u);  ap+=u;  sgn="+"; }
   sgn=" + ";
   for (i=kk; i>pp; i--) for (j=i-1; j>=pp; j--)
    { double u=uu[i-1][j-1];  printf(" %s %g",sgn,u);  ap+=u;  sgn="+"; }
   printf(" = %g\n",ap);
     }


/* Return the umbrella statistic  Ap  using Mann-Whitney differences */
/*   stored in  uu[][] */
/* Note p=1,2,...,kk  while i,j offsets are 0,1,..,kk-1 */

double get_aap (int pp, double uu[MAXK][MAXK], int kk)
 { int i,j;
   double aap=0.0;
   for (i=0; i<pp; i++)  for (j=i+1; j<pp; j++)
     aap += uu[i][j];
   for (i=pp-1; i<kk; i++)  for (j=i+1; j<kk; j++)
     aap += uu[j][i];
   return aap; }


/* Show main umbrella table */

/* Prototype for a function to return the mean-zero variance-one 
     scaled version of Ap
   The function code itself is later in the program. */

double get_astar(int pp, double aap, const int nn[MAXK], int kk,
      double *meanptr, double *varptr);

void show_umbrella_table(double uu[MAXK][MAXK], int kk,
    const int nn[MAXK], int pp)
 { int r;
   printf ("\nThe umbrella statistics Ap with "
      "normal approximations Astar:\n");
   for (r=1; r<=kk; r++)
    { char *starflag=" ";
      double ap, astar, mean,var;
      ap=get_aap(r,uu,kk);
      if (r==pp) starflag="*";
      astar=get_astar(r, ap,nn,kk, &mean,&var);
      printf (" %sp=%d  Ap=%5.1f   Mean=%5.1f  Var=%.2f   "
       "Astar=%6.3f   P=%7.5f (1-sided)\n",
        starflag, r, ap,mean,var, astar, 1-normint(astar)); }
           }



/* Return the large-sample normal approximation  Apstar  for  Ap. */
/* Return  mean=E(Ap)  and  var=Var(ap)  if requested */

double get_astar(int pp, double aap, const int nn[MAXK], int kk,
      double *meanptr, double *varptr)
 { int i, np,tn1,tn2, nn2,nn3, ntot;
   double astar, mean,var;
   /* Astar = (Ap-E(Ap))/sqrt(Var(Ap));  */
   /* First, calculate Mean=E(Ap) and Var=Var(Ap) */
   np=nn[pp-1];   /* p=1,2,...,kk  but  offsets=0,1,...,kk-1 */
   tn1=tn2=nn2=nn3=0;
   for (i=0; i<pp; i++) tn1+=nn[i];
   for (i=pp-1; i<kk; i++) tn2+=nn[i];
   for (i=0; i<kk; i++)
    { int ns=nn[i];  nn2+=ns*ns;  nn3+=ns*ns*(2*ns+3); }
   mean=(double)(tn1*tn1 + tn2*tn2 - nn2 - np*np)/4;
   for (ntot=i=0; i<kk; i++) ntot+=nn[i];
   var=2*(tn1*tn1*tn1 + tn2*tn2*tn2) + 3*(tn1*tn1 + tn2*tn2)
     - nn3 - np*np*(2*np+3) + 12*np*(tn1*tn2 - np*ntot);
   var/=72;
   /* Astar is normalized Ap */
   astar = (aap-mean)/sqrt(var);
   /* Always return astar (as function value) */
   /* Return mean and var if requested */
   if (meanptr!=NULL) *meanptr=mean;
   if (varptr!=NULL) *varptr=var;
   return astar; }





void show_pval (int nbig, int nsims)
 { double pp,pwid;
   pp=(double)nbig/nsims;
   pwid=1.960*sqrt(pp*(1-pp)/nsims);
   printf ("  P = %d/%s = %6.4f  95%% CI (%6.4f, %6.4f, %6.4f) (1-sided)\n",
     nbig, commnum(nsims), pp, pp-pwid, pp, pp+pwid); }



/* The Umbrella Test with UNKNOWN p (Chen-Wolfe form) (see text) */

void do_max_umbrella(const double ww[MAXTOT], const int nfirst[MAXK],
  const int nn[MAXK], int kk)
 { int i,r, p,pmax, ntot, nbig,ns;
   double uu[MAXK][MAXK], vv[MAXK][MAXK], wtemp[MAXTOT];
   double obs_max_astar;
   printf ("\nMultiple-umbrella test:\n");
   printf (
   "Is there a `peak' in treatment-group medians at SOME p, and if so,\n"
   "  at which treatment group?\n");

   /* Put the Mann-Whitney differences of  ww  into uu[][] */
   make_mannwhitney_diffs(uu, ww,nfirst,nn,kk);

   /* Find the largest Astar statistic and see where it is located */
   obs_max_astar=-1.0;  pmax=-1;
   for (p=1; p<=kk; p++)
    { double ap,zz;
      ap=get_aap(p,uu,kk);
      zz=get_astar(p, ap,nn,kk, NULL,NULL);
      if (p==1 || zz>obs_max_astar) { obs_max_astar=zz;  pmax=p; }}
   printf ("MAX ASTAR=%.3f  at  p=%d\n", obs_max_astar, pmax);
   printf ("\nSimulating the P-value of MaxAstar=%.3f "
    "using %s permutations:\n", 
     obs_max_astar, commnum(nsims));

   /* Copy the data in  ww[]  to  wtemp[]  for permutations */
   for (ntot=i=0; i<kk; i++) ntot+=nn[i];
   for (r=0; r<ntot; r++)  wtemp[r]=ww[r];
   /* Find P-value of MaxAstar for permuted copies of the data: */
   for (nbig=ns=0; ns<nsims; ns++)
    { double mstar;
      shuffle_dub(wtemp,ntot);
      /* Put the Mann-Whitney differences of  wtemp[]  in vv[][] */
      make_mannwhitney_diffs(vv, wtemp,nfirst,nn,kk);
      /* Find the maximum of kk Astar values */
      for (p=1; p<=kk; p++)
       { double aap,zz;
         aap=get_aap(p,vv,kk);
         zz=get_astar(p,aap,nn,kk, NULL,NULL);
         if (p==1 || zz>mstar)  mstar=zz; }
      /* Allow for a small amount of roundoff error */
      if (mstar>=obs_max_astar-(1e-12)) nbig++; }
   printf ("The multiple-comparison-corrected P-value "
    "for MaxAstar=%.3f is:\n", obs_max_astar);
   /* Display the estimated P-value and a CI for the P-value */
   show_pval(nbig,nsims); }