library(Matrix)
library(MASS)





# GEE function just for logistic regression, no other link function

# working correlations: corr="ind" ... independence, robust variance based on sandwich estimator by Lumley and Thas et al.
# corr="simple1", i.e. alpha1 !=0, and alpha2=0, model based variance (based on that assumption)
# corr="simple2"
# corr="exch" (exchangeable)



###############################################################################

GEE<- function(beta0=NULL,Y,X,corr=1,A=NULL,R0=NULL,iter.max=1e2,link="logit"){

  # link="logit" or link="probit"  
  
  eps<-1e-6 
  
N<-length(Y)
ind<-indexfct(n)
p<-dim(X)[2]



  n<-round(abs(polyroot(c(-2*N,-1,1)))[1])
  
  # get sparse matrix that indicates where correlations are non-zero

  # if beta not given or when 
  if(is.null(beta0) | corr==0){
    if(link=="logit"){
    res.logistic<-glm(Y~X-1,family=binomial)
    }else{
      res.logistic<-glm(Y~X-1,family=binomial(link="probit")) 
      
    }
    
    beta.hat <- res.logistic$coefficients
    se.beta.hat <- sqrt(diag(vcov(res.logistic)))
    
    
    
    # get U.matrix and XtWX.inv for robust variance
    eta<- as.numeric(X%*%beta.hat)
    
    
    
    
    #Varsqrtinv<-Diagonal(N,1/sqrt(Var))
    #Varsqrt<-Diagonal(N,sqrt(Var)) 
    
    if(link=="logit"){
      mu<-expit(eta)
      Var<- as.numeric(mu*(1-mu))
      D<-Diagonal(N,Var)
    }
    if(link=="probit"){
      mu<-pnorm(eta)
      Var<- as.numeric(mu*(1-mu))
      D<-Diagonal(N,pnorm(eta))
    }
  
     
    
    
    res<-Y-mu
    
    U.matrix <-  Diagonal(N,1/Var)%*%res  # U's without correlation and without XD
    
    DX<-D%*%X
    
    XtWX.inv <- solve(t(DX)%*%Diagonal(N,1/Var)%*%DX)
    
    beta0<-beta.hat
    R<-alpha<-alpha1<-alpha2<-NULL 
  }
  
  
  if(corr!=0){
  
    #print(beta0)
      
    beta.new<-beta.old<-beta0
    # could be improved  
    
    if(is.null(A)){
      A<-get.A(n)   
    }#end 
    
    
    
    
    
    eps<-1e-6
    Delta.beta <- 1
    iter<-0
    while(sum(abs(Delta.beta))>eps & iter<iter.max){
      
      beta.old<-beta.new
      
      iter<-iter+1
      eta<- Xbeta <- X%*%beta.old
      eta<-as.numeric(eta)
      
      if(link=="logit"){
      mu<-expit(eta)
      Var<- as.numeric(mu*(1-mu))
      D<-Diagonal(N,Var)
      }
      
      if(link=="probit"){
        mu<-pnorm(eta)
        Var<- as.numeric(mu*(1-mu))
        D<-Diagonal(N,dnorm(eta))  # dmu/deta=d (Phi (eta))/deta= phi(eta)
      }
      
      
      
    
      mu<-pmax(mu,1e-10)
      mu<-pmin(1-1e-10,mu)
      
      res<-Y-mu
      Var<- as.numeric(mu*(1-mu))
      
      Varsqrtinv<-Diagonal(N,1/sqrt(Var))
      Varsqrt<-Diagonal(N,sqrt(Var)) 
      
      
      ## D*Varsqrtinv=Varsqrtinv
      
      if(corr!=4){
      # corr=1 (alpha1),2 (alpha2),3 (exch)
      hilf<-estimate.R(res,Var,A,corr=corr)
      R<-hilf$R
      alpha1<-hilf$alpha1
      alpha2<-hilf$alpha2
      alpha<-hilf$alpha
      }else{
        R<-R0
      }
      
      R<-forceSymmetric(R)
      
      
      
      DX<-D%*%X
      
      
      
      # XtWX.inv <- solve(t(DX)%*%Diagonal(N,1/Var)%*%DX)  # only for independent data
      
      
      Varsqrtinvres<-Varsqrtinv%*%res  
      
      RinvVarsqrtinvres <- solve(R,Varsqrtinvres)
      
      
      VarsqrtinvDX<-Varsqrtinv%*%DX
      
      U<-  t(VarsqrtinvDX)%*%RinvVarsqrtinvres
      
      
    
      RinvVarsqrtDX <- solve(R,VarsqrtinvDX)
      XtWX<-t(VarsqrtinvDX)%*%RinvVarsqrtDX
      
      
      
      # U<- t(X)%*%D%*%Vsqrtinv%*%R.inv%*%Vsqrtinv%*%res 
      # d eta / d mu = log(mu/(1-mu))=1/(mu*(1-mu))
      
      # D : dmu/deta ==>> dmu/deta= Var for logistic regression 
      
      
      #CALCULATE NEW BETA
      
      #update beta uisng the formula in the paper
      #beta.new<-beta.old +  0.1*solve(t(X)%*%diag(D.inv)%*%C11.inv%*%diag(D.inv)%*%X)%*%t(X)%*%diag(D.inv)%*%C11.inv%*%res
      
    
      
      #XtWX.inv<- solve(XtWX)
      XtWX.inv<- ginv(as.matrix(XtWX))
      
      Delta.beta <- XtWX.inv%*%U
      
      beta.new<- beta.old + Delta.beta
      
      if(0){
      if(iter%%1==1){
      cat("iter:=",iter," |Delta|",sum(abs(Delta.beta)),"\n")
      cat("beta=",as.numeric(beta.new),"\n")
      }
      }
      #Delta.beta
    }#end while
    
    if(sum(abs(Delta.beta))>eps){cat("Not Converged too large Delta\n");return(NULL)}
    cat("Converged after",iter," Iterations\n")
    
    
    if(iter>=iter.max){cat("Not Converged too many operations\n");return(NULL)}
    
    beta.hat<-beta.new
    se.beta.hat<-sqrt(diag(XtWX.inv))
    
    
   
    U.matrix <-   Varsqrtinv%*%RinvVarsqrtinvres 
  
  }# if(corr!=0){
  
    #get robust variance estimator under sparse correlation
   
   U.matrix <-  matrix(U.matrix,N,p,byrow=F)*DX
   
   middle<-   t(U.matrix)%*%((A+Diagonal(N))%*%U.matrix) 
   #middle1<-   t(U.matrix)%*%matrix(c(matrix(1,1,N)%*%U.matrix),N,p,byrow=T)
   #middle1<-   t(U.matrix)%*%matrix(1,N,N)%*%U.matrix
   
   se.robust.beta.hat <- try(sqrt(diag(XtWX.inv%*% middle %*%XtWX.inv)))
   #se.robust1.beta.hat <- try(sqrt(diag(XtWX.inv%*% middle1 %*%XtWX.inv)))
    
    list(beta.hat=beta.hat,se.beta.hat=se.beta.hat,se.robust.beta.hat=se.robust.beta.hat,R=R,alpha=alpha,alpha1=alpha1,alpha2=alpha2)
    
    
  }# apply GEE algorithm for specific R-model
  
  ##########################################
  
  #end
  
  
  

if(0){
  #calculate Variance estimator
  
  
  # exchangable structure
  
  rho<-0.1
  
  off.diag.R.inv <- 1/((1+(N-1)*rho)*10)-(1/(N*(1-rho)))
  diag.R.inv<- 1/(1-rho) + off.diag.R.inv 
  
  M <- matrix(1,N,N);diag(B)<-0;
  
  R.inv<-  diag.R.inv*diag(N) +  off.diag.R.inv*M
  
  # robust according to sparse correlation
   dUdbeta <-           # A-Matrix for estimating equations
   B       <- t(U)%*%M%*%U  # B-Matrix for estimating equations
  
  
  
}# if(1/0)





estimate.R<-function(res,Var,A,corr=1){
N<-length(res)  
  
  
# model 1: just alpha1 \neq 0, alpha2=0  
# model 2: just alpha1 \neq 0, alpha2 \neq=0    
# model 3: exchangable  
  
r<-res/sqrt(Var)  
  
#number of links in A

#ones<-matrix(1,N,1)

N.corr.alpha2<-N.corr.alpha1<-alpha<-alpha1<-alpha2<-NULL

if(corr<=2){
N.corr.alpha1<-sum(A)
alpha1<- as.numeric(t(r)%*%A%*%r/N.corr.alpha1)
R<-Diagonal(N)
R<-R + alpha1*A
}

if(corr==2){
  A.c <- 1- A 
  diag(A.c)<-0
  N.corr.alpha2<-  sum(A.c) # or N*(N-1) - N.corr.alpha1
  
  
  alpha2<- as.numeric(t(r)%*%A.c%*%r/N.corr.alpha2)  
  R<-R + alpha2*A.c
}


if(corr==3){
  B<- matrix(1,N,N)
  diag(B)<-0
  alpha <- as.numeric(t(r)%*%B%*%r/(N*(N-1)))
  R<-Diagonal(N)
  R<-R + alpha*B
}

# R[30:40,30:40]  
  
return(list(R=R,N.corr.alpha2=N.corr.alpha2,N.corr.alpha1=N.corr.alpha1,alpha1=alpha1,alpha2=alpha2,alpha=alpha))
  
}#end GEE function

#########################################################################





##################################################################
# function for indexing odds
indexfct<-function(c){
  LODDS<-c*(c-1)/2;
  hilf<-NULL;
  for(i in 1:(c-1)){
    for(j in (i+1):c){
      hilf<-cbind(hilf,matrix(c(i,j),2,1))
    }
  }
  hilf
}#end indexfct
########################################################


abs.diff<-function(x,y){
  abs(x-y)  
}


get.A<-function(n){
ind<-indexfct(n)
A<-(outer(t(ind),(ind),FUN="abs.diff"))  # calculates differences of pairs, if zero index must match
#A[1,,,69]
#A <- Matrix(A, sparse = TRUE)
A<- 1*(apply(A,c(1,4),min)==0)
#A[1,69]
diag(A)<-0
#A <- as( A, "dgCMatrix")    
A <- Matrix(A, sparse = TRUE)
return(A)
}

get.A.rec<-function(n){
  
  if(n==2){
    return(Matrix(0,nrow=1,ncol=1, sparse = TRUE))
  }#if(n==1){
  
  
    
  #n>2, i.e. n=3,4,...
    
  A <- matrix(1,n-1,n-1);
  diag(A)<-0
  A <- Matrix(A, sparse = TRUE) 
  
  if(n>3){
  
    # obtain now the matrices A_i,j with i<j and put them all in B
  B<-Matrix(0,nrow=n-1,ncol=0, sparse = TRUE) 
  for(j in 2:(n-1)){
    B <- cBind(B,rBind(Matrix(0,nrow=j-2,ncol=n-j, sparse = TRUE),Matrix(1,1,n-j),Diagonal(n-j)))
  }#end for (i in 2:(n-1))
  }else{ ## essentially n==3
    
   B<- Matrix(1,nrow=n-1,ncol=1, sparse = TRUE)  
   B<- as(B, "CsparseMatrix")
     
  }#end 
  
  
  C<-rBind(cBind(A,B),cBind(t(B),get.A.rec(n-1)))  
  C<- as(C, "CsparseMatrix")
  return(C)
  
}# get.A.rec<-function(n){



#################################################################################
#################################################################################

output<-function(mat,row.names=NULL,digits=3,nsmall=2,special=NULL,every=1){
  hilf<-dim(mat)
  m<-hilf[1]
  n<-hilf[2]
  
  if(is.null(row.names)){row.names<-rownames(mat)}
  output1<-NULL
  for(i in 1:m){
    output1<-c(output1,row.names[i])
    for(j in 1:n){
      
      
      
      if(!is.na(mat[i,j]) & (mat[i,j]!=0 | j==1) ){  
        
        #if(j%%2==0 & j<=8){ 
        
        #  output1<-c(output1,"  ($",format(mat[i,j] ,nsmall=nsmall,digits=digits),"$)") 
        
        #}else{
        
        
        
        if(j==1){
          output1<-c(output1," & $",format(round(mat[i,j],1) ,nsmall=1,digits=1),"$")
          
        }else{
          output1<-c(output1," & $",format(mat[i,j] ,nsmall=nsmall,digits=digits),"$")  
          
        }
        
        #  }
        
      }else{
        if(j==7){
        output1<-c(output1," & -- ")
        }else{
        output1<-c(output1," & $",format(mat[i,j] ,nsmall=nsmall,digits=digits),"$")
        }
        
      }
      
    }  
    
    
    output1<-c(output1,"\\\\")
    #print(is.null(special))
    #print(j%%every)
    
    
    if(!is.null(special) &  i%%every==0){
      output1<-c(output1,special)
      
    }else{}
    
    output1<-c(output1,"\n  ")
    
  }#end   
  cat(output1,"\n ")
}#end

#################################################################################
#################################################################################