# bfvariance-functions.ssc
#
# function to calculate adjusted TFR with and without var effect
# 
#
# Author: Hans-Peter Kohler
# email:  kohler@demogr.mpg.de



# F.bf.calc is the primary function to calculate the adjustment
# of the parity-specific TFR with and without variance effects
#
# 

F.bf.calc <- function(nfx, age, years, i.prec=0.001){ 

  # calculation of the usual TFR,
  # mean age, variance, and
  # centralized third moment
  # of the fertility schedule
  p.age <- age+0.5
  tfr <- apply(nfx, 2, sum)
  mab <- apply(nfx * (age + 0.5), 2, sum) / tfr
  vab <- apply(nfx * (age + 0.5)^2, 2, sum) / tfr - mab^2
  age.0 <- sweep(matrix(p.age,length(p.age),length(years),byrow=F),2,mab,"-")
  kub <- apply(nfx*(age.0^3), 2, sum) / tfr
  names(tfr) <- paste("yr",years)
  names(mab) <- paste("yr",years)
  names(vab) <- paste("yr",years)

  # smooth the time series for TFR, mab, vab,kub
  tfr.s <- as.vector(smooth(tfr), twice=T)
  names(tfr.s) <- paste("yr",years)
  mab.s <- as.vector(smooth(mab), twice=T)
  names(mab.s) <- paste("yr",years)
  vab.s <- as.vector(smooth(vab), twice=T)
  names(vab.s) <- paste("yr",years)
  kub <- as.vector(smooth(kub), twice=T)
  names(kub) <- paste("yr",years)

  # implementation of BF formula
  r.bf <- c(NA, 0.5 * (mab.s[3:length(years)] - 
                     mab.s[1:(length(years)-2)]), NA)
  tfr.adj.bf <- tfr.s / (1 - r.bf)
  names(r.bf) <- paste("yr",years)
  names(tfr.adj.bf) <- paste("yr",years)


  # calclulate delta from observed variance
  # (using the second and second to last observation of delta
  # for the first and last year of the data to avoid a reduction
  # of length of the time series)

  delta <- 0.25 * log(vab.s[3:length(years)]/
                      vab.s[1:(length(years)-2)])
  delta <- c(delta[1], delta, delta[length(delta)])
  names(delta) <- paste("yr",years)

  # calculate bias corrected bf-r
  # based on a one-step correction (Result 11)
  delta.change <- c(NA, 0.5 * (delta[3:length(years)] - 
                               delta[1:(length(years)-2)]), NA)
  r.bf.biascorr <- r.bf + vab.s / (1 - r.bf) * (2*delta^2 + delta.change)

  #iterate estimation for the calculation of variance effects (Result 13)
  r.res <- F.iter.calc(mab.s,vab.s,kub,years,delta,i.prec)

  delta <- c(NA,delta[-c(1,length(delta))],NA)
  gamma <- r.res$gamma
  s2 <- r.res$s2
  a.bar <- r.res$a.bar

  # calculate adjusted TFR based on "correct" tempo gamma
  tfr.adj.var <- tfr.s/(1 - gamma)

  list(
       years = years,
       tfr = tfr,
       tfr.s = tfr.s,
       mab = mab,
       mab.s = mab.s,
       vab = vab,
       vab.s = vab.s,
       kub = kub,
       r.bf = r.bf,
       tfr.adj.bf = tfr.adj.bf,
       r.bf.biascorr = r.bf.biascorr,
       gamma = gamma,
       delta = delta,
       mab.adj.var = a.bar,
       vab.adj.var = s2,
       tfr.adj.var = tfr.adj.var
       )
}


# F.iter.calc performs the iteration in Result 13
# and iteratively calculates gamma and then corrects
# for the distortions caused by variance effects
F.iter.calc <- function(mab,vab,kub,yrs, delta.hat, prec=0.0001){
	a.hat <- mab
        s2.hat <- vab
	i <- 0
	g.diff <- 1

        #step 1
        gamma.hat <- 0.5*(mab[3:length(yrs)]-mab[1:(length(yrs)-2)])
	gamma.hat <- c(gamma.hat[1], gamma.hat, gamma.hat[length(gamma.hat)])
	names(gamma.hat) <- paste("yr",yrs)

        #iterate for calculation of gamma
	while (g.diff > prec){
          i <- i + 1

        #step 2
          s2.hat.old <- s2.hat
          s2.hat <- vab  +  (delta.hat/(1-gamma.hat) * s2.hat)^2 +
            delta.hat/(1-gamma.hat) * kub

        #step 3
          a.hat.old <- a.hat
          a.hat <- mab + delta.hat/(1-gamma.hat)*s2.hat
          
        #step 4
          gamma.hat.old <- gamma.hat
          gamma.hat <- 0.5*(a.hat[3:length(yrs)]-a.hat[1:(length(yrs)-2)])
          gamma.hat <- c(gamma.hat[1], gamma.hat, gamma.hat[length(gamma.hat)])
        # avoid estimates that are implausible
          gamma.hat[gamma.hat > 0.5] <- 0.5
          gamma.hat[gamma.hat < -0.5] <- -0.5
          names(gamma.hat) <- paste("yr",yrs)

        # calculate convergence criterion
          diff.vec <- abs(c(a.hat,s2.hat,gamma.hat) - 
                          c(a.hat.old,s2.hat.old,gamma.hat.old))
          g.diff <- max(diff.vec[!is.na(diff.vec)])

          print(paste("iteration",i,"parameters for gamma"))
          print(g.diff)
          print(gamma.hat, digits=4)

        }


        # put back missing values
	gamma.hat <- c(NA,gamma.hat[-c(1,length(gamma.hat))],NA)
	a.hat <- c(NA,a.hat[-c(1,length(a.hat))],NA)
	s2.hat <- c(NA,s2.hat[-c(1,length(s2.hat))],NA)

        
        names(a.hat) <- paste("yr",yrs)
	names(s2.hat) <- paste("yr",yrs)
	names(gamma.hat) <- paste("yr",yrs)
	
	list(a.bar = a.hat, s2 = s2.hat,
		gamma = gamma.hat, delta = delta.hat)
}
	

F.allorder <- function(nfx.tot, age, os.result){
  #direct calculation from nFx with all orders
  tfr <- apply(nfx.tot, 2, sum)
  mab <- apply(nfx.tot * (age + 0.5), 2, sum) / tfr
  vab <- apply(nfx.tot * (age + 0.5)^2, 2, sum) / tfr - mab^2
  f.years <- os.result[[1]]$years

  #indirect calculation from order specific results
  i.tfr <- (os.result[[1]]$tfr.s + os.result[[2]]$tfr.s +
            os.result[[3]]$tfr.s + os.result[[4]]$tfr.s)
  i.mab <- (os.result[[1]]$tfr.s * os.result[[1]]$mab.s +
            os.result[[2]]$tfr.s * os.result[[2]]$mab.s +
            os.result[[3]]$tfr.s * os.result[[3]]$mab.s +
            os.result[[4]]$tfr.s * os.result[[4]]$mab.s) / i.tfr
  i.vab <- (os.result[[1]]$tfr.s * os.result[[1]]$vab.s +
            os.result[[2]]$tfr.s * os.result[[2]]$vab.s +
            os.result[[3]]$tfr.s * os.result[[3]]$vab.s +
            os.result[[4]]$tfr.s * os.result[[4]]$vab.s) / i.tfr + 
              (os.result[[1]]$tfr.s * (os.result[[1]]$mab.s-i.mab)^2 +
               os.result[[2]]$tfr.s * (os.result[[2]]$mab.s-i.mab)^2 +
               os.result[[3]]$tfr.s * (os.result[[3]]$mab.s-i.mab)^2 +
               os.result[[4]]$tfr.s * (os.result[[4]]$mab.s-i.mab)^2) / i.tfr

  tfr.adj.var <- (os.result[[1]]$tfr.adj.var + 
                  os.result[[2]]$tfr.adj.var +
                  os.result[[3]]$tfr.adj.var + 
                  os.result[[4]]$tfr.adj.var)

  mab.adj.var <- (os.result[[1]]$tfr.s * os.result[[1]]$mab.adj.var +
                  os.result[[2]]$tfr.s * os.result[[2]]$mab.adj.var +
                  os.result[[3]]$tfr.s * os.result[[3]]$mab.adj.var +
                  os.result[[4]]$tfr.s * os.result[[4]]$mab.adj.var) / i.tfr

  gamma <- (os.result[[1]]$tfr.s * os.result[[1]]$gamma +
            os.result[[2]]$tfr.s * os.result[[2]]$gamma +
            os.result[[3]]$tfr.s * os.result[[3]]$gamma +
            os.result[[4]]$tfr.s * os.result[[4]]$gamma) / i.tfr

  delta <- (os.result[[1]]$tfr.s * os.result[[1]]$delta +
            os.result[[2]]$tfr.s * os.result[[2]]$delta +
            os.result[[3]]$tfr.s * os.result[[3]]$delta +
            os.result[[4]]$tfr.s * os.result[[4]]$delta) / i.tfr

  r.bf <- (os.result[[1]]$tfr.s * os.result[[1]]$r.bf +
           os.result[[2]]$tfr.s * os.result[[2]]$r.bf +
           os.result[[3]]$tfr.s * os.result[[3]]$r.bf +
           os.result[[4]]$tfr.s * os.result[[4]]$r.bf) / i.tfr

  vab.adj.var <- (os.result[[1]]$tfr.s * os.result[[1]]$vab.adj.var +
                  os.result[[2]]$tfr.s * os.result[[2]]$vab.adj.var +
                  os.result[[3]]$tfr.s * os.result[[3]]$vab.adj.var +
                  os.result[[4]]$tfr.s * os.result[[4]]$vab.adj.var) / i.tfr + 
                    (os.result[[1]]$tfr.s * (os.result[[1]]$mab.s-mab.adj.var)^2 +
                     os.result[[2]]$tfr.s * (os.result[[2]]$mab.s-mab.adj.var)^2 +
                     os.result[[3]]$tfr.s * (os.result[[3]]$mab.s-mab.adj.var)^2 +
                     os.result[[4]]$tfr.s * (os.result[[4]]$mab.s-mab.adj.var)^2) /
                       i.tfr

  list(
       years = f.years,
       tfr.s=	i.tfr,
       tfr = tfr,
       mab.s = i.mab,
       mab = mab,
       r.bf = r.bf,
       vab.s = i.vab,
       vab = vab,
       tfr.adj.bf = (os.result[[1]]$tfr.adj.bf + 
                     os.result[[2]]$tfr.adj.bf +
                     os.result[[3]]$tfr.adj.bf + 
                     os.result[[4]]$tfr.adj.bf),  
       tfr.adj.var = tfr.adj.var,
       mab.adj.var = mab.adj.var,
       gamma = gamma,
       delta= delta,
       vab.adj.var = vab.adj.var
       )
}