##########################################################
##                                                      ##
##  Simulation study for "Estimating the retransformed  ##
##  mean in a heteroscedastic two-part model" by Welsh  ##
##  and Zhou.                                           ##
##                                                      ##
##  Developed by Hao Cheng for Andrew Zhou              ##
##  Date: December 10, 2004                             ##
##                                                      ##
##########################################################



#----- clear environment -----
rm(list=ls())
print(Sys.time())



##########################################################
##                                                      ##
## simulation setting                                   ##
##                                                      ##
##########################################################

#----- input data, generate covariates and response -----
inputfile <- ""
while(inputfile=="")
{
    cat("Please specify the name of input data file: ")
    inputfile <- readLines(n=1)
    inputfile <- gsub(" ", "", inputfile)
}
if(!file.exists(inputfile))
    stop("File ", inputfile, " does not exist. Stop!")

dataset <- as.matrix(read.table(inputfile))
X <- cbind(1, dataset[,-1])
Y <- as.vector(dataset[,1])

#----- sample size -----
nSample <- nrow(X)

#----- number of covariates -----
#      nCovariate = 1 + (number of actual covariates)
#      covariates are given in format (1, covariate_1, ..., covariate_n),
#      where n = nCovariate - 1
nCovariate <- ncol(X)

#----- number of data sets -----
nDataset <- 1

#----- option on estimating parameters beta and gamma -----
#      option.bg = 0: maximum-likelihood estimator, solve estimating equation
#      option.bg = 1: maximum-likelihood estimator, maximize log-likelihood function
cat("\n")
cat("Please choose method for estimating parameters beta and theta\n")
cat("  1: MLE estimator, maximize log-likelihood function (default)\n")
cat("  2: MLE estimator, solve estimating equation\n")
option.bg <- readLines(n=1)
option.bg <- as.numeric(gsub(" ", "", option.bg))
if(is.na(option.bg) | (option.bg!=1 & option.bg!=2))
{
    option.bg<-1
    cat("Not a valid selection. Use the default method 1.\n")
}

#----- real parameter values -----
#      For real data, they are initial guess of parameters
#      For simulation, they are real parameter values used to generate data sets
cat("\n")
cat("Do you want to give initial guess of parameters (default: all zeros)? (Y/N)")
answer <- readLines(n=1)
bStop <- FALSE
while(!bStop)
{
    answer <- gsub(" ", "", answer)
    answer <- unlist(strsplit(answer, split=NULL))[1]
    if(is.na(answer)) answer<-' '
    if(answer=='n' | answer=='N')
    {
        bStop <- TRUE
        bInitpar <- FALSE
    }
    else if(answer=='y' | answer=='Y')
    {
        bStop <- TRUE
        bInitpar <- TRUE
    }
    else
    {
        cat("Unknown answer. Please input 'y' or 'n': ")
        answer <- readLines(n=1)
    }
}

if(bInitpar)
{
    cat("\n")
    cat("Please input initial guess of parameter alpha:\n")
    alpha.0 <- scan(quiet=TRUE, n=nCovariate)
    temp <- length(alpha.0)
    if(temp<nCovariate) alpha.0<-c(alpha.0,rep(0,nCovariate-temp))
    cat("\n")
    cat("Please input initial guess of parameter beta:\n")
    beta.0 <- scan(quiet=TRUE, n=nCovariate)
    temp <- length(beta.0)
    if(temp<nCovariate) beta.0<-c(beta.0,rep(0,nCovariate-temp))
    cat("\n")
    cat("Please input initial guess of parameter theta:\n")
    gamma.0 <- scan(quiet=TRUE, n=nCovariate)
    temp <- length(gamma.0)
    if(temp<nCovariate) gamma.0<-c(gamma.0,rep(0,nCovariate-temp))
} else
{
    alpha.0 <- rep(0, nCovariate)
    beta.0 <- rep(0, nCovariate)
    gamma.0 <- rep(0, nCovariate)
}

cat("alpha.0 =", alpha.0, "\n")
cat("beta.0  =", beta.0, "\n")
cat("gamma.0 =", gamma.0, "\n")

#----- (1-r) confidence interval -----
cat("\n")
cat("Please specify (1-r) confidence interval (default r=0.05):\n")
cat("  r = ")
confidence.level <- readLines(n=1)
confidence.level <- as.numeric(gsub(" ", "", confidence.level))
if(is.na(confidence.level))
{
    confidence.level<-0.05
    cat("  r = ", confidence.level, "\n")
}

#----- option on calculating confidence interval of mean estimators -----
#      option.ConfInt = 0: use formulas given in the paper
#      option.ConfInt = 1: use boostrap
cat("\n")
cat("Do you want to calculate bootstrap confidence interval? (Y/N)")
answer <- readLines(n=1)
bStop <- FALSE
while(!bStop)
{
    answer <- gsub(" ", "", answer)
    answer <- unlist(strsplit(answer, split=NULL))[1]
    if(is.na(answer)) answer<-' '
    if(answer=='n' | answer=='N')
    {
        bStop <- TRUE
        option.ConfInt <- 0
    }
    else if(answer=='y' | answer=='Y')
    {
        bStop <- TRUE
        option.ConfInt <- 1
    }
    else
    {
        cat("Unknown answer. Please input 'y' or 'n': ")
        answer <- readLines(n=1)
    }
}

if(option.ConfInt==1)
{
    #----- size of bootstrap samples -----
    cat("\n")
    cat("Please specify bootstrap sample size (default 100): ")
    nBootstrap <- readLines(n=1)
    nBootstrap <- as.numeric(gsub(" ", "", nBootstrap))
    if(is.na(nBootstrap))
    {
        nBootstrap<-100
        cat("Not a valid number. Use deault bootstrap sample size 100.\n")
    }
    
    #----- random number generator seed -----
    cat("\n")
    cat("Do you want to assign seed for random number generator? (Y/N)")
    answer <- readLines(n=1)
    bStop <- FALSE
    while(!bStop)
    {
        answer <- gsub(" ", "", answer)
        answer <- unlist(strsplit(answer, split=NULL))[1]
        if(is.na(answer)) answer<-' '
        if(answer=='n' | answer=='N')
        {
            bStop <- TRUE
            bSeed <- FALSE
        }
        else if(answer=='y' | answer=='Y')
        {
            bStop <- TRUE
            bSeed <- TRUE
        }
        else
        {
            cat("Unknown answer. Please input 'y' or 'n': ")
            answer <- readLines(n=1)
        }
    }
    if(bSeed)
    {
        cat("\n")
        cat("Please assign seed for random number generator: ")
        seed <- readLines(n=1)
        seed <- as.numeric(gsub(" ", "", seed))
        if(is.na(seed)) bSeed<-FALSE
    }
    if(bSeed) set.seed(seed) else cat("Seed not assigned. \n")
}

#----- define covariates used for zero response (alpha) computation -----
#
#      nAlpha: number of covariates used to calculate alpha
#      nAlpha <= nCovariate
#
#      XalphaIndicator is a vector that has nCovariate components.
#      It shows which components of covariates are used
#      when estimating parameter alpha
#      sum(indicator) must equal nAlpha
nAlpha <- nCovariate
XalphaIndicator <- rep(1,nCovariate)

#----- define covariates used for estimating mean in transformed space (beta) -----
nBeta <- nCovariate
XbetaIndicator <- rep(1,nCovariate)

#----- define covariates used for estimating variance in transformed space (gamma) -----
nGamma <- nCovariate
XgammaIndicator <- rep(1,nCovariate)


#----- load functions -----
source("estmean.r")

#----- check paramer setting validality -----
if(!check.data()) stop()




##########################################################
##                                                      ##
## specify transformation function h()                  ##
##                                                      ##
##########################################################

# transformation function h()
fh <- function(x)
{
    return(exp(x))
}


# inverse of transformation function h()
fhinv <- function(x)
{
    return(log(x))
}


# derivative of transformation function h()
dfh <- function(x)
{
    return(exp(x))
}



##########################################################
##                                                      ##
## specify mean function m() in transformed space       ##
##                                                      ##
## mean is only related with parameter beta             ##
## both beta and x have nBeta componets                 ##
##                                                      ##
##########################################################

# mean in transformed space
# m = beta*x
fm <- function(beta, xbeta)
{
    return(xbeta %*% beta)
}

# partial derivative of fm on beta
dfm.beta <- function(beta, xbeta)
{
    return(xbeta)
}



##########################################################
##                                                      ##
## specify vatiance function g() in transformed space   ##
##                                                      ##
## variance is related with parameters beta and gamma   ##
## both beta and xbeta have nBeta componets             ##
## both gamma and xgamma have nGamma components         ##
##                                                      ##
##########################################################

# variance in transformed space
# g = exp( gamma*x / 2 )
fg <- function(beta, gamma, xbeta, xgamma)
{
    fval <- exp(0.5 * xgamma %*% gamma)
    return(fval)
}

# partial derivative of fg on beta
dfg.beta <- function(beta, gamma, xbeta, xgamma)
{
    y <- xbeta * 0
    return(y)
}

# partial derivative of fg on gamma
dfg.gamma <- function(beta, gamma, xbeta, xgamma)
{
    y <- xgamma
    y <- y * as.vector(0.5 * exp(0.5 * xgamma %*% gamma))
    return(y)
}

# 2nd order partial derivative of fg on beta
d2fg.beta <- function(beta, gamma, xbeta, xgamma)
{
    y <- matrix(0, nBeta, nBeta)
    return(y)
}

# 2nd order partial derivative of fg on beta and gamma
d2fg.betagamma <- function(beta, gamma, xbeta, xgamma)
{
    y <- matrix(0, nBeta, nGamma)
    return(y)
}

# 2nd order partial derivative of fg on gamma and beta
d2fg.gammabeta <- function(beta, gamma, xbeta, xgamma)
{
    y <- matrix(0, nGamma, nBeta)
    return(y)
}

# 2nd order partial derivative of fg on gamma
d2fg.gamma <- function(beta, gamma, xbeta, xgamma)
{
    k1 <- as.vector(0.25 * exp(0.5 * xgamma %*% gamma))
    D <- k1 * outer(xgamma, xgamma)
    return(D)
}

# log of fg
logfg <- function(beta, gamma, xbeta, xgamma)
{
    return(0.5 * xgamma %*% gamma)
}



##########################################################
##                                                      ##
## execute computation                                  ##
##                                                      ##
##########################################################

# allocate memory to store results for each dataset
if(option.ConfInt==0)
{
    result <- matrix(0, nDataset, 8)
}
if(option.ConfInt==1)
{
    result <- matrix(0, nDataset, 12)
    tscore <- matrix(0, nBootstrap, 2)
}

# generate covariates
Xalpha <- X[, XalphaIndicator>0]
Xbeta <- X[, XbetaIndicator>0]
Xgamma <- X[, XgammaIndicator>0]

##
## estimate parameters
##
alpha.init <- alpha.0[XalphaIndicator>0]
beta.init <- beta.0[XbetaIndicator>0]
gamma.init <- gamma.0[XgammaIndicator>0]
par <- estpar(Y, Xalpha, Xbeta, Xgamma, alpha.init, beta.init, gamma.init)

parfile <- paste(unlist(strsplit(inputfile,".",fixed=TRUE))[1], "_par.dat", sep="")
sink(parfile)
outpar(par)
sink()


resultfile <- paste(unlist(strsplit(inputfile,".",fixed=TRUE))[1], "_result.dat", sep="")
file.remove(resultfile)
sink(resultfile)
cat("Source data file:       ", inputfile, "\n")
cat("Number of observations: ", nSample, "\n")
cat("Number of covariates:   ", nCovariate-1, "\n")
cat("----------------------------------------------------------\n\n")
sink()


#----- covariates value -----
#cat("\n")
#cat("Please input covariates values:\n")
#incov <- scan(quiet=TRUE, n=nCovariate-1)
#cov <- c(1, incov)
#estmean(par, result, tscore)
#cat("\n")
#cat("Do you want to see estimation results for another covariate? (Y/N) ")
answer <- 'Y'
bStop <- FALSE
while(!bStop)
{
    answer <- gsub(" ", "", answer)
    answer <- unlist(strsplit(answer, split=NULL))[1]
    if(is.na(answer)) answer<-' '
    if(answer=='n' | answer=='N')
    {
        bStop <- TRUE
    }
    else if(answer=='y' | answer=='Y')
    {
        cat("\n")
        cat("Please input covariates values:\n")
        incov <- scan(quiet=TRUE, n=nCovariate-1)
        temp <- length(incov)
        if(temp<nCovariate-1) incov<-c(incov,rep(0,nCovariate-1-temp))
        cat("Covariate = ", incov, "\n")
        cov <- c(1, incov)
        estmean(par, result, tscore)
        cat("\n")
        cat("Do you want to see estimation results for another covariate? (Y/N) ")
        answer <- readLines(n=1)
    }
    else
    {
        cat("Unknown answer. Please input 'y' or 'n': ")
        answer <- readLines(n=1)
    }
}
print(Sys.time())