library(matrixStats) # for numerically stable logsumexps

options(amelia.parallel="no",
        amelia.ncpus=1)
library(Amelia)

source("pl_methods.R")
source("measerr_methods_2.R") ## for my more generic function.

run_simulation_depvar <- function(df, result, outcome_formula = y ~ x + z, coder_formulas = list(y.obs.0 ~ 1, y.obs.1 ~ 1), proxy_formula = w_pred ~ y.obs.1+y.obs.0+y){
    (accuracy <- df[,mean(w_pred==y)])
    result <- append(result, list(accuracy=accuracy))
    (error.cor.z <- cor(df$x, df$w_pred - df$z))
    (error.cor.x <- cor(df$x, df$w_pred - df$y))
    (error.cor.y <- cor(df$y, df$y - df$w_pred))
    result <- append(result, list(error.cor.x = error.cor.x,
                                  error.cor.z = error.cor.z,
                                  error.cor.y = error.cor.y))

    model.null <- glm(y~1, data=df,family=binomial(link='logit'))
    (model.true <- glm(y ~ x + z, data=df,family=binomial(link='logit')))
    (lik.ratio <- exp(logLik(model.true) - logLik(model.null)))

    true.ci.Bxy <- confint(model.true)['x',]
    true.ci.Bzy <- confint(model.true)['z',]


    result <- append(result, list(lik.ratio=lik.ratio))

    result <- append(result, list(Bxy.est.true=coef(model.true)['x'],
                                  Bzy.est.true=coef(model.true)['z'],
                                  Bxy.ci.upper.true = true.ci.Bxy[2],
                                  Bxy.ci.lower.true = true.ci.Bxy[1],
                                  Bzy.ci.upper.true = true.ci.Bzy[2],
                                  Bzy.ci.lower.true = true.ci.Bzy[1]))
                                  
    (model.naive <- lm(y~w_pred+z, data=df))
    
    naive.ci.Bxy <- confint(model.naive)['w_pred',]
    naive.ci.Bzy <- confint(model.naive)['z',]

    result <- append(result, list(Bxy.est.naive=coef(model.naive)['w_pred'],
                                  Bzy.est.naive=coef(model.naive)['z'],
                                  Bxy.ci.upper.naive = naive.ci.Bxy[2],
                                  Bxy.ci.lower.naive = naive.ci.Bxy[1],
                                  Bzy.ci.upper.naive = naive.ci.Bzy[2],
                                  Bzy.ci.lower.naive = naive.ci.Bzy[1]))



    loa0.feasible <- glm(y.obs.0 ~ x + z, data = df[!(is.na(y.obs.0))], family=binomial(link='logit'))

    loa0.ci.Bxy <- confint(loa0.feasible)['x',]
    loa0.ci.Bzy <- confint(loa0.feasible)['z',]

    result <- append(result, list(Bxy.est.loa0.feasible=coef(loa0.feasible)['x'],
                                  Bzy.est.loa0.feasible=coef(loa0.feasible)['z'],
                                  Bxy.ci.upper.loa0.feasible = loa0.ci.Bxy[2],
                                  Bxy.ci.lower.loa0.feasible = loa0.ci.Bxy[1],
                                  Bzy.ci.upper.loa0.feasible = loa0.ci.Bzy[2],
                                  Bzy.ci.lower.loa0.feasible = loa0.ci.Bzy[1]))


    ## df.loa0.mle <- copy(df)
    ## df.loa0.mle[,y:=y.obs.0]
    ## loa0.mle <- measerr_mle_dv(df.loa0.mle, outcome_formula=outcome_formula, proxy_formula=proxy_formula)
    ## fisher.info <- solve(loa0.mle$hessian)
    ## coef <- loa0.mle$par
    ## ci.upper <- coef + sqrt(diag(fisher.info)) * 1.96
    ## ci.lower <- coef - sqrt(diag(fisher.info)) * 1.96

    ## result <- append(result, list(Bxy.est.loa0.mle=coef['x'],
    ##                               Bzy.est.loa0.mle=coef['z'],
    ##                               Bxy.ci.upper.loa0.mle = ci.upper['x'],
    ##                               Bxy.ci.lower.loa0.mle = ci.lower['x'],
    ##                               Bzy.ci.upper.loa0.mle = ci.upper['z'],
    ##                               Bzy.ci.lower.loa0.mle = ci.upper['z']))

    loco.feasible <- glm(y.obs.0 ~ x + z, data = df[(!is.na(y.obs.0)) & (y.obs.1 == y.obs.0)], family=binomial(link='logit'))

    loco.feasible.ci.Bxy <- confint(loco.feasible)['x',]
    loco.feasible.ci.Bzy <- confint(loco.feasible)['z',]

    result <- append(result, list(Bxy.est.loco.feasible=coef(loco.feasible)['x'],
                                  Bzy.est.loco.feasible=coef(loco.feasible)['z'],
                                  Bxy.ci.upper.loco.feasible = loco.feasible.ci.Bxy[2],
                                  Bxy.ci.lower.loco.feasible = loco.feasible.ci.Bxy[1],
                                  Bzy.ci.upper.loco.feasible = loco.feasible.ci.Bzy[2],
                                  Bzy.ci.lower.loco.feasible = loco.feasible.ci.Bzy[1]))

    
    ## df.double.proxy <- copy(df)
    ## df.double.proxy <- df.double.proxy[,y.obs:=NA]
    ## df.double.proxy <- df.double.proxy[,y:=NA]
    
    ## double.proxy.mle <- measerr_irr_mle_dv(df.double.proxy, outcome_formula=y~x+z, outcome_family=binomial(link='logit'), coder_formulas=list(y.obs.0 ~ y),  proxy_formula=w_pred ~ y.obs.0 + y, proxy_family=binomial(link='logit'))
    ## print(double.proxy.mle$hessian)
    ## fisher.info <- solve(double.proxy.mle$hessian)
    ## coef <- double.proxy.mle$par
    ## ci.upper <- coef + sqrt(diag(fisher.info)) * 1.96
    ## ci.lower <- coef - sqrt(diag(fisher.info)) * 1.96

    ## result <- append(result, list(Bxy.est.double.proxy=coef['x'],
    ##                               Bzy.est.double.proxy=coef['z'],
    ##                               Bxy.ci.upper.double.proxy = ci.upper['x'],
    ##                               Bxy.ci.lower.double.proxy = ci.lower['x'],
    ##                               Bzy.ci.upper.double.proxy = ci.upper['z'],
    ##                               Bzy.ci.lower.double.proxy = ci.lower['z']))


    df.triple.proxy <- copy(df)
    df.triple.proxy <- df.triple.proxy[,y.obs:=NA]
    df.triple.proxy <- df.triple.proxy[,y:=NA]
    
    triple.proxy.mle <- measerr_irr_mle_dv(df.triple.proxy, outcome_formula=outcome_formula, outcome_family=binomial(link='logit'), coder_formulas=coder_formulas, proxy_formula=proxy_formula, proxy_family=binomial(link='logit'))
    print(triple.proxy.mle$hessian)
    fisher.info <- solve(triple.proxy.mle$hessian)
    print(fisher.info)
    coef <- triple.proxy.mle$par
    ci.upper <- coef + sqrt(diag(fisher.info)) * 1.96
    ci.lower <- coef - sqrt(diag(fisher.info)) * 1.96

    result <- append(result, list(Bxy.est.triple.proxy=coef['x'],
                                  Bzy.est.triple.proxy=coef['z'],
                                  Bxy.ci.upper.triple.proxy = ci.upper['x'],
                                  Bxy.ci.lower.triple.proxy = ci.lower['x'],
                                  Bzy.ci.upper.triple.proxy = ci.upper['z'],
                                  Bzy.ci.lower.triple.proxy = ci.lower['z']))

    ## df.loco.mle <- copy(df)
    ## df.loco.mle[,y.obs:=NA]
    ## df.loco.mle[(y.obs.0)==(y.obs.1),y.obs:=y.obs.0]
    ## df.loco.mle[,y.true:=y]
    ## df.loco.mle[,y:=y.obs]
    ## print(df.loco.mle[!is.na(y.obs.1),mean(y.true==y,na.rm=TRUE)])
    ## loco.mle <- measerr_mle_dv(df.loco.mle, outcome_formula=outcome_formula, proxy_formula=proxy_formula)
    ## fisher.info <- solve(loco.mle$hessian)
    ## coef <- loco.mle$par
    ## ci.upper <- coef + sqrt(diag(fisher.info)) * 1.96
    ## ci.lower <- coef - sqrt(diag(fisher.info)) * 1.96

    ## result <- append(result, list(Bxy.est.loco.mle=coef['x'],
    ##                               Bzy.est.loco.mle=coef['z'],
    ##                               Bxy.ci.upper.loco.mle = ci.upper['x'],
    ##                               Bxy.ci.lower.loco.mle = ci.lower['x'],
    ##                               Bzy.ci.upper.loco.mle = ci.upper['z'],
    ##                               Bzy.ci.lower.loco.mle = ci.lower['z']))



    ## my implementatoin of liklihood based correction
    mod.zhang <- zhang.mle.dv(df.loco.mle)
    coef <- coef(mod.zhang)
    ci <- confint(mod.zhang,method='quad')

    result <- append(result,
                     list(Bxy.est.zhang = coef['Bxy'],
                          Bxy.ci.upper.zhang = ci['Bxy','97.5 %'],
                          Bxy.ci.lower.zhang = ci['Bxy','2.5 %'],
                          Bzy.est.zhang = coef['Bzy'],
                          Bzy.ci.upper.zhang = ci['Bzy','97.5 %'],
                          Bzy.ci.lower.zhang = ci['Bzy','2.5 %']))

    

    print(df.loco.mle)

    # amelia says use normal distribution for binary variables.
    tryCatch({
        amelia.out.k <- amelia(df.loco.mle, m=200, p2s=0, idvars=c('y','ystar','w','y.obs.1','y.obs.0','y.true'))
        mod.amelia.k <- zelig(y.obs~x+z, model='ls', data=amelia.out.k$imputations, cite=FALSE)
        (coefse <- combine_coef_se(mod.amelia.k, messages=FALSE))
        est.x.mi <- coefse['x','Estimate']
        est.x.se <- coefse['x','Std.Error']
        result <- append(result,
                         list(Bxy.est.amelia.full = est.x.mi,
                              Bxy.ci.upper.amelia.full = est.x.mi + 1.96 * est.x.se,
                              Bxy.ci.lower.amelia.full = est.x.mi - 1.96 * est.x.se
                              ))

        est.z.mi <- coefse['z','Estimate']
        est.z.se <- coefse['z','Std.Error']

        result <- append(result,
                         list(Bzy.est.amelia.full = est.z.mi,
                              Bzy.ci.upper.amelia.full = est.z.mi + 1.96 * est.z.se,
                              Bzy.ci.lower.amelia.full = est.z.mi - 1.96 * est.z.se
                              ))

    },
    error = function(e){
        message("An error occurred:\n",e)
        result$error <- paste0(result$error,'\n', e)
    })
    ## mle.irr <- measerr_irr_mle( df, outcome_formula = outcome_formula, rater_formula = rater_formula, proxy_formula=proxy_formula, truth_formula=truth_formula)

    ## fisher.info <- solve(mle.irr$hessian)
    ## coef <- mle.irr$par
    ## ci.upper <- coef + sqrt(diag(fisher.info)) * 1.96
    ## ci.lower <- coef - sqrt(diag(fisher.info)) * 1.96
    
    ## result <- append(result,
    ##                  list(Bxy.est.mle = coef['x'],
    ##                       Bxy.ci.upper.mle = ci.upper['x'],
    ##                       Bxy.ci.lower.mle = ci.lower['x'],
    ##                       Bzy.est.mle = coef['z'],
    ##                       Bzy.ci.upper.mle = ci.upper['z'],
    ##                       Bzy.ci.lower.mle = ci.lower['z']))

    return(result)

}
