#library(statnet)
library(ergm)
library(MASS)
library(tensor)
library(magic)
source("fcts.Lazega.r")

source("marginal.ERGM.fct.r")



# Dataset: Lazega (2001) - lawyers  in a New Engalnd Law Firm
# Description available at http://www.stats.ox.ac.uk/~snijders/siena/Lazega_lawyers_data.htm




#Lazege Lawyers


n<-36           #number lawyers
Ny<-n*(n-1)/2   #number dyads

#loading data set
#first dyads
Y<-as.matrix(read.table("ELwork36.txt"))
work<-as.network(Y,directed=FALSE)
Y<-as.matrix.network(work)

#read the attributes
h2<-read.table("ELattr.txt")

seniority<-h2[1:n,1]
status<-h2[1:n,2]
gender<-h2[1:n,3]
office<-h2[1:n,4]
years<-h2[1:n,5]
age<-h2[1:n,6]
practice<-h2[1:n,7]
lawschool<-h2[1:n,8]

#create main effects
seniority.sum<-matrix(seniority,n,n,byrow=F)+ matrix(seniority,n,n,byrow=T)
practice.sum<-matrix(practice,n,n,byrow=F) + matrix(practice,n,n,byrow=T)
diag(seniority.sum)<-0
diag(practice.sum)<-0

#similarity effects
practice.sim <-  1*((matrix(practice,n,n,byrow=F)- matrix(practice,n,n,byrow=T))==0)
gender.sim   <-  1*((matrix(gender,n,n,byrow=F)- matrix(gender,n,n,byrow=T))==0)
office.sim   <-  1*((matrix(office,n,n,byrow=F)- matrix(office,n,n,byrow=T))==0)
diag(practice.sim)<-0
diag(gender.sim)<-0
diag(office.sim)<-0

#model formula for standard ERGM
formula1<-as.formula("work ~ edges + edgecov(seniority.sum) + edgecov(practice.sum) + edgecov(practice.sim) +
edgecov(gender.sim) + edgecov(office.sim)  + gwesp(fixed=FALSE)")

#fit standard ERGM for Lazega data set, same model as in HH06
r<-ergm(formula1,MCMCsamplesize=20000)
#output of log-likelihood, coefficients, etc.
summary(r)
r$loglikelihood
r$coef
r$MCMCtheta
r$mle.lik
1-pnorm(abs(summary(r)$coefs[,1]/summary(r)$coefs[,2]))
summary(r)$coefs[,4]

#gradient	 The value of the gradient vector of the approximated loglikelihood function, evaluated at the maximizer. This vector should be very close to zero.
#covar	 Approximate covariance matrix for the MLE, based on the inverse Hessian of the approximated loglikelihood evaluated at the maximizer.
#samplesize	 The size of the MCMC sample
#failure	 Logical: Did the MCMC estimation fail?
#mc.se

r$gradient
r$covar
r$samplesize
r$failure
r$mc.se

#t-statistic for coefficients
tval <- r$coef/summary(r)$coefs[,2]
#p-value for coefficients
pval <- 2 * pt(q = abs(tval), df = 630-8, lower.tail = FALSE)


#index for lower triangular matrix
ind.lt<-lower.tri(Y)

#creating design matrix X for marginal model
X<-cbind(1,c(seniority.sum[ind.lt]),c(practice.sum[ind.lt]),c(practice.sim[ind.lt]),c(gender.sim[ind.lt]),c(office.sim[ind.lt]))


datafr<-data.frame(Y=c(Y[ind.lt]),seniority.sum=X[,2],practice.sum=X[,3],
practice.sim=X[,4],gender.sim=X[,5],office.sim=X[,6])


#fit a standard GLM
r.glm<-glm( Y ~ seniority.sum + practice.sum + practice.sim + gender.sim + office.sim, data=datafr,family=binomial(link = "logit"))


#predicted marginal probabilities
piY.glm<-r.glm$fitted.values
#beta coefficients
beta.glm<-r.glm$coefficients

#Y is now vector of Ny binary observations
Y.obs<-c(Y[ind.lt])



l.beta<-length(beta.glm)
#1 network statistic describing joint distribution
l.p<-1





# load parameters from fitted MERGM, i.e. Psi.old (nuisance parameters), beta.old (marginal model) and theta.old (ERGM parameters)
load(file="Beta_Theta_Psi.RData")
# these are the same parameters as in the Manuscript

# alternatively one might start with theta.old<-c(0,1) and beta.old<-beta.glm and Psi.old<-X%*%beta.glm,
# which would unrealistically mean that dyads would be independent





beta.new<-beta.old
Psi.new<-Psi.old
theta.new<-theta.old


#ERGM formula without the Psi
#ESP is the edgewise shared partner statistic, ESP obs corresponds to Z^obs
ESP.delta.obs<-get.shared.partner.own(as.matrix.network(work),n)
ESP.obs<-summary(work~esp(1:(n-2)))






     #M: MCMCsamplesize
     MCMCsamplesize<-10000 #use e.g. 30,000 if more accuracy required (but slower)
     M<-MCMCsamplesize
     MCMCstep<-1e3  # MCMC step size or thinning factor, increase to obtain higher independence between sampled networks (but slower)
     
     
    ############################################
    ## BEGIN STEP 2
    #############################################

     
     beta.old<-beta.new
     Psi.old<-Psi.new
     theta.old<-theta.new

    # gwesp tau is theta_1 - GWESP coefficient or natural parameter
    gwesp.tau.old<-theta.old[1]
    # gwesp alpha is theta_2 - GWESP rate parameter
    gwesp.alpha.old<-theta.old[2]

    mu.X.beta<-expit(X%*%beta.old)


    # Now follows Step 2 of the main algorithm
    # don't use Psi.old when change in beta and gwesp is too large
    # Options  fit="NR","direct" or "simple"
    # NR refers to Newton-Raphson, but it is the same as non-linear least squares in this instance
    # The other two options, "direct" or "simple" don't work properly at the moment, so just leave "NR"

    # also only applies to the ERGM model with GWESP distribution
    # The MCMC sampler - own function "get.MCMC" -  uses an internal ergm function called "ergm.getMCMCsample"
    # our experience is that it uses far less time compared to "simulate.ergm"

    #to adjust some tuning parameters have a look at "solve.Psi.GWESP"
    #STEP 2 takes a long time, usually up to iter.max MCMC samples are generated, each  takes between 2-5 minutes
    
    
    hilf<-solve.Psi.GWESP(X,beta.hat=beta.old,gwesp.tau=gwesp.tau.old,gwesp.alpha=gwesp.alpha.old,ESP=ESP.delta.obs,Y.obs,nw0=work,Psi0=Psi.old,fit="NR",MCMCsamplesize=M,iter.max=30,MCMCstep)

    #some output is always stored in  "SolvePsi.RData"  in case program is interrupted

    # END STEP 2

    #hilf is a list with the following elements:
    #list(Z.Psi=Z.Psi,Y.MCMC=Y.MCMC,eta0=eta0,Z=Z,Clist=Clist1,MCMCparams=MCMCparams1,eta01=eta01)

    ## POSTPROCESSING INFORMATION FROM STEP 2 NEEDED FOR OBTAINING NEW BETA AND THETA

    Y.MCMC<-hilf$Y.MCMC
    Z.Psi.old<-hilf$Z.Psi
    Z<-hilf$Z
    Z.obs<-hilf$Z.obs
    w<-hilf$w
    Psi.new<-hilf$Psi.new
    U<-hilf$U




    #compute ESP for MCMC sample

    ESP.MCMC<-matrix(0,MCMCsamplesize,n-2)
    start.time<-proc.time()[3]
    for(i in 1:MCMCsamplesize){
    Y.matrix.i<-Psi.matrix(Y.MCMC[i,],n)
    ESP.MCMC[i,]<-summary(as.network.matrix(Y.matrix.i,directed=F) ~ esp(1:(n-2)))
    if(i%%200==0 | i==1){end.time<-show.elapsed.time(start.time,string=paste("i: ",i, "out of ",MCMCsamplesize,sep=""))}
    }#end  for(i in 1:MCMCsamplesize){




    # calculate terms needed for the ML equations

    hilf1<-cov.w(ESP.MCMC,Y.MCMC,w)
    hilf2<-mean.var.w(ESP.MCMC,w,binary=FALSE)
    hilf3<-mean.var.w(Y.MCMC,w,binary=TRUE)

    #mean of ESP, i.e. E(Z)
    mu.ESP<-hilf2$m.Y
    #Cov(Z)
    C22<-cov.stats<-hilf2$cov.Y

    #EY
    mu.Y<-hilf3$m.Y
    #Cov(Y)
    C11<-cov.Y<-hilf3$cov.Y
    v.Y<-hilf3$v.Y
    #v.Y<-mu.Y*(1-mu.Y)



    v.Y.sqrt<-sqrt(v.Y)
    v.Y.sqrt.inv<-1/v.Y.sqrt

    #v.Y.sqrt.Xbeta<-sqrt(c(mu.X.beta)*(1-c(mu.X.beta)))
    #v.Y.sqrt.inv.Xbeta<-1/sqrt(c(mu.X.beta)*(1-c(mu.X.beta)))

    R11<- v.Y.sqrt.inv*t(v.Y.sqrt.inv*C11) #==diag(v.Y.sqrt.inv)%*%cov.Y%*%diag(v.Y.sqrt.inv)
    R11.inv<-solve(R11)
    #C11     <-  v.Y.sqrt*t(v.Y.sqrt*R11)
    C11.inv <-  v.Y.sqrt.inv*t(v.Y.sqrt.inv*R11.inv)
    #C11.inv.copy<-solve(C11)

    #Cov(Y,Z)
    C12<-cov.stats.Y<-hilf1$cov.Y1.Y2




    #C11.inv.Y.muY<-C11.inv%*%(Y.obs-mu.X.beta)
    C11.inv.Y.muY<-C11.inv%*%(Y.obs-mu.X.beta)

    D<-1/v.Y; #D = d eta / d mu = log(mu/(1-mu))=1/(mu*(1-mu))
    D.inv<-v.Y
    A.inv<-v.Y.sqrt.inv
    res<-Y.obs-mu.Y;


    #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

    #instead we re-express it as iterative re-weighted least squares, see Preisser and Qaqish (1996) "Deletion Diagnostics for Generalised Estimating Equations" Biometrika, 83(3): 551-562.
    #both are equivalent
    E<-D*res; #res=marg - mu
    #steps-size for beta is effectively 0.1
    Z.pseudo<- X%*%beta.old + 0.1*E;
    W.inv<-diag(D.inv*A.inv)%*%R11.inv%*%diag(D.inv*A.inv)
    beta.new<-solve(t(X)%*%W.inv%*%X)%*%t(X)%*%W.inv%*%Z.pseudo
    
    
    
    #CALCULATE NEW THETA
    #GR is d eta/d theta and stands for gradient
    GR<-rbind(gwesp(gwesp.tau=1,gwesp.alpha=gwesp.alpha.old,n),gwesp.der(gwesp.tau=gwesp.tau.old,gwesp.alpha=gwesp.alpha.old,n))






    theta.info<- GR%*%(C22 - t(C12)%*%C11.inv%*%C12)%*%t(GR)


    delta.theta <- solve(theta.info)%*%GR%*%(ESP.obs-mu.ESP - t(C12)%*%C11.inv.Y.muY)
    #step size for theta here is 0.1
    theta.new <- theta.old + 0.1*delta.theta

    #usually delta.theta is larger than delta.beta, increase  MCMCsamplesize to get smaller delta.theta

    #compute standard errors
     X.D.inv.C11.inv <- t(X)%*%diag(D.inv)%*%C11.inv
     C21.C11.inv <- t(C12)%*%C11.inv
     l.theta<-length(theta.new)
     Wbeta<-matrix(0,M,l.beta)
     Wtheta<-matrix(0,M,l.theta)
     start.time<-proc.time()[3]
     for(m in 1:M){
      Wbeta[m,]<-    X.D.inv.C11.inv%*%(Y.obs-Y.MCMC[m,])
      Wtheta[m,] <- GR%*%(ESP.obs-ESP.MCMC[m,]-C21.C11.inv%*%(Y.obs-Y.MCMC[m,]))
     if(m%%200==0 | m==1){end.time<-show.elapsed.time(start.time,string=paste("m: ",m, "out of ",M,sep=""))}
     }

      hilf.auto<-auto.cov(cbind(Wbeta,Wtheta),40)
      V<-hilf.auto$cov.Y





      hilf.auto1<-acf(cbind(Wbeta,Wtheta),lag.max=40,type="covariance",plot=F)
      V1<-hilf.auto1$acf[1,,]
      hilf<-2*tensor( matrix(1,1,40),hilf.auto1$acf[2:41,,] , alongA = 2, alongB = 1)
      V1<-V1 + hilf[1,,]

     Info.inv<- adiag(solve(t(X)%*%W.inv%*%X),solve(theta.info))
     Info.inv %*% V  %*% Info.inv

     #standard errors accounting for MCMC error
     se.MCMC<-sqrt(diag(Info.inv %*% V  %*% Info.inv))

     #standard errors without accounting for MCMC error
     se<-sqrt(diag(Info.inv))

      Z.MCMC<-c(beta.new,theta.new)/se.MCMC
      p.MCMC<-2*(1-pnorm(abs(Z.MCMC)))

      Z.st<-c(beta.new,theta.new)/se
      p.st<-2*(1-pnorm(abs(Z.st)))

     results<-cbind(c(beta.new,theta.new),Z.MCMC,se.MCMC,p.MCMC,Z.st,se,p.st)
     rownames(results)<-c("Intercept","seniority.sum","practice.sum","practice.sim","gender.sim","office.sim","gwesp.tau","gwesp.alpha")
     colnames(results)<-c("estimates","Z stat MMCM"," s.e. MCMC","p-value MCMC","Z","s.e.","p-value")

     print(results)

     gwesp.tau.new<-theta.new[1]
     gwesp.alpha.new<-theta.new[2]

     #obtain inaccurate estimate of log-likelihood for better estimates, see Hunter and Handcock (2006)

     Z.Psi.new<-get.updated.Psi.statistic(Y.MCMC=Y.MCMC,n,Psi.new=Psi.new)
     hilf<-obtain.weights(Z.eta=Z,eta.new=gwesp.tau.new,eta0=gwesp.tau.new,Z.Psi.new=Z.Psi.new,Z.Psi0=Z.Psi.new)
     U<-hilf$U

     ratio.0.eta.old<-ERGM.ratio(U,Z.obs,eta.new=gwesp.tau.new,eta.old=0,Y.obs,Psi.new=Psi.new,Psi.old=rep(0,Ny),Z)
     ratio.eta.new.eta.old<-ERGM.ratio(U,Z.obs,eta.new=gwesp.tau.new,eta.old=gwesp.tau.new,Y.obs,Psi.new=Psi.new,Psi.old=Psi.new,Z)
     lik<-ratio.eta.new.eta.old - ratio.0.eta.old - Ny*log(2)

     #Plot Z-statisticsto check model degeneracy
     plot(Z,type="l")


#NOW go back to BEGIN STEP 2



     
     

