## (c) Simon N Wood
## occupancy likelihood corrected
## Version allowing the generation time to be manipulated while keeping the
## disease progression times fixed.

lfun <- function(theta,link) {
## applies link to theta...
  if (is.matrix(link)) {
    a1 <- link[,2];a2 <- link[,3];link <- link[,1]
  } else {
    a1 <- rep(0,length(theta));a2 <- rep(1,length(theta))
  }
  ii <- link==1
  if (length(ii)) theta[ii] <- log(theta[ii])
  ii <- link==2
  if (length(ii)) theta[ii] <- binomial()$linkfun((theta[ii]-a1[ii])/(a2[ii]-a1[ii]))
  theta
}

ilink <- function(eta,link) {
## apply link inverse and compute dtheta/deta - eta is the unconstrained parameter
  if (is.matrix(link)) {
    a1 <- link[,2];a2 <- link[,3];link <- round(link[,1])
  } else {
    a1 <- rep(0,length(eta));a2 <- rep(1,length(eta))
  }
  theta <- eta
  dtheta <- eta*0+1
  ii <- link==1
  if (length(ii)) {
    dtheta[ii] <- theta[ii] <- exp(eta[ii])
  }
  ii <- link==2
  if (length(ii)) {
    theta[ii] <- binomial()$linkinv(eta[ii])*(a2[ii]-a1[ii]) + a1[ii]
    dtheta[ii] <- binomial()$mu.eta(eta[ii])*(a2[ii]-a1[ii])
  }
  list(theta=theta,dtheta=dtheta)
}

ll <- function(theta,gamma,S0,iaux,daux,dat,Sb=NULL,link=rep(0,length(theta)),plot=FALSE,
               step=1,nbk=rep(2,5),getR=FALSE,lab="",xlab="day") {
## log likelihood for the rep41 covid SEIR+hospital model.  nbk is a vector of overdispersion parameters.
## rep41 uses neg bin kappa = 2 for all - this is very hard to justify for deaths.
  nb <- iaux[1];np <- iaux[2];nc <- iaux[3];ns <- iaux[4]
  nx <- ns*nc*(1+np);
  nyo <- 14
  ny <- nyo*(1+np)
  x <- rep(0,nx)
  for (i in 1:nc) x[(i-1)*ns+1] <- S0[i]
  ii <- 1:length(gamma)+(nc-1)^2
  daux[ii] <- gamma
  n <- max(dat$day)
  th <- ilink(theta,link)
  thetaw <- theta ## original scale
  theta <- th$theta ## parameters on dynamic model specification scale
  if (any(!is.finite(theta))) return(list(l=NA,dl=NA))
  vm <- .C("RK4",yout = as.double(rep(0,ny*n)),xout=as.double(rep(0,nx*n)),x=as.double(x),
          theta=as.double(theta),h=as.double(1/step),t0=as.double(0),n=as.integer(n),nx=as.integer(nx),ny=as.integer(ny),
          ostep =as.integer(step),work=as.double(rep(0,5*nx)),iaux=as.integer(iaux),daux=as.double(daux))
  yout <- matrix(vm$yout,ny,n)

  ## hospital deaths...

  k <- nbk[4]
  y <- dat$death3;
  ynb <- matrix(NA,max(dat$day),7);munb <- matrix(NA,max(dat$day),9)
  ii <- !is.na(y); y<- y[ii];ii <- dat$day[ii]
  mu <- yout[4,ii]
  ynb[ii,4] <- y;munb[ii,4] <- mu
  n.death <- sum(mu)
  ld <- sum(k*log(k/(k+mu)) + y*log(mu/(mu+k)) + lgamma(k+y) - lgamma(k) - lgamma(y+1))
  dld.dmu <- y/mu - (y+k)/(mu+k)
  dld <- drop(yout[1:np*nyo+4,ii]%*%dld.dmu)
  if (plot&&length(y)) {
    ylim=c(range(y,mu));xlim <- range(ii)
    maj <- 2.2;mas <- .7
    pch=19;pex=.4
    plot(ii,y,ylab="",col="grey",ylim=ylim,xlim=xlim,xlab="",pch=pch,cex=pex);
    mtext("deaths",2,maj,cex=mas);mtext(xlab,1,maj,cex=mas)
    lines(ii,mu)
    text(220,ylim[2]*.86,lab)
  }

  ## care home deaths...

  k <- nbk[5]
  y <- dat$death2-dat$death3;ii <- !is.na(y); y<- y[ii];ii <- dat$day[ii]
  mu <- yout[5,ii];n.death <- n.death+ sum(mu)
  ynb[ii,5] <- y;munb[ii,5] <- mu
  lg <- sum(k*log(k/(k+mu)) + y*log(mu/(mu+k)) + lgamma(k+y) - lgamma(k) - lgamma(y+1))
  dlg.dmu <- y/mu - (y+k)/(mu+k)
  dlg <- drop(yout[1:np*nyo+5,ii]%*%dlg.dmu)
  if (plot&&length(y)) {
    points(ii,y,col=2,pch=pch,cex=pex);lines(ii,mu,col=4)
  }

  ## occupancy...

  k <- 100000

  ## total occupancy...
  y0 <- dat$phe_patients;i0 <- dat$day
  ii <- !is.na(y0);y <- sum(y0[ii]);i0 <- i0[ii]
  mu <- sum(yout[2,i0])

  ## log likelihood total occupancy....
  lot <- sum(k*log(k/(k+mu)) + y*log(mu/(mu+k)) + lgamma(k+y) - lgamma(k) - lgamma(y+1))
  dlot.dmu <- y/mu - (y+k)/(mu+k)
  dlot <- rowSums(yout[1:np*nyo+2,i0]) * dlot.dmu

  ## get change in occupancy over the day - difference between in and outflow
  k <- nbk[2] ## overdispersion parameter
  y <- diff(dat$phe_patients);ii <- !is.na(y); y <- y[ii];ii <- dat$day[ii]
  mu1 <- yout[11,ii] # general ward inflow
  mu2 <- yout[12,ii] # general ward outflow
  wi <- rep(1:ceiling(length(y)/7),each=7)[1:length(y)]
  y <- tapply(y,wi,sum) ## aggregate differences by week to avoid unmodelled weekly discharge cycle
  mu1 <- tapply(mu1,wi,sum);mu2 <- tapply(mu2,wi,sum)

  iw <- ii[1]+1:length(y)*7 - 4 ## output to allow over-dispersion update
  ynb[iw,6] <- y;munb[iw,6] <- mu1;munb[iw,7] <- mu2;
  alpha <- y-mu1+mu2;sig <- mu1+mu2

  lo <- sum(-alpha^2/(2*k*sig) - log(k*sig)/2) ## log likelihood occupancy change
  dlo.dmu1 <- alpha/(k*sig) + alpha^2/(2*k*sig^2) - .5/sig
  dlo.dmu2 <- -alpha/(k*sig) + alpha^2/(2*k*sig^2) -.5/sig
  dlo <- t(apply(yout[1:np*nyo+11,ii],1,function(y,wi) tapply(y,wi,sum),wi=wi))%*%dlo.dmu1 +
	 t(apply(yout[1:np*nyo+12,ii],1,function(y,wi) tapply(y,wi,sum),wi=wi))%*%dlo.dmu2
  lo <- lo + lot;
  dlo <- drop(dlo) + dlot

  ## occupancy itself for plotting and return..
  y <- dat$phe_patients;ii <- !is.na(y); y<- y[ii];ii <- dat$day[ii]
  mu <- yout[2,ii]
  ynb[ii,2] <- y;munb[ii,2] <- mu

  if (plot&&length(y)) {
    plot(ii,y,col="grey",ylim=range(c(y,mu)),xlim=xlim,ylab="",xlab="",pch=pch,cex=pex)
    mtext("occupancy",2,maj,cex=mas);mtext(xlab,1,maj,cex=mas)
    lines(ii,mu)
  }

  ## ICU occupancy...
  k <- 20000

  y0 <- dat$phe_occupied_mv_beds;i0 <- dat$day
  ii <- !is.na(y0);y <- sum(y0[ii]);i0 <- i0[ii]
  mu <- sum(yout[3,i0])
  ## log likelihood total icu occupancy
  lut <- sum(k*log(k/(k+mu)) + y*log(mu/(mu+k)) + lgamma(k+y) - lgamma(k) - lgamma(y+1))
  dlut.dmu <- y/mu - (y+k)/(mu+k)
  dlut <- rowSums(yout[1:np*nyo+3,i0]) * dlut.dmu

  ## get likelihood for icu occupancy changes...
  k <- nbk[3] ## overdispersion parameter
  y <- diff(dat$phe_occupied_mv_beds);ii <- !is.na(y); y<- y[ii];ii <- dat$day[ii]
  mu1 <- yout[13,ii] # ICU inflow
  mu2 <- yout[14,ii] # ICU outflow
  wi <- rep(1:ceiling(length(y)/7),each=7)[1:length(y)]
  y <- tapply(y,wi,sum) ## aggregate differences by week to avoid unmodelled weekly discharge cycle
  mu1 <- tapply(mu1,wi,sum);mu2 <- tapply(mu2,wi,sum)

  iw <- ii[1]+1:length(y)*7 - 4
  ynb[iw,7] <- y;munb[iw,8] <- mu1;munb[iw,9] <- mu2;

  alpha <- y-mu1+mu2;sig <- mu1+mu2
  lu <- sum(-alpha^2/(2*k*sig) - log(k*sig)/2) ## log likelihood - icu occupancy change
  dlu.dmu1 <- alpha/(k*sig) + alpha^2/(2*k*sig^2) - .5/sig
  dlu.dmu2 <- -alpha/(k*sig) + alpha^2/(2*k*sig^2) -.5/sig
  dlu <- t(apply(yout[1:np*nyo+13,ii],1,function(y,wi) tapply(y,wi,sum),wi=wi))%*%dlu.dmu1 +
	 t(apply(yout[1:np*nyo+14,ii],1,function(y,wi) tapply(y,wi,sum),wi=wi))%*%dlu.dmu2
  #dlu <- drop(yout[1:np*nyo+13,ii]%*%dlu.dmu1)+drop(yout[1:np*nyo+14,ii]%*%dlu.dmu2)
  lu <- lu + lut
  dlu <- dlu + dlut
  # ICU occupancy itself for plotting and return
  y <- dat$phe_occupied_mv_beds;ii <- !is.na(y); y<- y[ii];ii <- dat$day[ii]
  mu <- yout[3,ii]
  ynb[ii,3] <- y;munb[ii,3] <- mu
  if (plot&&length(y)) {
    points(ii,y,col=2,pch=pch,cex=pex)#xlab="day",ylab="ICU",col="grey",ylim=range(c(y,mu)));
    lines(ii,mu,col=4)
  }

  ## hospitalization...

  y <- dat$phe_admissions; ii <- !is.na(y); y<- y[ii]
  ii <- dat$day[ii] ## cols of yout corresponding to observed data
  k <- nbk[1] ## theta parameter
  mu <- yout[1,ii]
  ynb[ii,1] <- y;munb[ii,1] <- mu
  lh <- sum(k*log(k/(k+mu)) + y*log(mu/(mu+k)) + lgamma(k+y) - lgamma(k) - lgamma(y+1))
  dlh.dmu <- y/mu - (y+k)/(mu+k)
  dlh <- drop(yout[1:np*nyo+1,ii]%*%dlh.dmu)
  #if (plot) par(mfrow=c(2,3));
  if (plot&&length(y)) {
    plot(ii,y,col="grey",ylim=range(c(y,mu)),xlim=xlim,ylab="",xlab="",pch=pch,cex=pex);
    mtext("admissions",2,maj,cex=mas);mtext(xlab,1,maj,cex=mas)
    lines(ii,mu)
  }

  ## serology samples...

  N <- sum(S0[4:13])
  y <- dat$n_positive;ii <- !is.na(y); y<- y[ii];
  n <- dat$total_samples[ii]
  ii <- dat$day[ii]
  mu <- yout[6,ii]/N
  ls <- sum(dbinom(y,n,mu,log=TRUE))
  dls.dmu <- y/mu - (n-y)/(1-mu)
  dls <- drop(yout[1:np*nyo+6,ii]%*%dls.dmu)/N
  if (plot&&length(y)) {
    plot(ii,y/n,ylab="",ylim=range(c(y/n,mu)),col="grey",xlim=xlim,xlab="",pch=pch,cex=pex)
    mtext("P(sero)",2,maj,cex=mas);mtext(xlab,1,maj,cex=mas)
    lines(ii,mu,col=1)
  }

  ## REACT PCR samples...

  N <- sum(S0[2:18])
  y <- dat$react_pos;ii <- !is.na(y); y<- y[ii];
  n <- dat$react_samples[ii]
  ii <- dat$day[ii]
  mu <- yout[7,ii]/N
  lr <- sum(dbinom(y,n,mu,log=TRUE))
  dlr.dmu <- y/mu - (n-y)/(1-mu)
  dlr <- drop(yout[1:np*nyo+7,ii]%*%dlr.dmu)/N
  if (plot&&length(y)) {
    plot(ii,y/n,ylab="",col="grey",xlim=xlim,xlab="",pch=pch,cex=pex);lines(ii,mu,col=1)
    mtext("P(pcr)",2,maj,cex=mas);mtext(xlab,1,maj,cex=mas)
  }

  l=lh+lo+lu+ld+lg+lr+ls ## total log likelihood
  dl=dlh+dlo+dlu+dld+dlg+dlr+dls ## corresponding derivative vector

  if (!is.null(Sb)) { ## add smoothing penalty
    Sb <- drop(Sb %*% theta[1:nb])
    l <- l - sum(theta[1:nb]*Sb)/2
    dl[1:nb] <- dl[1:nb] - Sb
  }
  R <- yout[9,]; R[R<0] <- NA; Rd <- NULL
  if (plot||getR) { ## get the Diekmann version of R...
    gamma_A <- gamma[2];gamma_C <- gamma[3];p_c <- gamma[4]
    Rd <- R
    beta <- yout[10,]
    C <- matrix(daux[1:(nc-1)^2],nc-1,nc-1)[-(nc-1),-(nc-1)]
    n <- max(dat$day)
    xout <- matrix(vm$xout,nx,n)[1:(nc*ns),]
    S <- xout[0:(nc-3)*ns+1,]
    for (i in 1:n) {
      C0 <- S[,i]*C
      K <- beta[i]*rbind(cbind((1-p_c)*C0/gamma_A,(1-p_c)*C0/gamma_C),cbind(p_c*C0/gamma_A,p_c*C0/gamma_C))
      Rd[i] <- abs(eigen(K,only.values = TRUE)$values[1])
    }
  } else xout <- NULL
  if (ncol(link)>3) { ## prior on dynamic params too
     ii <- (nb+1):np
     sig.p <- link[ii,5]
     l <- l - sum((thetaw[ii]-link[ii,4])^2/(2*sig.p^2))
     dprior <- c(rep(0,nb),-(thetaw[ii]-link[ii,4])/sig.p^2)
  } else dprior <- 0
  list(l=l,dl=dl*th$dtheta+dprior,R=R,Rd=Rd,fi=yout[8,],xout=xout,n.death=n.death,munb=munb,ynb=ynb)
} ## ll

thfix <- function(theta,fixi,fixv) {
## if fixi!=NULL then it indexes parameters that have been dropped from the optimization of
## theta, while fixv gives their fixed values. This routine inserts the fixed values into
## theta in the correct place...
  if (!is.null(fixi)) {
    if (length(fixi)!=length(fixv)) stop("fixv wrong length")
    th1 <- rep(0,length(theta)+length(fixi))
    ii <- rep(FALSE,length(th1))
    ii[fixi] <- TRUE
    th1[ii] <- fixv
    th1[!ii] <- theta
    theta <- th1
  }
  theta
} ## thfix

nll <- function(theta,gamma,S0,iaux,daux,dat,Sb,link,step=1,nbk=rep(2,5),fixi=NULL,fixv=NULL) {
## if fixi!=NULL then it indexes parameters that have been dropped from the optimization of
## theta, while fixv gives their fixed values.
 iaux[7] <- 0; ## no derivatives needed
 theta <- thfix(theta,fixi,fixv)
 -ll(theta,gamma,S0,iaux,daux,dat,Sb=Sb,link=link,plot=FALSE,step=step,nbk=nbk)$l
}

nll0 <- function(theta,gamma,S0,iaux,daux,dat,Sb,link,step=1,nbk=rep(2,5),fixi=NULL,fixv=NULL) {
## if fixi!=NULL then it indexes parameters that have been dropped from the optimization of
## theta, while fixv gives their fixed values.
 iaux[7] <- 0; ## no derivatives needed
 theta <- thfix(theta,fixi,fixv)
 ll(theta,gamma,S0,iaux,daux,dat,Sb=Sb,link=link,plot=FALSE,step=step,nbk=nbk)
}

nllg <- function(theta,gamma,S0,iaux,daux,dat,Sb,link,step=1,nbk=rep(2,5),fixi=NULL,fixv=NULL) {
## grad function corresponding to nll
 theta <- thfix(theta,fixi,fixv)
 g <- -ll(theta,gamma,S0,iaux,daux,dat,Sb=Sb,link=link,plot=FALSE,step=step,nbk=nbk)$dl
 if (!is.null(fixi)) g <- g[-fixi]
 g
}

hessian <- function(theta,gamma,S0,iaux,daux,dat,Sb,link,step=1,nbk=rep(2,5),fixi=NULL,fixv=NULL) {
  ## compute Hessian by finite differencing
  nt <- length(theta)
  eps <- 1e-7
  g0 <- nllg(theta,gamma,S0,iaux,daux,dat,Sb,link,step,nbk,fixi,fixv)
  H <- matrix(0,nt,nt)
  for (i in 1:nt) {
    theta1 <- theta;theta1[i] <- theta[i] + eps
    g1 <- nllg(theta1,gamma,S0,iaux,daux,dat,Sb,link,step,nbk,fixi,fixv)
    H[,i] <- (g1-g0)/eps;
  }
  .5*(t(H)+H)
} ## hessian

nll1 <-  function(theta,gamma,S0,iaux,daux,dat,Sb,link,step=1,nbk=rep(2,5),fixi=NULL,fixv=NULL) {
## unused nlm version
 theta <- thfix(theta,fixi,fixv)
 l0 <- ll(theta,gamma,S0,iaux,daux,dat,Sb=Sb,link=link,plot=FALSE,step=step,nbk=nbk)
 nll <- -l0$l
 attr(nll,"gradient") <- -l0$dl
 nll
}

ginv <- function(V,warn=FALSE) { ## generalized inverse
  if (!all(is.finite(V))) return(NULL)
  ev <- eigen(V,symmetric=TRUE)
  d <- ev$values;ii <- d>0
  d[ii] <-1/d[ii];d[!ii] <- 0
  if (warn&&sum(!ii)) warning("indefinite V\n")
  ev$vectors%*%(d*t(ev$vectors))
}

tf <- function(x,a1=.00,a2=.1) { ## transform as applied to b(t)
  ii <- x<0
  x[ii] <- exp(x[ii])
  x[!ii] <- exp(-x[!ii])
  x[ii] <- a1 + (a2-a1)*x[ii]/(1+x[ii])
  x[!ii] <- a1 + (a2-a1)*1/(1+x[!ii])
  x
}

itf <- function(y,a1=.00,a2=.1) { ## transform as applied to b(t)
  y <- (y-a1)/(a2-a1)
  log(y/(1-y))
}

setup.model <- function(location="london",Xtype=2,beta.off = .25,rep41C=FALSE) {
## sets up the information required for a model fit, including fixed parameters,
## data, initial parameter values and limits.

  source("params.r") ## some fixed parameters

  region <- c("east_of_england","london","midlands","north_east_and_yorkshire",
  "north_west","south_east","south_west")

  if (!location %in% region) stop("location wrong")

  source("data.r",local=TRUE) ## obtain the fit data and some other parameters

  load(paste("C_",location,".rda",sep=""))
  ## or use m from sircovid package - problem is that it only matches what is documented
  ## up to age 70 (parameters change but results unchanged)...
  if (rep41C) C <- ch$m[1:18,1:18]

  n <- max(rtr$day) ## number of output steps
  step <- 1 ## steps between output
  h <- 1/step ## step length
  hx <- h/2 ## step at which basis evaluation required
  t <- h*0:(2*step*n)/2 ## when basis evaluation required
  nb <- 80 ## number of b(t) coeffs
  if (Xtype==1) {## adaptive
    t0 <- 45 ## constant b(t) before here
    ii <- t >= t0 ## set up adaptive smooth from t0
    smb <- smoothCon(s(t,k=nb,bs="ad"),data=data.frame(t=t[ii]))[[1]]
    ni <- sum(!ii) ## need to pad model matrix to fill in constant prior to t0
    X0 <- matrix(smb$X[1,],ni,ncol(smb$X),byrow=TRUE)
    Xb <- rbind(X0,smb$X) ## b(t) basis matrix
    Sb <- smb$S ## penalty matrix
    attr(Sb,"rank") <- smb$rank
  } else if (Xtype==2) { ## force b(t) to be constant at start
    ne <- 5
    smb <- smoothCon(s(t,k=nb+ne,bs="bs"),data=data.frame(t=t))[[1]]
    T <- diag(nb+ne)[,-(1:ne)];T[1:ne,1] <- 1
    Xb <- smb$X%*%T ## b(t) basis matrix
    Sb <- t(T)%*%smb$S[[1]]%*%T ## penalty matrix
  } else { ## rep 41
    nb <- length(bk)
    Xb <- matrix(0,length(t),nb)
    Xb[,1] <- as.numeric(t<=bk[1])
    for (i in 1:nb) {
      yk <- bk*0;yk[i] <- 1
      Xb[,i] <- approx(bk,yk,t,rule=2)$y
    }
    Sb <- NULL
  }

  b0 <- c(rep(-2.4,5),-2.404542,-2.468239,-2.676757,-3.068759,-3.510927,
  -3.833448,-3.971020,-3.999736,rep(-4,10),-3.999977,-3.987959,-3.920592,
  -3.752347,-3.506797,-3.277703,-3.147313,-3.104178,rep(-3.5,14))
  t0 <- 0:44*7.5+7
  b0 <- exp(approx(t0,b0,t,rule=2)$y)
  a1 <- 0; a2 <- 0.1
  b0[b0>a2] <- a1 + (a2-a1)*.95
  b0[b0<a1] <- a1 + (a2-a1)*.05
  b0 <- itf(b0,a1,a2)
  beta <- coef(lm(b0~Xb-1)) ## initial b(t) parameters
  beta <- beta - beta.off

  ## set up param vector
  eps <- 0.1; m_chw <- 4e-6; m_chr <- 4e-6
  t0 <- 20;p_g <- 0.6;p_h_max <- 0.45
  mu_ic <- .9; mu_d <- .9; p_ic_max <- .3
  p_icd_max <- .6;p_wd_max <- .3;p_hd_max <- .37
  gamma_G <- .4;
  gamma_IC_Wr <- .0641;gamma_Wr <- .16;gamma_Hr<- 0.0935

  theta <- c(beta,eps,m_chw,m_chr,t0,p_h_max,p_g,mu_ic,p_ic_max,p_icd_max,mu_d,p_wd_max,p_hd_max,gamma_G,
             gamma_IC_Wr,gamma_Wr,gamma_Hr)
  names(theta)[(nb+1):length(theta)] <- c("eps","m_chw","m_chr","t0", "p_h_max","p_g","mu_ic","p_ic_max",
                             "p_icd_max","mu_d","p_wd_max","p_hd_max","gamma_G","gamma_IC_Wr","gamma_Wr","gamma_Hr")

  link <- c(rep(0,nb),1,1,1,0,2,2,2,2,2,2,2,2,1,1,1,1) ## link types: 0 - I(), 1 - log, 2 - logit
  a1 <-   c(rep(0,nb),0,0,0,0,0,0,.66,.14,.57,.66,.25,.36,0,0,0,0) ## lower logistic limit
  a2 <-   c(rep(0,nb),0,0,0,0,1,1, 1 ,.36,.77,  1,.46,.56,0,0,0,0)  ## upper logistic limit
  ### prior mean on working scale...
  mup <- c(rep(0,nb),-.7,-12.7,-12.7,30,1.1,1.1,-1,-.2,0,-1,0,0,-1,-2.75,-1.8,-2.37)
  sdp <- c(rep(NA,nb),2,2,2,10,2,2,2,2,2,2,1,1,2,2,2,2) ## prior s.d. on working scale
  link <- cbind(link,a1,a2,mup,sdp)
  ## fixed param vector
  p_c <- 0.6 # .6 ## Knock et al
  inoc_sig <- 7.5 # arbitrary
  p_star <- 0.25 ## Knock et al
  f_start <- 92 ## improvement start
  f_end <- 152  ## improvement end
  gamma_U <- .5 ## Table S7
  ## gamma_E, gamma_C and gamma_ph should keep mean time from infection to hospital at
  ## 12.3 days to match cited literature
  #gamma_E <- .44 ## Table S2 or use ch$gamma_E based on early lauer paper
  #gamma_E <- 1/2.9 ## based on McAloon meta-analysis!
  gamma_E <- 1/1.5 ## mean 3 days to infective
  #gamma_C <- .25 ## Table S2 - based on equating the median and mean of an exponential!
  #gamma_C <- 1/6.5 ## Based on cited Docherty paper *mean* for males and females (see supp)
  gamma_C <- 1/4 ## mean 4 days infective => generation time 6.2 (brought down by asymptomatics)
  gamma_ph <- 1/5.3 ## 5.3 is the length of the post I_C pre hospital phase (2.8 without correction)
  gamma_A <- ch$gamma_asympt
  gamma_IC_Wr <- .0641 ## Table S2
  gamma_IC_Wd <- .14 ## Table S2
  gamma_IC_d <- .17 ## Table S2
  gamma_Wr <- .16  ## Table S2
  gamma_Wd <- .12  ## Table S2
  gamma_Hr <- .0935 ## Table S2
  gamma_Hd <- .19 ## Table S2
  gamma_S <- 1/13 ## Table S7
  p_S <- .85 ## Table S7

  gamma0 <- c(gamma_E,gamma_A,gamma_C,p_c,inoc_sig,gamma_ph,p_star,
           f_start,f_end,ch$gamma_ICU_D,gamma_U,gamma_IC_Wr,gamma_IC_Wd,gamma_IC_d,
	   gamma_Wr,gamma_Wd,gamma_Hr,gamma_Hd,ch$gamma_PCR_pre,
           ch$gamma_PCR_pos,gamma_S,p_S)

  names(gamma0) <- c("gamma_E","gamma_A","gamma_C","p_c","inoc_sig","gamma_ph","p_star",
  "f_start","f_end","gamma_IC_pre","gamma_U","gamma_IC_Wr","gamma_IC_Wd","gamma_IC_d",
  "gamma_Wr","gamma_Wd","gamma_Hr","gamma_Hd","gamma_P","gamma_Po","gamma_S","p_S")

  gamma <- c(gamma0,psi.h,psi.ic,psi.icd,psi.wd,psi.hd) ## 19 vector
  ng <- length(gamma)
  np <- length(theta)
  nc <- 19 ## classes
  ns <- 37 ## states per class
  nk <- prod(dim(Xb))
  deriv <- 1
  iaux <- c(nb,np,nc,ns,ng,nk,deriv,-1)
  daux <- c(C,gamma,t(Xb),hx,rep(0,nb+2*(nc+nc*np)))
  S0=ch$N_tot
  thetaw <- lfun(theta,link=link);fixi <- fixv <- NULL
  arg <- list(location=location,gamma=gamma,S0=S0,thetaw=thetaw,iaux=iaux,daux=daux,dat=rtr,Sb=Sb,link=link,step=step)
  arg
} ## setup.model

nlnb <- function(theta,y,mu) {
## negative log lik for NB
  ii <- is.na(y)|is.na(mu)
  y <- y[!ii];mu <- mu[!ii]
  -sum(theta*log(theta/(theta+mu)) + y*log(mu/(mu+theta)) + lgamma(theta+y) - lgamma(theta))
}

fit.model <- function(arg) {
## fit model using Fellner Schall sp updates...
  if (is.null(arg$Sb)) { m <- 0; Sb <- NULL} else
  if (is.list(arg$Sb)) {
    lambda <- rep(10,length(arg$Sb))
    Sb <- arg$Sb[[1]]*lambda[1]
    m <- length(arg$Sb)
    if (m>1) for (i in 2:m) Sb <- Sb +  arg$Sb[[i]]*lambda[i]
    if (m==1) arg$Sb <- arg$Sb[[1]]
  } else {
    m <- 1
    Sb <- arg$Sb * 10
  }
  thetaw <- if (is.null(arg$fixi)) arg$thetaw else arg$thetaw[-arg$fixi]
  er <- optim(thetaw,nll,gr=nllg,method="BFGS",control=list(maxit=400),gamma=arg$gamma,S0=arg$S0,
            iaux=arg$iaux,daux=arg$daux,dat=arg$dat,Sb=Sb,link=arg$link,step=arg$step,
	    nbk=rep(2,5),fixi=arg$fixi,fixv=arg$fixv)
  cat(arg$location," ")

  er$par -> thetaw -> thetaw0

  if (is.null(arg$Sb)) { m <- 0; Sb <- NULL} else
  if (is.list(arg$Sb)) {
    lambda <- rep(10,length(arg$Sb))
    Sb <- arg$Sb[[1]]*lambda[1]
    m <- length(arg$Sb)
    if (m>1) for (i in 2:m) Sb <- Sb +  arg$Sb[[i]]*lambda[i]
  } else {
    m <- 1; lambda <- 10
    Sb <- arg$Sb * lambda
  }
  nrep <- if (m>0) 20 else 1 ## based on prelim run

  nbk <- c(2,2,2,2000,50) ## admissions, ward and icu occupancy, hospital deaths, care home deaths
  k1 <- c(40,100,100,2000,2000) ## upper limits on overdispersion param
  for (i in 1:nrep) {
    er0 <- er
    er <- optim(thetaw,nll,gr=nllg,method="BFGS",control=list(maxit=400),gamma=arg$gamma,S0=arg$S0,
            iaux=arg$iaux,daux=arg$daux,dat=arg$dat,Sb=Sb,link=arg$link,step=arg$step,
	    nbk=nbk,fixi=arg$fixi,fixv=arg$fixv,hessian=FALSE)
    ## get Hessian without Sb first - check +ve semidef
    H <- hessian(er$par,arg$gamma,arg$S0,arg$iaux,arg$daux,arg$dat,Sb=NULL,link=arg$link,
                 step=arg$step,nbk=nbk,arg$fixi,arg$fixv)
    eh <- eigen(H)
    if (any(eh$values<0)) { ## find nearest +ve semi definite
      eh$values[eh$values<0] <- 0
      H <- eh$vectors %*% (eh$values*t(eh$vectors))
    }
    Hp <- H ## penalized version
    if (!is.null(Sb)) {
      nb <- ncol(Sb)
      Hp[1:nb,1:nb] <- Hp[1:nb,1:nb] + Sb
    }

    ## update nb theta...
    l0 <- nll0(thetaw,gamma=arg$gamma,S0=arg$S0,iaux=arg$iaux,daux=arg$daux,
            dat=arg$dat,Sb=Sb,link=arg$link,step=arg$step)
    for (i in c(1,5)) { ## optimize log lik
      nbk[i] <- optimize(nlnb,interval=c(2,k1[i]),y=l0$ynb[,i],mu=l0$munb[,i])$minimum
    }
    for (i in 2:3) { ## update occupancy rates overdispersion
      y <- l0$ynb[,i+4];ii <- !is.na(y);y <- y[ii]
      mu1 <- l0$munb[ii,i*2+2]
      mu2 <- l0$munb[ii,i*2+3]
      k.mle <- mean((y-mu1+mu2)^2/(mu1+mu2))
      k.mle <- if (k.mle < 1) 1 else { if (k.mle>k1[i]) k1[i] else k.mle}
      nbk[i] <- k.mle
    }
    er$par -> thetaw
    thetaw <- thfix(er$par,arg$fixi,arg$fixv)
    Vb <- ginv(Hp,warn=TRUE)
    if (is.null(Vb)) break
    nb <- arg$iaux[1]
    b <- thetaw[1:nb]
    if (m>1) {
      Si <- ginv(Sb)
      if (is.null(Si)) { Vb <- NULL; break}
      for (i in 1:m) {
        fac <- drop((sum(Si*arg$Sb[[i]]) - sum(Vb[1:nb,1:nb]*arg$Sb[[i]]))/t(b)%*%arg$Sb[[i]]%*%b)
	if (fac<0.2) fac <- .2
	if (fac>5) fac <- 5
        lambda[i] <- lambda[i]*fac
      }

      Sb <- arg$Sb[[1]]*lambda[1]
      for (i in 2:m) Sb <- Sb +  arg$Sb[[i]]*lambda[i]
    } else if (m==1) {
      fac <- drop(((nb-2)/lambda - sum(Vb[1:nb,1:nb]*arg$Sb))/t(b)%*%arg$Sb%*%b)
      if (fac<0.2) fac <- .2
      if (fac>5) fac <- 5
      lambda <- lambda*fac
      Sb <- lambda*arg$Sb
    }
    if (m>0) cat(arg$region,": ",lambda,"\n")
  }
  if (is.null(Vb)) {
    er <- er0
    er$par -> thetaw
    Vb <- ginv(er$hessian)
  }
  list(thetaw=thetaw,theta=ilink(thetaw,arg$link)$theta,Vb=Vb,H=Hp,lambda=lambda,nbk=nbk)
} ## fit.model

check.model <- function(arg,plot=TRUE,deriv=FALSE) {
## checks if initial params gtive something remotely reasonable
  arg$iaux[7] <- as.numeric(deriv)
  l0 <- ll(arg$thetaw,arg$gamma,arg$S0,arg$iaux,arg$daux,arg$dat,Sb=NULL,link=arg$link,plot=plot,step=arg$step)
}

check.derivs <- function(thetaw,arg,eps = 1e-6) {
## checks derivatives of penalized log likelihood. Note currently not checking with smoothing penalty.
  l0 <- ll(thetaw,arg$gamma,arg$S0,arg$iaux,arg$daux,arg$dat,Sb=NULL,link=arg$link,plot=FALSE,step=arg$step)
  fd <- l0$dl*0
  arg$iaux[7] <- 0
  for (i in 1:length(fd)) {
    theta1 <- thetaw;theta1[i] <- theta1[i] + eps
    l1 <- ll(theta1,arg$gamma,arg$S0,arg$iaux,arg$daux,arg$dat,Sb=NULL,link=arg$link,plot=FALSE,step=arg$step)
    fd[i] <- (l1$l-l0$l)/eps
    cat(".")
  }
  cat("\n")
  plot(l0$dl,fd);abline(0,1)
  list(dl=l0$dl,fd=fd,l=l0$l)
}

dRfi <- function(arg) {
## lazy evaluation of derivatives of R and incidence trajectories w.r.t. parameters, assuming thetaw and Vb added to arg
  eps = 1e-5
  thetaw <- arg$thetaw
  l0 <- ll(thetaw,arg$gamma,arg$S0,arg$iaux,arg$daux,arg$dat,Sb=NULL,link=arg$link,plot=FALSE,step=arg$step,getR=TRUE)
  dR <- df <- matrix(0,length(l0$Rd),length(thetaw))
  for (i in 1:length(thetaw)) {
    theta1 <- thetaw;theta1[i] <- theta1[i] + eps
    l1 <- ll(theta1,arg$gamma,arg$S0,arg$iaux,arg$daux,arg$dat,Sb=NULL,link=arg$link,plot=FALSE,step=arg$step,getR=TRUE)
    dR[,i] <- (log(l1$Rd)-log(l0$Rd))/eps
    df[,i] <- (log(l1$fi)-log(l0$fi))/eps
  }
  list(dR=dR,df=df,se.r=rowSums((dR%*%arg$Vb) * dR)^.5,se.f=rowSums((df%*%arg$Vb) * df)^.5,fi=l0$fi,Rd=l0$Rd)
} ## dRfi

plot.fit <- function(fit,arg) {
## quick plot of fit to data + R and incidence.
  l0 <- ll(fit$thetaw,arg$gamma,arg$S0,arg$iaux,arg$daux,arg$dat,Sb=NULL,link=arg$link,plot=TRUE,step=arg$step)
  X11()
  par(mfrow=c(2,2))
  plot(log(l0$Rd),xlim=c(50,150));abline(0,0);abline(v=84)
  plot(log(l0$Rd));abline(1,0);abline(v=84)
  plot(l0$Rd,xlim=c(50,150));abline(1,0);abline(v=84)

  plot(l0$fi,xlim=c(50,150));abline(v=84)

} ## plot.fit

###############################################################################
setwd("~sw283/research.notes/covid19/rep41/code/albus") ## Data & Code location
###############################################################################

library(mgcv);library(parallel)
dyn.load("albus-gl.so")

region <- c("east_of_england","london","midlands","north_east_and_yorkshire",
  "north_west","south_east","south_west")

rlab <- c("East","London","Midlands","N.E. & Yorkshire", "North West", "South East", "South West")

## setup regional fits...
arg <- list()
for (i in 1:7) arg[[i]] <-setup.model(region[i],Xtype=1,beta.off=.75,rep41C=FALSE)
names(arg) <- rlab
load("thetaw0-g.rda")
for (i in 1:7) {
  arg[[i]]$region <- rlab[i]
  arg[[i]]$thetaw <- thetaw0 ## initialize from old fit to midlands to speed up
}

## par(mfrow=c(2,3));check.model(arg[[1]])
## plot.fit(er,arg[[1]])
## system.time(fd <- check.derivs(arg[[1]]$thetaw,arg[[1]]));fd$dl[-(1:80)];fd$fd[-(1:80)]

system.time(er <- mclapply(arg,fit.model,mc.cores=7))

for (i in 1:7) { ## add fit results to arg list
  arg[[i]]$thetaw <- er[[i]]$thetaw
  arg[[i]]$Vb <- er[[i]]$Vb
}
## get derivatives of R and fi (incidence) w.r.t. parameters
system.time(dd <- mclapply(arg,dRfi,mc.cores=7))

# for (i in 1:7) arg[[i]]$iaux[8] <- 83 ## freeze b(t) at b(83)

## check all fits and compute total incidence and average R -
## R is number of new infections per existing infection, so
## must be weighted by existing infectives.
ns <- arg[[1]]$iaux[4]
nc <- arg[[1]]$iaux[3]
ii <- rep((0:(nc-3))*ns,each=2)+4:5 ## index of all infectives outside care homes
n.deaths <- 0

ps <- FALSE
if (ps) postscript("fits.eps",horizontal=F,height=9,width=11)
par(mfrow=c(7,5),mar=c(3.5,4,0,0))# par(mfcol=c(5,7),mar=c(4,4,1,1))
xlab=""
for (k in 1:7) {
  arg[[k]]$iaux[7] <- 0 ## turn off derivatives
  if (k==7) xlab="day"
  l0 <- ll(er[[k]]$thetaw,arg[[k]]$gamma,arg[[k]]$S0,arg[[k]]$iaux,arg[[k]]$daux,arg[[k]]$dat,Sb=NULL,
        link=arg[[k]]$link,plot=TRUE,step=arg[[k]]$step,getR=TRUE,lab=arg[[k]]$region,xlab=xlab)
  n.deaths <- n.deaths + l0$n.death
  arg[[k]]$iaux[7] <- 1
  Ik <- colSums(l0$xout[ii,])
  rep <- 1000;n <- length(l0$Rd)
  if (k==1) {
    R <- exp(log(dd[[1]]$Rd) + matrix(rnorm(n*rep),n,rep)*dd[[1]]$se.r)*Ik
    f <- exp(log(dd[[1]]$fi) + matrix(rnorm(n*rep),n,rep)*dd[[1]]$se.f)
    Rd <- dd[[k]]$Rd*Ik
    ftot <- dd[[k]]$fi
    Itot <- Ik
  } else {
    R <- R + exp(log(dd[[k]]$Rd) + matrix(rnorm(n*rep),n,rep)*dd[[k]]$se.r)*Ik
    f <- f + exp(log(dd[[k]]$fi) + matrix(rnorm(n*rep),n,rep)*dd[[k]]$se.f)
    Itot <- Itot + Ik
    Rd <- Rd+ dd[[k]]$Rd*Ik
    ftot <-ftot + dd[[k]]$fi
  }
  #text(200,.01,names(arg)[k])
  #X11()
}
if (ps) dev.off()

Rd <- Rd/Itot
Rav <- R/Itot
Rci <- apply(Rav,1,quantile,prob=c(.025,.975))
fci <- apply(f,1,quantile,prob=c(.025,.975))

for (j in 1:2) {
if (j==1) {
  xlim <- c(50,150);xpos <- 110
  fn <- "fw-gc.eps"
} else {
  xlim <- c(60,320);xpos <-150
  fn <- "w2-gc.eps"
}
vt <- c(84,216,245,258,288,310)

if (ps) postscript(fn) else X11()
par(mfrow=c(2,2),mar=c(5,5,1,1))
t <- 1:length(Rd)
c1 <- 1.3
plot(ftot,xlim=xlim,type="l",ylim=c(0,max(fci[2,50:150])),xlab="day",ylab="incidence",cex.lab=c1)
lines(t,fci[1,],lty=2);lines(t,fci[2,],lty=2)
abline(v=vt)

plot(dd[[1]]$fi,xlim=xlim,ylim=c(0,70000),type="l",lwd=3,xlab="day",ylab="incidence",cex.lab=c1)
lines(t,exp(log(dd[[1]]$fi)+dd[[1]]$se.f*2),lty=2)
lines(t,exp(log(dd[[1]]$fi)-dd[[1]]$se.f*2),lty=2)
text(xpos,20000,arg[[1]]$region)
for (k in 2:7) {
  lines(dd[[k]]$fi,col=k,lwd=3)
  text(xpos,10000+k*10000,arg[[k]]$region,col=k)
  lines(t,exp(log(dd[[k]]$fi)+dd[[k]]$se.f*2),lty=2,col=k)
  lines(t,exp(log(dd[[k]]$fi)-dd[[k]]$se.f*2),lty=2,col=k)
}
abline(v=vt)
plot(Rd,xlim=xlim,ylim=c(0,4),type="l",lwd=1,xlab="day",ylab="R",cex.lab=c1)
lines(t,Rci[1,],lty=2);lines(t,Rci[2,],lty=2)

abline(1,0);abline(v=vt)

plot(log(dd[[1]]$Rd),xlim=xlim,ylim=c(-2,2),type="l",lwd=3,xlab="day",ylab="log(R)",cex.lab=c1)
lines(t,log(dd[[1]]$Rd)+dd[[1]]$se.r*2,lty=2)
lines(t,log(dd[[1]]$Rd)-dd[[1]]$se.r*2,lty=2)
text(xpos,.35,arg[[1]]$region)
for (k in 2:7) {
  lines(log(dd[[k]]$Rd),col=k,lwd=3)
  lines(t,log(dd[[k]]$Rd)+dd[[k]]$se.r*2,lty=2,col=k)
  lines(t,log(dd[[k]]$Rd)-dd[[k]]$se.r*2,lty=2,col=k)
  text(xpos,.1+k*.25,arg[[k]]$region,col=k)
}
abline(0,0);abline(v=vt)
if (ps) dev.off()
}


  # clear
  cat("\014")
  rm(list = ls())
  dev.off(dev.list()["RStudioGD"])

  # Libraries
  library(seqinr)
  library(MASS)
  library(ggplot2)
  library(ggpubr)
  library(stringr)
  library(combinat)

  #--------------------- SELECT DATA ---------------------------------------------
  # 'artificial' and 'artificial-non-uni' for artificial dataset, 'real' for real-world data (biom_data_150bp.rda)
  dataset_type <- "real"

  #---------------- FUNCTIONS ----------------------------------------------------

  # Returns prime numbers up to integer n
  sieve <- function(n) {
    n <- as.integer(n)
    if(n > 1e6) stop("n too large")
    primes <- rep(TRUE, n)
    primes[1] <- FALSE
    last.prime <- 2L
    for(i in last.prime:floor(sqrt(n)))
    {
      primes[seq.int(2L*last.prime, n, last.prime)] <- FALSE
      last.prime <- last.prime + min(which(primes[(last.prime+1):n]))
    }
    which(primes)
  }

  # Reads sequence from input file
  seqList <- function(x, type) {
    if (type == "DNA") {
      as.list(read.fasta(x,
                         as.string = TRUE,
                         seqtype="DNA",
                         seqonly = TRUE,
                         strip.desc = TRUE))
    }
    else {
      if (type == "AA") {
        as.list(read.fasta(x,
                           as.string = FALSE,
                           seqtype="AA",
                           seqonly = FALSE,
                           strip.desc = TRUE))
      }
      else {
        NULL
      }
    }
  }

  # Generates random DNA sequences
  createRandomSequencesBasedOnDistr <- function(count, length, prob=c(0.25,0.25,0.25,0.25), fileNameRandSeqs) {

    sink(fileNameRandSeqs)
    for (i in 1:count) {
      cat(">Seq", i, "\n", sep = "")
      seqX <- sample(c("A","C","G","T"), length, rep=TRUE, prob)
      cat(paste(seqX,collapse=""), "\n", sep = "")
    }
    sink()
  }

  # Generating random permutations of godel number assignments
  createRandomSequenceValues <- function(seedList, type) {

    if ( type == "DNA" ) {
      results <- matrix(, nrow = length(seedList), ncol = 4)
    }
    else if ( type == "AA" ) {
      results <- matrix(, nrow = length(seedList), ncol = 23)
    }

    for (i in 1:length(seedList)) {
      if ( type == "DNA" ) {
        results[i,] <- sample(seq(from = 1, to = 4, by = 1), size = 4, replace = FALSE)
      }
      else if ( type == "AA" ) {
        results[i,] <- sample(seq(from = 1, to = 23, by = 1), size = 23, replace = FALSE)
      }
    }
    return(results)
  }

  assignSets <- function(randSequenceValues, type) {

    stringValues <- ""

    if (type == "DNA") {
      #DNA Set
      cat("Sequence Values: (")

      for (i in 1:4) {
        stringValues <- paste(stringValues, randSequenceValues[i], ", ")
      }
      cat(stringValues)
      cat(")\n")
      assign("A", randSequenceValues[1], envir = .GlobalEnv)
      assign("C", randSequenceValues[2], envir = .GlobalEnv)
      assign("G", randSequenceValues[3], envir = .GlobalEnv)
      assign("T", randSequenceValues[4], envir = .GlobalEnv)
    }
    else if ( type == "AA") {
      # AA Set 1
      cat("Sequence Values: (")
      for (i in 1:23) {
        stringValues <- paste(stringValues, randSequenceValues[i], ", ")
      }
      cat(stringValues)
      cat(")\n")
      assign("A", randSequenceValues[1], envir = .GlobalEnv) # alanine, ala
      assign("R", randSequenceValues[2], envir = .GlobalEnv) # arginine, arg
      assign("N", randSequenceValues[3], envir = .GlobalEnv) # asparagine, asn
      assign("D", randSequenceValues[4], envir = .GlobalEnv) # aspartic acid, asp
      assign("B", randSequenceValues[5], envir = .GlobalEnv) # sparagine or aspartic acid, asx
      assign("C", randSequenceValues[6], envir = .GlobalEnv) # cysteine, cys
      assign("E", randSequenceValues[7], envir = .GlobalEnv) # glutamic acid, glu
      assign("Q", randSequenceValues[8], envir = .GlobalEnv) # glutamine, gln
      assign("Z", randSequenceValues[9], envir = .GlobalEnv) # glutamine or glutamic acid, glx
      assign("G", randSequenceValues[10], envir = .GlobalEnv) # glycine, gly
      assign("H", randSequenceValues[11], envir = .GlobalEnv) # histidine, his
      assign("I", randSequenceValues[12], envir = .GlobalEnv) # isoleucine, ile
      assign("L", randSequenceValues[13], envir = .GlobalEnv) # leucine, leu
      assign("K", randSequenceValues[14], envir = .GlobalEnv) # lysine, lys
      assign("M", randSequenceValues[15], envir = .GlobalEnv) # methionine, met
      assign("F", randSequenceValues[16], envir = .GlobalEnv) # phenylalanine, phe
      assign("P", randSequenceValues[17], envir = .GlobalEnv) # proline, pro
      assign("S", randSequenceValues[18], envir = .GlobalEnv) # serine, ser
      assign("T", randSequenceValues[19], envir = .GlobalEnv) # threonine, thr
      assign("W", randSequenceValues[20], envir = .GlobalEnv) # tryptophan, trp
      assign("Y", randSequenceValues[21], envir = .GlobalEnv) # tyrosine, tyr
      assign("V", randSequenceValues[22], envir = .GlobalEnv) # valine, val
      assign("X", randSequenceValues[23], envir = .GlobalEnv) # undetermined
    }

    stringValues
  }

  # Statistics function
  godelStatistics <- function(x) {

    if (logOutput) {
      cat(paste("  Min  : ", summary(x)[1], "\n"))
      cat(paste("1st Qu.: ", summary(x)[2], "\n"))
      cat(paste("Median : ", summary(x)[3], "\n"))
      cat(paste(" Mean  : ", summary(x)[4], "\n"))
      cat(paste("3rd Qu.: ", summary(x)[5], "\n"))
      cat(paste("  Max  : ", summary(x)[6], "\n"))
      cat(paste("St. Dev: ", sd(x), "\n"))
    }

    results <- list(minG = summary(x)[1], firstQG = summary(x)[2],
         medianG = summary(x)[3], meanG = summary(x)[4],
         thirdQG = summary(x)[5], maxG = summary(x)[6],
         stdG = sd(x))

    return(results)
  }

  calculateGodelNumbers <- function(sequences, primes, encoding){

    chars = str_split(sequences, "", simplify = TRUE)

    chars.enc = matrix(data = 0, nrow = nrow(chars), ncol = ncol(chars))

    chars.enc = chars.enc + (chars == "A") * encoding[1]
    chars.enc = chars.enc + (chars == "C") * encoding[2]
    chars.enc = chars.enc + (chars == "G") * encoding[3]
    chars.enc = chars.enc + (chars == "T") * encoding[4]

    primes = log(primes)

    chars.enc = chars.enc %*% diag(primes)
    godel_nums <- rowSums(chars.enc) # gn_new
    return(godel_nums)

  }

  #----------------------------- Main code ---------------------------------------
  primes <- sieve(20000)            # length of primes should be >= max sequence length
  logOutput <- FALSE                # debug output
  replicate <- TRUE                 # if we need to set specific seed numbers
  type <- "DNA"                     # type of data

  # encodings
  encodings <- permn(c(1,2,3,4))
  numberOfEncodings <- 24;              # number of different assignments of letters for Godel numbers (we have DNA data)

  # The following values should execute within 6 minutes with a good enough resolution
  seqLengthLimit <- 150
  numberOfArtificialSeqs <- 90000     # "A","C","G","T"

  # Creating artificial data set
  # createRandomSequencesBasedOnDistr(numberOfArtificialSeqs, seqLengthLimit, c(0.3,0.3,0.2,0.2), "data/artificialSeq-non-unis.fasta")

  # Loading a fasta file - artificial dataset
  if (dataset_type == "artificial"){
    ompGene.list <- seqList("data/artificialSeqs.fasta", type)
    ompGene.list.raw <- ompGene.list
  } else if (dataset_type == "real"){
    ompGene.list <- seqList("data/realSeqs_biom_150bp.fasta", type)
  } else if (dataset_type == "artificial-non-uni"){
    ompGene.list <- seqList("data/artificialSeq-non-unis.fasta", type)
  } else {
    print("Incorrect dataset type.")
  }
  sequences <- unlist(ompGene.list)

  #encodingValues <- createRandomSequenceValues(seedValuesList, type) - we don't neet this function
  encodingValues <- matrix(0, nrow = length(encodings), ncol = 4)
  for (i in 1:nrow(encodingValues)){
    encodingValues[i,] <- encodings[[i]]
  }

  sizeExp <- length(ompGene.list)
  selectedSequences <- c(1:sizeExp)

  godelValuePoints <- data.frame(matrix(0, ncol = numberOfEncodings, nrow = sizeExp))
  namesList <- list()
  for (i in 1:numberOfEncodings) {
    namesList[i] <- paste('godel_log_pos', i, sep = '')
  }
  godelValuePoints <- setNames(godelValuePoints, namesList)
  godelValuePoints$seqNames <- unlist(attributes(ompGene.list)$name)

  stringAssignValues <- vector(mode="character", length=numberOfEncodings)

  len_of_sequences <- str_length(sequences[1])
  primes <- primes[1:len_of_sequences]
  primes <- as.numeric(primes)

  # main loop - godel number calculation
  for (indexPos in 1:numberOfEncodings) {

    stringAssignValues[indexPos] <- assignSets(encodingValues[indexPos,], type)
    encoding <- encodingValues[indexPos,]
    godelValuePoints[, paste('godel_log_pos', indexPos, sep = '')] <- calculateGodelNumbers(sequences, primes, encoding)

  }

  statsPos <- data.frame(matrix(0, ncol = 7, nrow = numberOfEncodings))
  statsPos <- setNames(statsPos, c("minG", "firstQG", "medianG", "meanG", "thirdQG", "maxG", "stdG"))
  for (indexPos in 1:numberOfEncodings) {
    statsPos[indexPos, ] <- godelStatistics(godelValuePoints[, paste('godel_log_pos', indexPos, sep = '')])
  }

  statsPos_rownames <- c()
  for (i in 1:numberOfEncodings){
    this_enc <- paste(encodingValues[i,], collapse = '')
    statsPos_rownames <- c(statsPos_rownames, paste("enc_", this_enc, sep = ''))
  }

  rownames(statsPos) <- statsPos_rownames

  # Theoretical distribution parameters
  P1 <- sum(log(primes[1:seqLengthLimit]))
  P2 <- sum((log(primes[1:seqLengthLimit]))^2)
  theoreticalMeanEqual <- P1*2.5
  theoreticalStdEqual <- sqrt(P2*1.25)
  theoretical_dist <- rnorm(90000, theoreticalMeanEqual, theoreticalStdEqual)

  # comparisons
  para <- matrix(nrow=numberOfEncodings, ncol = 2)
  for (indexPos in 1:numberOfEncodings) {
    fit <- fitdistr(godelValuePoints[, indexPos], "normal")
    para[indexPos, 1] <- fit$estimate["mean"]
    para[indexPos, 2] <- fit$estimate["sd"]
  }

  statsPos$stringAssignValues <- as.factor(stringAssignValues)
  statsPos$fit_estimate <- ((para[,1]-theoreticalMeanEqual)/para[,1])*100
  statsPos$fit_sd <- ((para[,2]-theoreticalStdEqual)/para[,2])*100

  statsPos$theoretical_mean <- theoreticalMeanEqual
  statsPos$theoretical_std <-theoreticalStdEqual

  # t-test
  statsPos$p.value <- NaN
  for (i in 1:numberOfEncodings){
    t_test_results <- t.test(godelValuePoints[,i], theoretical_dist)
    statsPos[i,]$p.value <- t_test_results$p.value
  }

  #-------------------- Main histogram plotting ----------------------------------
  dir.create('plots')
  indexPos = 1 # Specify which encoding to plot
  binwidthPlot = 5
  this_enc <- paste(encodingValues[indexPos,], collapse = '')

  # comparisons file save
  filepath <- paste('plots/comparisons_',dataset_type, '.csv', sep = '' )
  write.csv(as.data.frame(statsPos), file = filepath)

  this_title= expression(paste('Histogram and Density plot of ln(Godel numbers) - encoding 1234', sep = '' ))

  # Plots
  filepath <- paste('plots/histogram_encoding_', this_enc,'_',dataset_type, '.png', sep = '' )
  png(file = filepath, width=1200, height=800)
  my_plot <- ggplot(godelValuePoints, aes(x=godelValuePoints[, indexPos]), environment = environment()) +
    geom_histogram(aes(y=..density..), colour="black", fill="white", binwidth=binwidthPlot)+
    geom_density(alpha=.2, fill="#FF6666") +
    stat_function(fun = dnorm, args = list(mean = theoreticalMeanEqual, sd = theoreticalStdEqual), color = "darkblue", size = 1, linetype = "longdash") +
    stat_function(fun = dnorm, args = list(mean = para[indexPos,1], sd = para[indexPos,2]), color = "darkred", size = 2, linetype = "dotdash") +
    labs(title = this_title, x=expression('ln(Godel Numbers)'), y = "Density")+
    xlim(1000,3000) +
    ylim(0, 0.009) +
    scale_color_brewer(palette="Accent") +
    theme(title  = element_text(size=20),
          axis.text.x = element_text(size = 20, face = 'bold'), axis.text.y = element_text(size = 15),
          axis.title.x =  element_text(size=20))
  print(my_plot)

  dev.off()

  # ---------------- Other plots -------------------------------------------------

  p1 <- ggplot(statsPos, aes(x = reorder(stringAssignValues, fit_estimate), y = fit_estimate, group=1)) +
    geom_point() +
    geom_line() +
    theme(axis.text.x = element_text(angle = 90, hjust = 1)) +
    labs(title= "Difference between theoretical and observed value of Mean (%)",x="Assignments", y = "Difference in mean (%)")

  p2 <- ggplot(statsPos, aes(x = reorder(stringAssignValues, fit_estimate), y = fit_sd, group=1)) +
    geom_point() +
    geom_line() +
    theme(axis.text.x = element_text(angle = 90, hjust = 1)) +
    labs(title= "Difference between theoretical and observed value of standard deviation (%)",x="Assignments", y = "Difference in  standard deviation (%)")

  ggarrange(p1 + rremove("x.text"), p2,
            labels = c("A", "B"),
            ncol = 1, nrow = 2)

  # a t g c
  lapply(which(statsPos$maxG == min(statsPos$maxG)),
         function(x) assignSets(encodingValues[x,], "DNA"))
  # artificial dataset: 4 , 1 , 3 , 2 (min of meanG)
  # artificial dataset: 4 , 2 , 1 , 3 (min of minG)
  # artificial dataset: 2 , 1 , 3 , 4 (min of maxG)

  # a t g c
  lapply(which(statsPos$maxG == max(statsPos$maxG)),
         function(x) assignSets(encodingValues[x,], "DNA"))
  # artificial dataset: 1 , 4 , 2 , 3 (max of meanG)
  # artificial dataset: 4 , 3 , 2 , 1 (max of minG)
  # artificial dataset: 1 , 3 , 4 , 2 (max of maxG)


### About:
## The function is to estimate the mean and covariance function
## from a cluster of functions/longitudinal observations. The output
## will also predict each curve.

## Function arguments:
### "data" is a data frame with three arguments:
# (1) "argvals": observation times;
# (2) "subj": subject indices;
# (3) "y": values of observations;
# Note that: we only handle complete data, so missing values are not allowed at this moment

## "newdata" is of the same strucutre as "data".
## To predict at new values of "argvals", make the corresponding "y" NA

### "center" means if we want to compute population mean

### "argvals.new" if we want the estimated covariance function at "argvals.new"; if NULL,
### then 100 equidistant points in the range of "argvals" in "data"

### "knots" is the number of interior knots for B-spline basis functions to be used;

face.sparse.inner <- function(data, newdata = NULL, W = NULL,
                              center = TRUE, argvals.new = NULL,
                              knots = 7, knots.option = "equally-spaced",
                              p = 3, m = 2, lambda = NULL, lambda_mean = NULL,
                              search.length = 14,
                              lower = -3, upper = 10,
                              calculate.scores = FALSE, pve = 0.99) {

  #########################
  #### step 0: read in data
  #########################
  check.data(data)
  if (!is.null(newdata)) {
    check.data(newdata, type = "predict")
  }

  y <- data$y
  t <- data$argvals
  subj <- data$subj
  tnew <- argvals.new
  if (is.null(tnew)) tnew <- seq(min(t), max(t), length = 100)

  fit_mean <- NULL

  knots.initial <- knots
  #########################
  #### step 1: demean
  #########################
  r <- y
  mu.new <- rep(0, length(tnew))
  if (center) {
    fit_mean <- pspline(data, argvals.new = tnew, knots = knots.initial, lambda = lambda_mean)
    mu.new <- fit_mean$mu.new
    r <- y - fit_mean$fitted.values
  }
  #########################
  #### step 2:raw estimates
  #########################
  indW <- F # whether identity W
  if (is.null(W)) indW <- T

  raw <- raw.construct(data.frame("argvals" = t, "subj" = subj, "y" = as.vector(r)))
  C <- raw$C
  st <- raw$st
  N <- raw$st
  N2 <- raw$N2
  if (indW) W <- raw$W
  n0 <- raw$n0
  # indicator where ti = tj
  delta <- Matrix((st[, 1] == st[, 2]) * 1) # sparse

  #########################
  #### step 3: smooth
  #########################

  # set up
  knots <- construct.knots(t, knots, knots.option, p)
  ## construct design & penalty on column dimension
  List <- pspline.setting(st[, 1], knots = knots, p, m, type = "simple", knots.option = knots.option)
  B1 <- List$B
  B1 <- Matrix(B1)
  DtD <- List$P
  ## construct design on row dimension
  B2 <- spline.des(knots = knots, x = st[, 2], ord = p + 1, outer.ok = TRUE, sparse = TRUE)$design
  c <- dim(B1)[2]
  c2 <- c * (c + 1) / 2
  ## combine tensor product design
  B <- Matrix(t(KhatriRao(Matrix(t(B2)), Matrix(t(B1)))))
  G <- Matrix(duplication.matrix(c))

  BtWB <- matrix(0, nrow = c^2, ncol = c^2)
  Wdelta <- c()
  WC <- c()
  for (i in 1:n0) {
    seq <- (sum(N2[1:i]) - N2[i] + 1):(sum(N2[1:i])) # select combos for object i
    B3 <- Matrix(matrix(B[seq, ], nrow = length(seq)))
    W3 <- W[[i]] # don't form a large W
    BtWB <- BtWB + crossprod(B3, W3 %*% B3)
    Wdelta <- c(Wdelta, as.matrix(W3 %*% delta[seq]))
    WC <- c(WC, as.matrix(W3 %*% C[seq]))
  }

  GtBtWBG <- crossprod(G, BtWB %*% G)

  BG <- B %*% G # sparse
  detWde <- crossprod(delta, Wdelta) # detWde = sum(delta)
  GtBtWdelta <- crossprod(BG, Wdelta)
  XtWX <- rbind(cbind(GtBtWBG, GtBtWdelta), cbind(t(GtBtWdelta), detWde))
  X <- cbind(BG, delta)
  # precalculate (2 eigens)
  eSig <- eigen(XtWX, symmetric = TRUE)
  V <- eSig$vectors
  E <- eSig$values
  E <- E + 0.000001 * max(E)
  Sigi_sqrt <- matrix.multiply(V, 1 / sqrt(E)) %*% t(V)
  ## penalty
  P <- crossprod(G, Matrix(suppressMessages(kronecker(diag(c), DtD)))) %*% G

  Q <- bdiag(P, 0)
  tUQU <- crossprod(Sigi_sqrt, (Q %*% Sigi_sqrt))
  Esig <- eigen(tUQU, symmetric = TRUE)
  # s, f, ftilde, F(A here)
  U <- Esig$vectors
  s <- Esig$values
  A0 <- Sigi_sqrt %*% U
  A <- as.matrix(X %*% A0) # F=XA dense

  AtA <- crossprod(A) # diff
  f <- crossprod(A, C) # diff
  ftilde <- crossprod(A, WC) # diff

  c2 <- c2 + 1
  g <- rep(0, c2)
  G1 <- matrix(0, c2, c2)
  mat_list <- list()
  # G, g, Li(AitAi here), Litilde(Li here)
  for (i in 1:n0) {
    seq <- (sum(N2[1:i]) - N2[i] + 1):(sum(N2[1:i]))
    Ai <- matrix(A[seq, ], nrow = length(seq))
    AitAi <- crossprod(Ai) # t(Ai)%*%Ai
    Wi <- W[[i]]

    fi <- crossprod(Ai, C[seq]) # t(Fi)Ci
    Ji <- crossprod(Ai, Wi %*% C[seq])
    Li <- crossprod(Ai, Wi %*% Ai)
    g <- g + Ji * fi
    G1 <- G1 + AitAi * (Ji %*% t(ftilde))

    LList <- list()
    LList[[1]] <- AitAi
    LList[[2]] <- Li
    mat_list[[i]] <- LList
  }

  # select smoothing parameter
  Lambda <- seq(lower, upper, length = search.length)
  Gcv <- 0 * Lambda
  gcv <- function(x) {
    lambda <- exp(x)
    d <- 1 / (1 + lambda * s)
    ftilde_d <- ftilde * d
    cv0 <- -2 * sum(ftilde_d * f)
    cv1 <- sum(ftilde_d * (AtA %*% ftilde_d))
    cv2 <- 2 * sum(d * g)
    cv3 <- -4 * sum(d * (G1 %*% d))
    cv4 <- sum(unlist(sapply(mat_list, function(x) {
      a <- x[[1]] %*% ftilde_d
      b <- x[[2]] %*% ftilde_d
      2 * sum(a * b * d)
    })))
    cv <- cv0 + cv1 + cv2 + cv3 + cv4
    return(cv)
  }
  if (is.null(lambda)) {
    Lambda <- seq(lower, upper, length = search.length) # construct lambda grid
    Length <- length(Lambda)
    Gcv <- rep(0, Length)
    for (i in 1:Length) {
      Gcv[i] <- gcv(Lambda[i])
    }
    i0 <- which.min(Gcv)
    lambda <- exp(Lambda[i0])
  }
  # parameter estimated (sym)
  alpha <- matrix.multiply(A0, 1 / (1 + lambda * s)) %*% ftilde
  Theta <- G %*% alpha[1:(c2 - 1)] # wrong index?
  Theta <- matrix(Theta, c, c)
  sigma2 <- alpha[c2]
  if (sigma2 <= 0.000001) {
    warning("error variance cannot be non-positive, reset to 1e-6!")
    sigma2 <- 0.000001
  }
  # make sure Theta positive definite(2 eigens)
  Eigen <- eigen(Theta, symmetric = TRUE)
  Eigen$values[Eigen$values < 0] <- 0
  npc <- sum(Eigen$values > 0) # which.max(cumsum(Eigen$values)/sum(Eigen$values)>pve)[1]
  if (npc > 1) {
    Theta <- matrix.multiply(Eigen$vectors[, 1:npc], Eigen$values[1:npc]) %*% t(Eigen$vectors[, 1:npc])
    Theta_half <- matrix.multiply(Eigen$vectors[, 1:npc], sqrt(Eigen$values[1:npc]))
  }
  if (npc == 1) {
    Theta <- Eigen$values[1] * suppressMessages(kronecker(Eigen$vectors[, 1], t(Eigen$vectors[, 1])))
    Theta_half <- sqrt(Eigen$values[1]) * Eigen$vectors[, 1]
  }
  Eigen <- eigen(Theta, symmetric = TRUE)

  #########################
  #### step 4: calculate estimated covariance function
  ######################### (2 eigens)
  # prepare estimation dense design
  Bnew <- spline.des(knots = knots, x = tnew, ord = p + 1, outer.ok = TRUE, sparse = TRUE)$design
  Gmat <- crossprod(Bnew) / nrow(Bnew)
  eig_G <- eigen(Gmat, symmetric = T)
  G_half <- eig_G$vectors %*% diag(sqrt(eig_G$values)) %*% t(eig_G$vectors)
  G_invhalf <- eig_G$vectors %*% diag(1 / sqrt(eig_G$values)) %*% t(eig_G$vectors)

  Chat.new <- as.matrix(tcrossprod(Bnew %*% Matrix(Theta), Bnew)) # cov matrix
  Chat.diag.new <- as.vector(diag(Chat.new))
  Cor.new <- diag(1 / sqrt(Chat.diag.new)) %*% Chat.new %*% diag(1 / sqrt(Chat.diag.new)) # scale to cor matrix

  # what is these parts? (another npc)
  Eigen.new <- eigen(as.matrix(G_half %*% Matrix(Theta) %*% G_half), symmetric = TRUE)
  # Eigen.new = eigen(Chat.new,symmetric=TRUE)
  npc <- which.max(cumsum(Eigen$values) / sum(Eigen$values) > pve)[1] # determine number of PCs
  eigenfunctions <- matrix(Bnew %*% G_invhalf %*% Eigen.new$vectors[, 1:min(npc, length(tnew))], ncol = min(npc, length(tnew)))
  eigenvalues <- Eigen.new$values[1:min(npc, length(tnew))]
  # eigenfunctions = eigenfunctions*sqrt(length(tnew))/sqrt(max(tnew)-min(tnew))
  # eigenvalues = eigenvalues/length(tnew)*(max(tnew)-min(tnew))

  #########################
  #### step 5: calculate variance
  #########################
  var.error.hat <- rep(sigma2, length(t))
  var.error.new <- rep(sigma2, length(tnew))

  Chat.raw.new <- Chat.new + diag(var.error.new) # add variance sigma2
  Chat.raw.diag.new <- as.vector(diag(Chat.raw.new))
  Cor.raw.new <- diag(1 / sqrt(Chat.raw.diag.new)) %*% Chat.raw.new %*% diag(1 / sqrt(Chat.raw.diag.new))
  #########################
  #### step 6: prediction
  #########################
  if (!is.null(newdata)) {
    # means for newdata
    mu.pred <- rep(0, length(newdata$argvals))
    var.error.pred <- rep(sigma2, length(newdata$argvals))
    if (center) {
      mu.pred <- predict.pspline.face(fit_mean, newdata$argvals)
    }

    # setup prediction data
    subj.pred <- newdata$subj
    subj_unique.pred <- unique(subj.pred)
    y.pred <- newdata$y
    Chat.diag.pred <- 0 * y.pred
    se.pred <- 0 * y.pred

    scores <- list(
      subj = subj_unique.pred,
      scores = matrix(NA, nrow = length(subj_unique.pred), ncol = npc),
      u = matrix(NA, nrow = length(subj_unique.pred), ncol = nrow(Theta))
    )

    # prediction for each subject i
    for (i in 1:length(subj_unique.pred)) {
      # select the observations for subject i
      sel.pred <- which(subj.pred == subj_unique.pred[i])
      lengthi <- length(sel.pred)

      pred.points <- newdata$argvals[sel.pred]
      mu.predi <- mu.pred[sel.pred]
      var.error.predi <- var.error.pred[sel.pred]
      # center response
      y.predi <- y.pred[sel.pred] - mu.predi
      sel.pred.obs <- which(!is.na(y.predi))
      # select tij correspond to non na response
      obs.points <- pred.points[sel.pred.obs]
      if (!is.null(obs.points)) {
        var <- mean(var.error.predi[sel.pred.obs])
        if (var == 0 & length(sel.pred.obs) < npc) {
          stop("Measurement error estimated to be zero and there are fewer observed points thans PCs; scores
                cannot be estimated.")
        }
        B3i.pred <- spline.des(knots = knots, x = pred.points, ord = p + 1, outer.ok = TRUE, sparse = TRUE)$design
        B3i <- spline.des(knots = knots, x = obs.points, ord = p + 1, outer.ok = TRUE, sparse = TRUE)$design
        Chati <- tcrossprod(B3i %*% Theta, B3i)
        Chat.diag.pred[sel.pred] <- diag(Chati)
        if (length(sel.pred.obs) == 1) Ri <- var.error.predi[sel.pred.obs]
        if (length(sel.pred.obs) > 1) Ri <- diag(var.error.predi[sel.pred.obs])
        Vi.inv <- as.matrix(solve(Chati + Ri))

        Vi.pred <- tcrossprod(B3i.pred %*% Theta, B3i.pred)
        Hi <- as.matrix(B3i.pred %*% tcrossprod(Theta, B3i) %*% Vi.inv)
        ui <- tcrossprod(Theta, B3i) %*% Vi.inv %*% y.predi[sel.pred.obs]
        scores$u[i, ] <- as.vector(ui)
        y.pred[sel.pred] <- as.numeric(Hi %*% y.predi[sel.pred.obs]) + mu.predi # prediction
        temp <- as.matrix(B3i.pred %*% tcrossprod(Theta, B3i))
        if (length(sel.pred.obs) > 1) {
          se.pred[sel.pred] <- sqrt(diag(Vi.pred - temp %*% Vi.inv %*% t(temp)))
        }
        if (length(sel.pred.obs) == 1) {
          se.pred[sel.pred] <- sqrt(Vi.pred[1, 1] - Vi.inv[1, 1] * temp %*% t(temp))
        }

        ## predict scores
        if (calculate.scores == TRUE) {
          temp <- matrix(t(eigenfunctions), nrow = npc) %*% (as.matrix(Bnew) %*% ui) / sum(eigenfunctions[, 1]^2)
          temp <- as.matrix(temp)
          scores$scores[i, 1:npc] <- temp[, 1]
        }
      }
    }
  } ## if(is.null(newdata))
  if (is.null(newdata)) {
    y.pred <- NULL
    mu.pred <- NULL
    var.error.pred <- NULL
    Chat.diag.pred <- NULL
    se.pred <- NULL
    scores <- NULL
  }

  res <- list(
    newdata = newdata, W = W, y.pred = y.pred, Theta = Theta, argvals.new = tnew,
    mu.new = mu.new, Chat.new = Chat.new, var.error.new = var.error.new,
    Cor.new = Cor.new, eigenfunctions = eigenfunctions, eigenvalues = eigenvalues,
    Cor.raw.new = Cor.raw.new, Chat.raw.diag.new = Chat.raw.diag.new,
    scores = scores, calculate.scores = calculate.scores,
    mu.hat = fit_mean$fitted.values, var.error.hat = var.error.hat,
    mu.pred = mu.pred, var.error.pred = var.error.pred, Chat.diag.pred = Chat.diag.pred,
    se.pred = se.pred,
    fit_mean = fit_mean, lambda_mean = fit_mean$lambda,
    lambda = lambda, Gcv = Gcv, Lambda = Lambda, knots = knots, knots.option = knots.option, s = s, npc = npc, p = p, m = m,
    center = center, pve = pve, sigma2 = sigma2, r = r, DtD = DtD,
    U = Eigen.new$vectors[, 1:npc], G_invhalf = G_invhalf
  )

  class(res) <- "face.sparse"
  return(res)
}


#' Function for fast sampling of DP mixture of Linear regressions.
#'
#' This function takes in a training data.frame and optional testing data.frame and performs posterior sampling. It returns posterior mean regression line for training and test sets. The function is built for continuous outcomes.
#' This differs from NDPMix in the following ways: NDPMix returns draws from the posterior *predictive* distribution of the outcome, whereas fDPMix() returns the regression line. This will have the same mean as the NDPMix predictions, but lower variance. Finally, the back-end computation is different, with regression line evaluation at point X being evaluated as a weighted average of cluster-specific regression evaluations at X. This is faster than NDPMix which takes a Monte Carlo appraoch: assign X to one of the cluster specific regressions, then draw a predicted outcome from that cluster's regression.
#'
#' We recommend normalizing continuous covariates and outcomes via the \code{scale} function before running \code{fDPMix}
#'
#' Please see \url{https://stablemarkets.github.io/ChiRPsite/index.html} for examples and detailed model and parameter descriptions.
#'
#' @import stats
#'
#' @param d_train A \code{data.frame} object with outcomes and model covariates/features. All features must be \code{as.numeric} - either continuous or binary with binary variables coded using \code{1} and \code{0}. Categorical features are not supported. We recommend standardizing all continuous features. NA values are not allowed and each row should represent a single subject, longitudinal data is not supported.
#' @param d_test Optional \code{data.frame} object containing a test set of subjects containing all variables specifed in \code{formula}. The same rules that apply to \code{d_train} apply to \code{d_test}.
#' @param formula Specified in the usual way, e.g. for \code{p=2} covariates, \code{y ~ x1 + x2}. All covariates - continuous and binary - must be \code{as.numeric} , with binary variables coded as \code{1} or \code{0}. We recommend standardizing all continuous features. NA values are not allowed and each row should represent a single subject, longitudinal data is not supported.
#' @param burnin integer specifying number of burn-in MCMC draws.
#' @param iter integer greater than \code{burnin} specifying how many total MCMC draws to take.
#' @param init_k Optional. integer specifying the initial number of clusters to kick off the MCMC sampler.
#' @param phi_y Optional. Length two \code{as.numeric} vector specifying the mean and variance of the inverse gamma prior on the outcome variance. By default, \code{phi_y[1]} is set to \code{var(y)}, where \code{y} is the outcome, so the inverse gamma prior is centered around the marginal empirical variance. With prior variance given by \code{phi_y[2]}. By default, \code{phi_y[2]=.5*var(y)}. This value should be large to be "flat" and small to be "tight" around the marginal empirical variance.
#' @param beta_prior_mean Optional. If there are \code{p} covariates, a length \code{p+1} \code{as.numeric} vector specifying mean of the Gaussian prior on the outcome model's conditional mean parameter vector. Default is regression coefficients from running OLS on the outcomes.
#' @param beta_prior_var Optional. If there are \code{p} covariates, a length \code{p+1} \code{as.numeric} vector specifying variance of the Gaussian prior on the outcome model's conditional mean parameter vector. The full covarince of the prior is set to be diagonal. This vector specifies the diagonal enteries of this prior covariance. Default is estimated variances from running OLS on the outcome.
#' @param beta_var_scale Optional. A multiplicative constant that scales \code{beta_prior_var}. If you leave \code{beta_prior_mean} and \code{beta_prior_var} at their defaults, This constant toggles how wide new cluster parameters are dispersed around the observed data parameters.
#' @param mu_scale Optional. An \code{as.numeric}, scalar constant that controls how widely distributed new cluster continuous covariate means are distributed around the empirical covariate mean. Specifically, all continuous covariates are assumed to have Gaussian likelihood with Gaussian prior on their means. \code{mu_scale=2} specifies that the variance of the Gaussian prior is twice as large as the empirical variance.
#' @param tau_x Optional numeric length two vector for specifying the inverse gamma prior on each continuous covariate's prior variance. By default, the inverse gamma prior is centered around the empirical variance. \code{tau_x[1]} is the positive constant that multiplies this empirical variance. A value of 1 indicates that the inverse gamma is centered around the empirical variance. A value of one half centers the inverse gamma prior around one half the empirical covariate variance. The second argument of \code{tau_x} is the prior variance. For example, if\code{tau_x[1]=1} and \code{tau_x[2]=.001}, then the inverse gamma prior is tight around the empirical variance. If \code{tau_x[2]=100}, then the inverse gamma prior for this covariate is widely distributed around the empirical variance.
#' @return Returns \code{predictions$train} and \code{cluster_inds$train}. \code{predictions$train} returns an \code{nrow(d_train)} by \code{iter - burnin} matrix of posterior predictions. \code{cluster_inds$train} returns an \code{nrow(d_train)} by \code{iter - burnin} matrix of cluster assignment indicators, which can be input into the function \code{cluster_assign_mode()} to compute posterior mode assignment. \code{predictions$test} and \code{cluster_inds$test} are returned if \code{d_test} is specified.
#' @keywords Dirichlet Process Gaussian
#' @examples
#' set.seed(1)
#'
#' N = 200
#' x<-seq(1,10*pi, length.out = N) # confounder
#' y<-rnorm(n = length(x), sin(.5*x), .07*x )
#' d <- data.frame(x=x, y=y)
#' d$x <- as.numeric(scale(d$x))
#' d$y <- as.numeric(scale(d$y))
#'
#' plot(d$x,d$y, pch=20, xlim=c(min(d$x), max(d$x)+1 ), col='gray')
#'
#' d_test = data.frame(x=seq(max(d$x), max(d$x+1 ), .01 ))
#'
#'
#' res = fDPMix(d_train = d, d_test = d_test, formula = y ~ x,
#'              iter=100, burnin=50, tau_x = c(.01, .001) )
#'
#' ## in-sample
#' lines(d$x, rowMeans(res$train), col='steelblue')
#' lines(d$x, apply(res$train,1,quantile,probs=.05) , col='steelblue', lty=2)
#' lines(d$x, apply(res$train,1,quantile,probs=.90) , col='steelblue', lty=2)
#'
#' ## out of sample
#' lines(d_test$x, rowMeans(res$test), col='pink')
#' lines(d_test$x, apply(res$test,1,quantile,probs=.05) , col='pink', lty=2)
#' lines(d_test$x, apply(res$test,1,quantile,probs=.90) , col='pink', lty=2)
#'
#' @export
fDPMix<-function(d_train, formula, d_test=NULL, burnin=100,iter=1000, init_k=10,
                 phi_y=NULL,
                 beta_prior_mean=NULL, beta_prior_var=NULL,beta_var_scale=1000,
                 mu_scale=1, tau_x =c(.05,2) ){

  ###------------------------------------------------------------------------###
  #### 0 - Parse User Inputs                                                ####
  ###------------------------------------------------------------------------###

  # error checking user inputs
  if( missing(d_train) ){ stop("ERROR: must specify a training data.frame.") }
  x <- all.vars(formula[[3]]) # covariate names
  y <- all.vars(formula[[2]]) # outcome name

  nparams <- length(x) + 1

  y <- d_train[,y]

  if(is.null(phi_y)){
    phi_y = c(var(y), 1/2*var(y))
  }

  func_args<-mget(names(formals()),sys.frame(sys.nframe()))
  error_check(func_args,'fDP')

  if(!is.null(d_test)){
    xt <- model.matrix(data=rbind(d_train[,x, drop=F], d_test[,x, drop=F]),
                       object= as.formula(paste0('~ ',paste0(x, collapse = '+'))))
    nt <- nrow(xt)
    test_ind = c(rep(0, nrow(d_train)), rep(1, nrow(d_test)) )
  }else{
    xt <- model.matrix(data=d_train,
                       object= as.formula(paste0('~ ',paste0(x, collapse = '+'))))
    nt <- nrow(xt)
    test_ind = c( rep(0, nrow(d_train)) )
  }

  x_type <- vector(length = length(x))
  for(v in 1:length(x)){
    if( length(unique(d_train[, x[v] ]))==2 ){
      x_type[v] <- 'binary'
    }else if(length(unique(d_train[,x[v]]))>2){
      x_type[v] <- 'numeric'
    }
  }

  store_l <- iter - burnin
  x_names <- x
  x <- model.matrix(data=d_train, object = formula )

  n<-nrow(x)

  xall_names <- x_names
  nparams_all <- length(xall_names)
  xall <- d_train[,xall_names]

  all_type_map <- c(x_type)
  names(all_type_map) <- c(x_names)
  xall_type <- all_type_map[as.character(xall_names)]

  xall_names_bin <- xall_names[xall_type == 'binary']
  xall_names_num <- xall_names[xall_type == 'numeric']

  # parse numeric versus binary variables for outcome model
  n_bin_p <- sum(x_type=='binary')
  n_num_p <- sum(x_type=='numeric')

  ## calibrate priors
  if(is.null(beta_prior_mean)){
    reg <- lm(data=d_train, formula = formula)
    beta_prior_mean <- reg$coefficients
  }

  if(is.null(beta_prior_var)){
    reg <- lm(data=d_train, formula = formula)
    beta_prior_var <- beta_var_scale*diag(vcov(reg))
  }

  ## test let phi_y[1] equal prior mean of inv gamma prior for phi
  ## let phi_y[2] equal the prior variance.
  ## then set phy_i[1] = sample outcome variance
  ## and let phi_y[3] control flattness around sample variance.

  g1 = (1/phi_y[2])*(phi_y[1]^2) + 2
  b1 = (g1 - 1)*phi_y[1]

  prior_means <- apply(x[,xall_names_num, drop=F], 2, mean)
  names(prior_means) <- xall_names_num

  prior_var <- mu_scale*apply(x[,xall_names_num, drop=F], 2, var)
  names(prior_var) <- xall_names_num

  g2 = (1/tau_x[2])*( tau_x[1]*apply(x[,xall_names_num, drop=F], 2, var) + 2)
  b2 = (g2 - 1)*tau_x[1]*apply(x[,xall_names_num, drop=F], 2, var)
  names(g2) <- xall_names_num
  names(b2) <- xall_names_num

  K <- init_k
  a <- 1

  ###------------------------------------------------------------------------###
  #### 1 - Create Shells for storing Gibbs Results                          ####
  ###------------------------------------------------------------------------###

  alpha <- numeric(length = store_l)

  curr_class_list <- 1:K
  class_names <- paste0('c',curr_class_list)

  c_shell<-matrix(NA, nrow=n, ncol=1) # shell for indicators

  class_shell <- matrix(NA, nrow=n, ncol=store_l)

  psi_shell <- matrix(NA, nrow= 1, ncol = K) # shell for Y's cond. var
  colnames(psi_shell) <- class_names

  beta_shell <- matrix(NA, nrow = nparams, ncol = K)
  colnames(beta_shell) <- class_names
  rownames(beta_shell) <- colnames(x)

  c_b_shell <- vector(mode = 'list', length = n_bin_p) # -1 for trt
  names(c_b_shell) <- xall_names_bin

  c_n_shell <- vector(mode = 'list', length = n_num_p)
  names(c_n_shell) <- xall_names_num

  for( p in xall_names_bin ){
    c_b_shell[[p]] <- matrix(NA, nrow = 1, ncol = K)

    c_b_shell[[p]][1,] <- rep(.5, K)

    colnames(c_b_shell[[p]]) <- class_names
  }

  for( p in xall_names_num ){
    c_n_shell[[p]] <- list(matrix(NA, nrow=1, ncol=K),
                           matrix(NA, nrow=1, ncol=K))

    c_n_shell[[p]][[2]][1, ] <- rep(0, K)
    c_n_shell[[p]][[2]][1, ] <- rep(20^2, K)

    colnames(c_n_shell[[p]][[1]]) <- class_names
    colnames(c_n_shell[[p]][[2]]) <- class_names
  }

  c_b_new <- numeric(length = n_bin_p)
  names(c_b_new) <- xall_names_bin

  c_n_new <- matrix(NA, nrow = 2, ncol = n_num_p)
  colnames(c_n_new) <- xall_names_num

  if(is.null(d_test)){
    predictions <- vector(mode = 'list', length = 1)
    predictions[[1]] <- matrix(nrow = n, ncol = store_l)
    names(predictions) <- c('train')
  }else{
    predictions <- vector(mode = 'list', length = 2)
    predictions[[1]] <- matrix(nrow = n, ncol = store_l)
    predictions[[2]] <- matrix(nrow = nrow(d_test), ncol = store_l)
    names(predictions) <- c('train','test')
  }

  ###------------------------------------------------------------------------###
  #### 2 - Set Initial Values                                               ####
  ###------------------------------------------------------------------------###

  # initially assign everyone to one of K clusters randomly with uniform prob
  c_shell[,1] <- sample(x = class_names, size = n,
                        prob = rep(1/K,K), replace = T)

  beta_start <- matrix(0, nrow=nparams, ncol=K)
  colnames(beta_start) <- class_names

  for(k in class_names){
    beta_shell[,k] <- beta_prior_mean
  }

  for(p in xall_names_bin ){
    c_b_shell[[p]][1,] <- rep(1, K)
  }

  for(p in xall_names_num ){
    c_n_shell[[p]][[2]][ 1 , ] <- rep(0, K)
    c_n_shell[[p]][[2]][ 1 , ] <- rep(20^2, K)
  }

  psi_shell[1,1:K] <- rep(500000, K)

  ###------------------------------------------------------------------------###
  #### 3 - Gibbs Sampler                                                    ####
  ###------------------------------------------------------------------------###
  message_seq <- seq(2, iter, round(iter/20) )
  cat(paste0('Iteration ',2,' of ', iter,' (burn-in) \n'))
  message_step <- (iter)/10
  for(i in 2:iter){

    ## print message to user
    if( i %% message_step == 1 ) cat(paste0('Iteration ',i,' of ',
                                            iter,
                                            ifelse(i>burnin,' (sampling) ',
                                                   ' (burn-in)') ,
                                            '\n'))

    # compute size of each existing cluster
    class_ind<-c_shell[,1]
    nvec<-table(factor(class_ind,levels = class_names) )

    ###----------------------------------------------------------------------###
    #### 3.1 - update Concentration Parameter                               ####
    ###----------------------------------------------------------------------###

    a_star <- rnorm(n = 1, mean =  a, sd = 1)
    if(a_star>0){
      r_num <- cond_post_alpha(a_star, n, K, nvec)
      r_denom <- cond_post_alpha(a, n, K, nvec)
      r <- exp(r_num - r_denom)
      accept <- rbinom(1,1, min(r,1) )==1
      if(accept){ a <- a_star }
    }

    ###----------------------------------------------------------------------###
    #### 3.2 - update parameters conditional on classification              ####
    ###----------------------------------------------------------------------###
    K <- length(unique(class_names))
    for(k in class_names){ # cycle through existing clusters - update parameters

      ck_ind <- class_ind == k

      y_ck <- y[ck_ind]

      x_ck <- x[ck_ind,, drop=FALSE]

      psi_shell[1, k]<-rcond_post_psi(beta = beta_shell[,k, drop=F],
                                      y = y_ck, xm = x_ck,
                                      g1 = g1, b1 = b1)

      beta_shell[,k] <- rcond_post_beta(y=y_ck,
                                        xm=x_ck,
                                        psi = psi_shell[1,k],
                                        beta_prior_mean = beta_prior_mean,
                                        beta_prior_var = beta_prior_var)

      for( p in xall_names_bin ){
        c_b_shell[[p]][1, k] <- rcond_post_mu_trt(x_ck = x_ck[, p])
      }

      for( p in xall_names_num ){
        c_n_shell[[p]][[1]][1,k] <- rcond_post_mu_x(x_ck = x_ck[, p],
                                                    phi_x = c_n_shell[[p]][[2]][1,k],
                                                    lambda = prior_means[p],
                                                    tau = prior_var[p])

        c_n_shell[[p]][[2]][1,k] <- rcond_post_phi(mu_x = c_n_shell[[p]][[1]][1, k],
                                                   x = x_ck[ , p],
                                                   g2 = g2[p], b2 = b2[p])
      }

    }

    ###----------------------------------------------------------------------###
    #### 3.3 - update classification conditional on parameters              ####
    ###----------------------------------------------------------------------###

    ## draw parameters for potentially new cluster from priors
    ### draw beta, phi, covariate parameters

    beta_new <- mvtnorm::rmvnorm(n = 1, mean = beta_prior_mean,
                                 sigma = diag(beta_prior_var) )
    psi_new <- invgamma::rinvgamma(n = 1, shape = g1, rate = b1)

    for(p in xall_names_bin){
      c_b_new[p] <- rbeta(n = 1, shape1 = 1, shape2 = 1)
    }

    for(p in xall_names_num){
      c_n_new[1, p] <- rnorm(n = 1, prior_means[p], sqrt(prior_var[p]) )
      c_n_new[2, p] <- invgamma::rinvgamma(n = 1, shape = g2[p], rate = b2[p])
    }

    name_new <- paste0('c', max(as.numeric(substr(class_names, 2,10))) + 1)

    ## compute post prob of membership to existing and new cluster,
    ## then classify

    uniq_clabs = colnames(beta_shell)

    weights <- class_update_train_fdp(n = n, K = length(class_names), alpha = a,
                                      name_new =  name_new,
                                      uniq_clabs = uniq_clabs,
                                      clabs = class_ind,

                                      y = y, x = x,
                                      x_cat_shell = c_b_shell , x_num_shell = c_n_shell,

                                      ## col number of each covar...index with base 0
                                      cat_idx=xall_names_bin, num_idx=xall_names_num,

                                      beta_shell = beta_shell,
                                      psi_shell = psi_shell,

                                      beta_new=beta_new, psi_new=psi_new,
                                      cat_new = c_b_new, num_new=c_n_new)

    c_shell[,1] <- apply(weights, 1, FUN = sample,
                         x=c(colnames(beta_shell), name_new), size=1, replace=T)

    if( i > burnin ){

      ###--------------------------------------------------------------------###
      #### 4.0 - Predictions on a Training and Test set  & Store Results    ####
      ###--------------------------------------------------------------------###

      test_pred <- post_mean_test_fdp(n=nt,K=K,alpha=a, name_new=name_new,
                                      uniq_clabs = uniq_clabs, x=xt,
                                      x_cat_shell=c_b_shell, x_num_shell=c_n_shell,
                                      cat_idx=xall_names_bin, num_idx=xall_names_num,
                                      cat_new=c_b_new, num_new=c_n_new,
                                      clabs=class_ind, beta_shell, beta_new)
      if( is.null(d_test) ){
        predictions[['train']][, i-burnin  ] <- test_pred[test_ind==0]
      }else{
        predictions[['train']][, i-burnin  ] <- test_pred[test_ind==0]
        predictions[['test']][, i-burnin  ] <- test_pred[test_ind==1]
      }
    }

    ###----------------------------------------------------------------------###
    #### 5.0 - update shell dimensions as clusters die/ are born            ####
    ###----------------------------------------------------------------------###

    new_class_names <- unique(c_shell[,1])
    K_new <- length(new_class_names)

    # Grow shells by adding labels/parameters of new clusters
    new_class_name <- setdiff(new_class_names, class_names)
    if(length(new_class_name)>0){ # new cluster was born

      beta_shell <- cbind(beta_shell, t(beta_new) )
      colnames(beta_shell)[K+1] <- new_class_name

      psi_shell <- cbind(psi_shell, psi_new)
      colnames(psi_shell)[K+1] <- new_class_name

      for(p in xall_names_bin ){
        c_b_shell[[p]] <- cbind( c_b_shell[[p]], c_b_new[p] )
        colnames(c_b_shell[[p]])[K+1] <- new_class_name
      }

      for(p in xall_names_num ){
        c_n_shell[[p]][[1]] <- cbind(c_n_shell[[p]][[1]], c_n_new[1, p])
        colnames(c_n_shell[[p]][[1]])[K+1] <- new_class_name

        c_n_shell[[p]][[2]] <- cbind(c_n_shell[[p]][[2]], c_n_new[2, p])
        colnames(c_n_shell[[p]][[2]])[K+1] <- new_class_name
      }
    }

    # shrink shells by deleting labels of clusters that have died.
    dead_classes <- setdiff(class_names, new_class_names)
    if(length(dead_classes) > 0){ # clusters have died
      beta_shell <- beta_shell[,! colnames(beta_shell) %in% dead_classes, drop=F]

      psi_shell <- psi_shell[, ! colnames(psi_shell) %in% dead_classes, drop=F]

      for(p in xall_names_bin){
        c_b_shell[[p]] <- c_b_shell[[p]][, ! colnames(c_b_shell[[p]]) %in% dead_classes, drop=F]
      }

      for(p in xall_names_num){
        c_n_shell[[p]][[1]] <- c_n_shell[[p]][[1]][, ! colnames(c_n_shell[[p]][[1]]) %in% dead_classes, drop=F]
        c_n_shell[[p]][[2]] <- c_n_shell[[p]][[2]][, ! colnames(c_n_shell[[p]][[2]]) %in% dead_classes, drop=F]
      }
    }

    K <- K_new
    class_names <- new_class_names

  }

  results <- predictions
  cat('Sampling Complete.')
  return(results)
}