/*
** The survfit routine for raw data, not multi-state, not Cox
** See methods document, survival:survfitkm
**
**  For a variable "zed" used in the code, I use "zed2" as the argument
**    passed from R
**  y: 2 or 3 column matrix of survival times
**  weight: vector of case weights
**  sort1:  vector of sort indices for the start time (if y has 3 cols),
**    will be null otherwise
**  sort2: vector of sort indices for the end time
**  type: the type of survival compuation
**  id: used for clustered variance; idenitfies the group. 0,1,2,...
**  nid: number of unique groups
**  position: help for someone with multiple observations, e.g. (1,10), (10,14),
**    (14,20).   1*(first obs for a subject) + 2*(last obs for a subject).  
**    Someone with only one observation will have a value of 3.
**  influence: what influence matrices to return, if any.  
**       1*(cumhaz) + 2 * (survival)
**  reverse: 0= ordinary KM, 1= estimate of censoring distribution (not finished)
**  entry:   1= also keep rows with unique entry times
*/
#include <math.h>
#include "survS.h"
#include "survproto.h"

SEXP survfitkm(SEXP y2, SEXP weight2,  SEXP sort12, SEXP sort22, 
               SEXP type2, SEXP id2, SEXP nid2,   SEXP position2,
               SEXP influence2, SEXP reverse2,   SEXP entry2) {
              
    int i, i1, i2, j, k, person1, person2;
    int nused, nid, type, influence;
    int ny, ntime;
    int reverse, entry;
    double *time1=0, *time2, *status, *wt;
    double v1, v2, dtemp, haz;
    double temp, dtemp2, dtemp3, frac, btemp;
    double d0, d1, nrisk;
    int *sort1=0, *sort2, *id=0;
    static const char *outnames[]={"time", "n", "estimate", "std.err",
                                     "influence1", "influence2", ""};
    SEXP rlist;
    double *gwt=0, *inf1=0, *inf2=0, *inf3=0;  /* =0 to silence -Wall */
    int *gcount=0;
    int n1, n2, n3, n4;
    int *position=0;
    double wt1, wt2, wt3, wt4;
                      
    /* output variables */
    double  *n[8],  *dtime,
            *kvec, *nvec, *std[3], *imat1=0, *imat2=0;  /* =0 to silence -Wall*/
    double km, nelson;  /* current estimates */

    /* map the input data */
    ny = ncols(y2);     /* 2= ordinary survival 3= start,stop data */
    nused = nrows(y2);
    if (ny==3) { 
        time1 = REAL(y2);
        time2 = time1 + nused;
        sort1 = INTEGER(sort12);
    }
    else time2 = REAL(y2);
    status= time2 +nused;
    wt = REAL(weight2);
    sort2 = INTEGER(sort22);
    nused = LENGTH(sort22);
                   
    type = asInteger(type2);
    nid = asInteger(nid2);
    if (nid >0) id = INTEGER(id2);
    position = INTEGER(position2);
    influence = asInteger(influence2);
    reverse = asInteger(reverse2);
    entry   = asInteger(entry2);

    /* nused was used for two things just above.  The first was the length of
       the input data y, only needed for a moment to set up time1, time2, and
       status.  The second is the number of these observations we will actually
       use, which is the length of sort2.  This routine can be called multiple
       times with sort1/sort2 pointing to different subsets of the data while
       y, wt, id and position can remain unchanged
    */

    /* pass 1, get the number of unique times, needed for memory allocation 
      Number of xval groups (unique id values) has been supplied.  I count
      both unique starting and ending times.
    */
    if (entry ==1) {
        ntime =1; 
        temp = time1[sort1[0]];  /* less than any event/censor time */
	j =1;
        for (i=0; i<nused; i++) {
            i2 = sort2[i];
	    for (; j< nused && (time1[sort1[j]] < time2[i2]); j++) {
		i1 = sort1[j];
		if (time1[i1] != temp && (position[i1] & 1)==1) {
		    ntime++;
		    temp = time1[i1];
		}
	    }
            if ((time2[i2] != temp) && (position[i2] >1 || status[i2]>0)) {
                ntime++;
                temp = time2[i2];
            }
        }
    } else {
	/* count unique time2 values (with position >1)*/
        ntime = 1; 
        temp = time2[sort2[0]];  /* shortest ending time */
        for (i=1; i<nused; i++) {
            i2 = sort2[i];
            if ((position[i2] >1 || status[i2]>0) && time2[i2] != temp) {
                ntime++;
                temp = time2[i2];
            }
        }
    }
 
   /* Allocate memory for the output 
        n has 6 columns for number at risk, events, censor, then the 
        3 weighted versions of the same, then optionally two more for
        number added to the risk set (when entry=1)
    */
    PROTECT(rlist = mkNamed(VECSXP, outnames));
    
    dtime  = REAL(SET_VECTOR_ELT(rlist, 0, allocVector(REALSXP, ntime)));
    if (entry==0) j=6;  else j=8;
    n[0]  = REAL(SET_VECTOR_ELT(rlist, 1, allocMatrix(REALSXP, ntime, j)));
    for (i=1; i<j; i++) n[i] = n[0] + i*ntime;

    kvec  = REAL(SET_VECTOR_ELT(rlist, 2, allocMatrix(REALSXP, ntime, 2)));
    nvec  = kvec + ntime;  /* Nelson-Aalen estimate */
    /* std of survival, cumhaz, and auc, latter only if robust is used */
    std[0] = REAL(SET_VECTOR_ELT(rlist, 3, allocMatrix(REALSXP, ntime,ny)));
    std[1] = std[0] + ntime;
    if (ny ==3) std[2] = std[0] + ntime;

    if (nid >0 ) { /* robust variance */
        gcount = (int *) R_alloc(nid, sizeof(int));
        if (type <3) {  /* working vectors for the influence */
            gwt  = (double *) R_alloc(4*nid, sizeof(double)); 
            inf1 = gwt + nid;
            inf2 = inf1 + nid; 
	    inf3 = inf2 + nid;
            for (i=0; i< nid; i++) {
                gwt[i] =0.0;
                gcount[i] = 0;
                inf1[i] =0; /* per subject IJ, survival, cumhaz, auc */
                inf2[i] =0;
		inf3[i] =0;
            }
        }
        else {
            gwt = (double *) R_alloc(2*nid, sizeof(double));
            inf2 = gwt + nid;
            for (i=0; i< nid; i++) {
                gwt[i] =0.0;
                gcount[i] = 0;
                inf2[i] =0;
            }
        }

        /* these are not accumulated, so do not need to be zeroed */
        if (type <3) {
            if (influence==1 || influence ==3) 
                imat1 = REAL(SET_VECTOR_ELT(rlist, 4,
                                     allocMatrix(REALSXP, nid, ntime)));
            if (influence==2 || influence==3) 
                imat2 =  REAL(SET_VECTOR_ELT(rlist, 5,
                           allocMatrix(REALSXP, nid, ntime))); 
        }                
        else  if (influence !=0) 
                imat2 = REAL(SET_VECTOR_ELT(rlist, 5,
                               allocMatrix(REALSXP, nid, ntime))); 
    }

    R_CheckUserInterrupt();  /*check for control-C */

    /* 
    ** First fill in the list of unique times, a reprise of the code that
    **  counted the number of them.  (But now dtime has been allocated.)
    */
    if (entry==1) {
	temp = time1[sort1[0]];  /* less than any event/censor time */
	dtime[0] = temp;
	k =1; 
	j =1;
        for (i=0; i<nused; i++) {
            i2 = sort2[i];
	    for (; j< nused && (time1[sort1[j]] < time2[i2]); j++) {
		i1 = sort1[j];
		if (time1[i1] != temp && (position[i1] & 1)==1) {
		    temp = time1[i1];
		    dtime[k++] = temp;
		}
	    }
            if ((time2[i2] != temp) && (position[i2] >1 || status[i2]>0)) {
                temp = time2[i2];
		dtime[k++] = temp;
            }
        }
    } else {
	/* count ending times */
        temp = time2[sort2[0]];  /* shortest ending time */
	dtime[0]= temp;
	k=1;
        for (i=1; i<nused; i++) {
            i2 = sort2[i];
            if ((position[i2] > 1 || status[i2]>0) && time2[i2] != temp) {
                temp = time2[i2];
		dtime[k++] = temp;
            }
        }
    }
 
    /*
    ** Next compute all the counts
    ** Temp variables n1 = number at risk, n2= number of events, n3 = censors,
    **   n4 = additions.  wt1, wt2, wt3, wt4 = case weighted version of
    **   n1 to n4.  Store this in n[][0-5] for return to the R routine.
    ** person1, person2 track through sort1 and sort2, respectively
    **  likewise with i1 and i2
    */
    person1 = nused-1;  person2 = nused-1;
    n1=0; wt1=0;
    for (k=ntime-1; k>=0; k--) {
        n2=0; n3=0; wt2=0; wt3=0; 
        for (; person2 >=0; person2--) {
            i2= sort2[person2];
            if (time2[i2] < dtime[k]) break;

            n1++;             /* number at risk */
            wt1 += wt[i2];    /* weighted number at risk */
            if (status[i2] ==1) {
                n2++;            /* events */
                wt2 += wt[i2];
            } else if (position[i2]& 2) {
                /*
                ** if this is the last of a string (a,b](b,c](c,d].. for
                **  a subject (position[i2]=2 or 3), then it is a 'real' censor
                */
                n3++;
                wt3 += wt[i2];
            }
        }
        
        if (ny==3) {  /* go through entry times */
            n4 =0; wt4 =0;
            for (; person1 >=0; person1--) {
                i1 = sort1[person1];
                if (time1[i1] < dtime[k]) break;
                n1--; /* entry is >= dtime, remove from risk set */
                wt1 -= wt[i1];
                if ((entry==1) && (position[i1] & 1) && time1[i1] == dtime[k]) {
		    /* an entry at exactly this time */
		    n4++;
		    wt4 += wt[i1];
		}
            }
            if (entry==1) {
		n[6][k] = n4;
		n[7][k] = wt4;
	    }
        }

        n[0][k] = n1;  n[1][k]=n2;  n[2][k]=n3;
        n[3][k] = wt1; n[4][k]=wt2; n[5][k]=wt3; 
    }

    R_CheckUserInterrupt();  /*check for control-C */


    /*
    ** compute the survival and cumulative hazard, simple variance
    **  see the methods document, survfitkm: estimates
    */
    nelson =0.0; km=1.0; 
    v1=0; v2=0;
    for (i =0; i<ntime; i++) {
	if (reverse == 1) {    /* estimate G, the censoring curve */
	    d0  = n[2][i]; /* number of censors */
	    d1 = n[5][i]; /* weighted censors */
	    nrisk = n[3][i]- n[4][i]; /* at risk */
	} else {
	    d0 = n[1][i];    /* number of deaths, unweighted */
	    d1 = n[4][i];    /* number of deaths, weighted */
	    nrisk = n[3][i]; /* number at risk */
	}	
	
	if (type==1 || type==3) {  /* Nelson-Aalen hazard */
	    if (d0>0 && d1>0) {  /* at least one event with wt>0*/
		nelson += d1/nrisk;
		v2 += d1/(nrisk*nrisk);
	    }
	    nvec[i] = nelson;
	    std[1][i] = sqrt(v2);
	} else {              /* Fleming hazard */
	    for (j=0; j<d0; j++) {
		dtemp = nrisk - j*d1/d0;
		nelson += d1 /(d0* dtemp);
		v2 += d1/(d0*dtemp*dtemp);
	    }
	    kvec[i] = exp(-nelson);
	    nvec[i] = nelson;
	    std[1][i] = sqrt(v2);
	}
	
	if (type < 3) {  /* KM survival */
	    if (d0>0 && d1>0) {  /* at least one event */
		km *= (nrisk-d1)/nrisk;
		v1 += d1/(nrisk * (nrisk - d1)); /* Greenwood */
	    }
	    kvec[i] = km;
	    std[0][i] = sqrt(v1);
	} else {  /* exp survival */
	    kvec[i] = exp(-nvec[i]);
	    std[0][i] = std[1][i];
	}
    }

    if (nid > 0) {
	/* replace the variances above with IJ estimates, optionally save out
	**  the influence. See methods document, survfitkm: robust variance
	*/
	v1=0; v2 =0;
	person2=0; 
	if (ny==3) {
	    person1 =0;
	} else {
	    /* at the start, everyone is at risk */
	    for (i=0; i< nused; i++) {
		i2 = id[sort2[i]];
		gcount[i2]++;
		gwt[i2] += wt[sort2[i]];
            }
        }
            
	km =1 ; /* This lags behind */
        if (type==1) {  /* KM survival, NA hazard*/
            for (i=0; i< ntime; i++) {
		/* toss the outdated */
		for (; person2 < nused; person2++) {
		    i2 = sort2[person2];
		    if (time2[i2] >= dtime[i]) break;
		    gcount[id[i2]] --;
		    if (gcount[id[i2]]==0) gwt[id[i2]] =0;
		    else gwt[id[i2]] -= wt[i2];
		}
		if (ny == 3) {
                    /* add in new subjects */
                    for (; person1 < nused; person1++) {
                        i1 = sort1[person1];
                        if (time1[i1] >= dtime[i]) break;
			gcount[id[i1]]++;
			gwt[id[i1]] += wt[i1];
		    }
                }
         
                if (n[1][i] > 0 && n[4][i]>0) { /* need to update the sums */
                    haz = n[4][i]/n[3][i];
                    for (k=0; k< nid; k++) {
                        inf1[k] = inf1[k] *(1.0 -haz) + gwt[k]*km*haz/n[3][i];
                        inf2[k] -= gwt[k] * haz/n[3][i];
                     }
                    for (; person2<nused; person2++) { 
                        i2 = sort2[person2];
                        if (time2[i2] > dtime[i]) break; /* those at this time */
                        if (status[i2]==1) {
                            inf1[id[i2]] -= km* wt[i2]/n[3][i];
                            inf2[id[i2]] += wt[i2]/n[3][i];
                        }
			gcount[id[i2]] --;
			if (gcount[id[i2]] ==0) gwt[id[i2]] = 0.0;
			else gwt[id[i2]] -= wt[i2];
		    }
                    km *= (1-haz);
                    nelson += haz;
                   
                    v1=0; v2=0;
                    for (k=0; k<nid; k++) {
                        v1 += inf1[k]*inf1[k];
                        v2 += inf2[k]*inf2[k];
                    }
                } 
         
		std[0][i] = sqrt(v1);
                std[1][i] = sqrt(v2);
                if (influence==1 || influence ==3) 
                    for (k=0; k<nid; k++) *imat1++ = inf1[k];
                if (influence==2 || influence ==3)
                    for (k=0; k<nid; k++) *imat2++ = inf2[k];
            }
        }
        else if (type==2) {  /* KM survival, Fleming-Harrington hazard */
            for (i=0; i< ntime; i++) {
 		/* toss the outdated */
		for (; person2 < nused; person2++) {
		    i2 = sort2[person2];
		    if (time2[i2] >= dtime[i]) break;
		    gcount[id[i2]] --;
		    if (gcount[id[i2]]==0) gwt[id[i2]] =0;
		    else gwt[id[i2]] -= wt[i2];
		}
		if (ny == 3) {
                    /* add in new subjects */
                    for (; person1 < nused; person1++) {
                        i1 = sort1[person1];
                        if (time1[i1] >= dtime[i]) break;
			gcount[id[i1]]++;
			gwt[id[i1]] += wt[i1];
		    }
                }
        
                if (n[1][i] > 0 && n[4][i] >0) { /* need to update the sums */
                    dtemp =0;  /* the working denominator */
                    dtemp2=0;  /* sum of squares */
                    dtemp3=0;
                    temp = n[3][i] - n[4][i];  /* sum of weights for the non-deaths */
                    for (k=n[1][i]; k>0; k--) {
                        frac = k/n[1][i];
                        btemp = 1/(temp + frac*n[4][i]);  /* "b" in the math */
                        dtemp += btemp;
                        dtemp2 += btemp*btemp*frac;
                        dtemp3 += btemp*btemp;    /* non-death deriv */
                    }

                    dtemp /=  n[1][i];        /* average denominator */
                    if (n[4][i] != n[1][i]) { /* case weights */
                        dtemp2 *= n[4][i]/ n[1][i];
                        dtemp3 *= n[4][i]/ n[1][i];
                    }
                    nelson += n[4][i]*dtemp;

                    haz = n[4][i]/n[3][i];
                    for (k=0; k< nid; k++) {
                        inf1[k] = inf1[k] *(1.0 -haz) + gwt[k]*km*haz/n[3][i];
                        if (gcount[k]>0) inf2[k] -= gwt[k] * dtemp3;
                    }
                    for (; person2<nused; person2++) {                
                        /* catch the endpoints up to this event time */
                        i2 = sort2[person2];
                        if (time2[i2] > dtime[i]) break;
                        if (status[i2]==1) {
                            inf1[id[i2]] -= km* wt[i2]/n[3][i];
                            inf2[id[i2]] += wt[i2] *(dtemp + dtemp3 - dtemp2);
                         }
			gcount[id[i2]] --;
			if (gcount[id[i2]] ==0) gwt[id[i2]] = 0.0;
			else gwt[id[i2]] -= wt[i2];
		    }
                    km *= (1-haz);
                    
                    v1=0; v2=0;
                    for (k=0; k<nid; k++) {
                        v1 += inf1[k]*inf1[k];
                        v2 += inf2[k]*inf2[k];
                    }
                }
         
		std[0][i] = sqrt(v1);
                std[1][i] = sqrt(v2);
                if (influence==1 || influence ==3) 
                    for (k=0; k<nid; k++) *imat1++ = inf1[k];
                if (influence==2 || influence ==3)
                    for (k=0; k<nid; k++) *imat2++ = inf2[k];
            }
        }

        else if (type==3) {  /* exp() survival, NA hazard */
            for (i=0; i< ntime; i++) {
		/* toss the outdated */
		for (; person2 < nused; person2++) {
		    i2 = sort2[person2];
		    if (time2[i2] >= dtime[i]) break;
		    gcount[id[i2]] --;
		    if (gcount[id[i2]]==0) gwt[id[i2]] =0;
		    else gwt[id[i2]] -= wt[i2];
		}
                if (ny==3) {
                    /* add in new subjects */
                    for (; person1 < nused; person1++) {
                        i1 = sort1[person1];
                        if (time1[i1] >= dtime[i]) break;
			gcount[id[i1]]++;
			gwt[id[i1]] += wt[i1];
		    }
                }
                 
                if (n[1][i] > 0 && n[4][i]>0) { /* need to update the sums */
                    haz = n[4][i]/n[3][i];
                    for (k=0; k< nid; k++) {
                        inf2[k] -= gwt[k] * haz/n[3][i];
                    }
                    for (; person2<nused; person2++) { 
                        /* catch the endpoints up to this event time */
                        i2 = sort2[person2];
                        if (time2[i2] > dtime[i]) break;
                        if (status[i2]==1) {
                             inf2[id[i2]] += wt[i2]/n[3][i];
                        }
			gcount[id[i2]] --;
			if (gcount[id[i2]] ==0) gwt[id[i2]] = 0.0;
			else gwt[id[i2]] -= wt[i2];
		    }
                    nelson += haz;
                   
                    v2=0;
                    for (k=0; k<nid; k++) {
                        v2 += inf2[k]*inf2[k];
                    }
		}
         
                std[1][i] = sqrt(v2);
                std[0][i] = sqrt(v2);
                
                if (influence>0)
                    for (k=0; k<nid; k++) *imat2++ = inf2[k];
            }

        } else {  /* exp() survival,  Fleming-Harrington hazard */
            for (i=0; i< ntime; i++) {
 		/* toss the outdated */
		for (; person2 < nused; person2++) {
		    i2 = sort2[person2];
		    if (time2[i2] >= dtime[i]) break;
		    gcount[id[i2]] --;
		    if (gcount[id[i2]]==0) gwt[id[i2]] =0;
		    else gwt[id[i2]] -= wt[i2];
		}
                 if (ny==3) {
                    /* add in new subjects */
                    for (; person1 < nused; person1++) {
                        i1 = sort1[person1];
                        if (time1[i1] >= dtime[i]) break;
			gcount[id[i1]]++;
			gwt[id[i1]] += wt[i1];
		    }
                }
                if (n[1][i] > 0 && n[4][i] >0) { /* need to update the sums */
                        dtemp =0;  /* the working denominator */
                    dtemp2=0;  /* sum of squares */
                    dtemp3=0;
                    temp = n[3][i] - n[4][i];  /* sum of weights for the non-deaths */
                    for (k=n[1][i]; k>0; k--) {
                        frac = k/n[1][i];  
                        btemp = 1/(temp + frac*n[4][i]);  /* "b" in the math */
                        dtemp += btemp;
                        dtemp2 += btemp*btemp*frac;
                        dtemp3 += btemp*btemp;    /* non-death deriv */
                    } 
                    
                    dtemp /=  n[1][i];        /* average denominator */
                    if (n[4][i] != n[1][i]) { /* case weights */
                        dtemp2 *= n[4][i]/ n[1][i];
                        dtemp3 *= n[4][i]/ n[1][i];
                    }
                    nelson += n[4][i]*dtemp;

                    for (k=0; k< nid; k++) {
                        if (gcount[k]>0) inf2[k] -= gwt[k] * dtemp3;
                    }
                    for (; person2<nused; person2++) { 
                         i2 = sort2[person2];
                        if (time2[i2] > dtime[i]) break;
                        if (status[i2]==1) {
                            inf2[id[i2]] += wt[i2] *(dtemp + dtemp3 - dtemp2);
                        }
			gcount[id[i2]] --;
			if (gcount[id[i2]] ==0) gwt[id[i2]] = 0.0;
			else gwt[id[i2]] -= wt[i2];
		    }
            
                    v2=0;
                    for (k=0; k<nid; k++) v2 += inf2[k]*inf2[k];
                }
		
		std[1][i] = sqrt(v2);
                std[0][i] = sqrt(v2);
                
                if (influence>0)
                    for (k=0; k<nid; k++) *imat2++ = inf2[k];
            }
        }
    }
    
    UNPROTECT(1);
    return(rlist);
}
