### Quantile Rank Score Test for PLVC model
### Author: Huixia Judy Wang
###            Department of Statistics, North Carolina State University
### Last Modified On: Thursday, August 12, 2010 at 16:59
### Created Date:  July 25, 2010

library(splines)
library(quantreg)


kn.SIC <- function(y, design.VC, design.C, time, N, tau, degree)
{
# choose the optimal kn based on SIC
# design.VC: design matrix of varying coeff
# design.C: design matrix for the parametric coeff
    SIC = NULL
    for(kn in 0:5)       
    {
        q = kn+degree+1
        u = seq(0, 1, length=kn+2)[-c(1,kn+2)]
        Knots = as.numeric(quantile(time, u))
        pi.t = bs(time, knots=Knots, intercept=TRUE, degree=degree)[,1:(q)]
        Pi = apply(design.VC, 2, function(x) x*pi.t)
        Pi = matrix(Pi, nrow=N)
        G = cbind(Pi, design.C)
        pkn=ncol(G)
        rq1=rq(y~G-1, tau=tau)
        res = rq1$resid
        SIC = c(SIC, log(sum(res*(tau-1*(res<0)))) + 0.5*log(N)*pkn/N)    
    }
    ################# Fit the model again with the optimal kn
    kn = max(which(SIC == min(SIC))-1)
    return(kn)
}


QRS.PLVC = function(y, time, design.VC, design.C, ID, tau=0.5, test.VC=c("VC","C"), test.idx, degree=3, se=c("iid","nid"), restrcited="yes")
{
    # y: response variable (N by 1 vector)
    # time: measurement time (N by 1 vector)
    # design.VC: design matrix of covariates with varying coefficients, N by p matrix, 
    #            the first column is 1 corresponding to the varying intercept
    # design.C: design matrix of covariates with constant coefficients, N by q matrix
    # test.VC = "VC": test if the constancy of varying coefficients
    # test.VC = "C": test if the constant coefficients equal zero or not
    # test.idx:  the index (indices) of the coefficients to be tested, 
    #           for instance, if test.VC="VC" and test.idx=1: test if the intercept coefficient is constant or not
    #                         if test.VC="VC" and test.idx=2: test if the varying coefficient of the first predictor is constant or not
    #                         if test.VC="C" and test.idx=1: test if the first constant coefficient is zero or not
    # se="iid": assume homotscedastic errors
    # se="nid": assume heteroscedastic errors, need estimate f_{ij}.  
    # restricted="yes": estimate the intra-subject correlation by using residuals estimated under the alternative model
    # restricted="no": estimate the intra-subject correlation by using residuals estimated under the null model
    
    N = length(y)
    n = length(unique(ID))  
    subject = model.matrix(~-1+factor(ID)) 
    K <- apply(subject, 2, sum)  # number of repeats for each subject, median is 12
    degree=3  # cubic spline

    design.VC = as.matrix(design.VC)
    design.C = as.matrix(design.C)
    p = ncol(design.VC)
    q = ncol(design.C)
    
    ############## Use SIC to choose the number of knots, kn
    kn = kn.SIC(y, design.VC, design.C, time, N, tau, degree)
    cat("the number of knots is", kn, "\n")
     
    ###### generate the B-spline design matrix
    dn = kn+degree+1
    u = seq(0, 1, length=kn+2)[-c(1,kn+2)]
    Knots = as.numeric(quantile(time, u))
    pi.t = bs(time, knots=Knots, intercept=TRUE, degree=degree)[,1:(dn)]
        
    if(test.VC=="C")
    {
        Z = design.C[,test.idx]
        X = NULL
        for(j in 1:p)
        {
            X = cbind(X, design.VC[,j]*pi.t)
            }
        if(q>length(test.idx)) X = cbind(X, design.C[,-test.idx])
        }
    if(test.VC=="VC")
    {
        X = Z = NULL
        for(j in (1:p)[-test.idx])
        {
            X = cbind(X, design.VC[,j]*pi.t)
            }
        for(j in test.idx)
        {
            X = cbind(X, design.VC[,j])
            Z = cbind(Z, (design.VC[,j]*pi.t)[,-1])
            }
        }
 
 # calculate zstar
 if(se=="iid") 
    zstar = as.matrix(qr.resid(qr(X), Z))
   if(se=="nid")
   {
        x = X
        eps=.Machine$double.eps^(2/3)
        
        h <- bandwidth.rq(tau, N, hs = TRUE)
        bhi <- rq.fit.br(x, y, tau + h, ci = FALSE)
        bhi <- coefficients(bhi)
        blo <- rq.fit.br(x, y, tau - h, ci = FALSE)
        blo <- coefficients(blo)
        dyhat <- x %*% (bhi - blo)
        if (any(dyhat <= 0)) {
            pfis <- (100 * sum(dyhat <= 0))/N
            warning(paste(pfis, "percent fis <=0"))
        }
        f <- pmax(eps, (2 * h)/(dyhat - eps))
        B = diag(f)
        Pw = X%*%solve(t(X)%*%B%*%X)%*%t(X)%*%B
        zstar= Z - Pw%*%Z
    }
    
    # perform the inference using QRS methods.
    rq0 = rq(y~X-1, tau=tau)
    ehat = rq0$resid
    ehat = ehat - quantile(ehat, tau)     
    psi = tau - 1*(ehat<0)
    
    if(restricted=="no") 
    {
        rq1 = rq(y~X+Z-1, tau=tau)
        ehat = rq1$resid
        psi = tau - 1*(ehat<0) 
    }
    
    ## indep: assume independence
    ## QRS.g: unspecified correlation structure
    ## QRS.CS: assume compound symmetry covariance, exchangeable correlations
    test = QRS.Tn(X, zstar, psi,  subject, N, K, tau, n, nt) 
    Tn.g = test$Tn.g
    Tn.CS = test$Tn.CS
    p1 = length(test.idx)
    if(test.VC=="C")
    {
        QRS.g = 1-pchisq(Tn.g,p1)
        QRS.CS = 1-pchisq(Tn.CS,p1)
        }
    if(test.VC=="VC")
    {
        dd = (kn+degree)*p1
        tt.g = (Tn.g-dd)/sqrt(2*dd)
        tt.CS = (Tn.CS-dd)/sqrt(2*dd)
        QRS.g = 2*pnorm(-abs(tt.g))
        QRS.CS = 2*pnorm(-abs(tt.CS))
        }
        
    ### estimated coefficient curves and the effects of z under the full model
        X = NULL
        for(j in 1:p)
        {
            X = cbind(X, design.VC[,j]*pi.t)
            }
        X = cbind(X, design.C)
        
        theta = rq(y~X-1, tau)$coef
        coef.VC = NULL
        
        for(j in 1:p)
        {
            coef.VC=cbind(coef.VC, pi.t %*% theta[((j-1)*dn+1):(dn*j)])
        }
        coef.C = theta[-(1:p*dn)]
        
    out = list(QRS.g=QRS.g, QRS.CS=QRS.CS, Tn.g=Tn.g, Tn.CS=Tn.CS, coef.VC=coef.VC, coef.C=coef.C)
    return(out)
}
        
        
        

QRS.Tn = function(X, zstar, psi,  subject, N, K, tau, n, nt) 
{
    # calculate the QRS test statistic Tn
    # Tn.indep: assuming independent errors
    # Tn.g:  assuming an unspecified correlation structure
    # Tn.CS: assuming an exchangeable correlation
    
    # K is the number of repeated measurements of each subject
    P1 <- crossprod(zstar)
    P2 <- t(zstar)%*%subject%*%t(subject)%*%zstar - P1
    sn = 1/sqrt(N)*t(zstar)%*%psi

    # QRS assuming independence
    Qn.indep = 1/N*P1*tau*(1-tau) + 0.00001  
    Tn.indep = t(sn)%*%solve(Qn.indep)%*%sn
    if(length(zstar)>N)     q = ncol(zstar)
    if(length(zstar)==N) q=1
    
    ### calculate delta, assuming an exchangeable correlation
    an = tau-psi  #1*(ehat<0)
    if(length(X)==N) pkn=1
    else pkn=ncol(X)
    delta <- (t(an)%*%subject%*%t(subject)%*%(an) - crossprod(an))/(sum(apply(as.matrix(K),1,function(x)x*(x-1)))-pkn)
    Qn.CS <- Qn.indep + 1/N*P2 * as.numeric(-tau^2+delta)
    if(max(abs(Qn.CS))<=0.00001) Qn.CS = Qn.CS + 0.00001 
    Tn.CS = t(sn)%*%solve(Qn.CS)%*%sn  
    
    ### The QRS assuming an unspecified correlation structure
    # fill in those missed measurements (be careful!)
    if(q==1)
    {
        psi2 = zstar2 = rep(0, n*nt)
        for(i in 1:n)
        {
            ind = which(subject[,i]==1)
            ind2 = ((i-1)*nt+1): (i*nt)
            psi2[ind2[1:length(ind)]] = psi[ind] 
            zstar2[ind2[1:length(ind)]] = zstar[ind]
        }
        psi.m = matrix(psi2, ncol=n)
        zstar.m = matrix(zstar2, ncol=n)
        Obj = rbind(psi.m, zstar.m)
        Qn.g = 1/N*sum(apply(Obj, 2, function(x) crossprod(x[-(1:nt)],x[(1:nt)])^2))
    }
    if(q>1)
    {
        psi2 = rep(0, n*nt)
        zstar2 = matrix(0, nrow=n*nt, ncol=q)
        for(i in 1:n)
        {
            ind = which(subject[,i]==1)
            ind2 = ((i-1)*nt+1): (i*nt)
            psi2[ind2[1:length(ind)]] = psi[ind] 
            zstar2[ind2[1:length(ind)],] = zstar[ind,]
        }
        Obj = cbind(zstar2, psi2)
        Obj = array(Obj, dim=c(nt, n, q+1)) # the first [,,1:q] are zstar.m, and the last one is psi.m
        Qn.g = 1/N*matrix(apply(
                apply(Obj, 2, function(xx) 
                {
                    z1 = xx[,1:q]
                    ppsi = xx[,q+1]
                    out = t(z1)%*%ppsi%*%t(ppsi)%*%z1
                    return(out)
                }), 1, sum), ncol=q)
    }
    if(max(abs(Qn.g))<=0.00001) Qn.g = Qn.g + 0.00001   # modified this for the xyB for small nsubj=25
    Tn.g = t(sn)%*%solve(Qn.g)%*%sn 
    tmp.out = data.frame(Tn.indep=Tn.indep, Tn.g=Tn.g, Tn.CS=Tn.CS)    
    return(tmp.out)
}