#!/usr/bin/env Rscript
library(argparser)
library(mecor)
library(ggplot2)
library(data.table)
library(filelock)
library(arrow)
library(Amelia)
library(Zelig)
library(predictionError)
options(amelia.parallel="no",
        amelia.ncpus=1)
setDTthreads(40)

source("simulation_base.R")

simulate_data <- function(N,
                          m,
                          B0,
                          Bxy,
                          Bzx,
                          Bzy,
                          seed,
                          y_explained_variance = 0.025,
                          prediction_accuracy = 0.73,
                          y_bias = -0.8,
                          z_bias = 0,
                          x_bias = 0,
                          Px = 0.5,
                          Pz = 0.5,
                          accuracy_imbalance_difference = 0.3) {
  set.seed(seed)
  # make w and y dependent
  z <- rbinom(N, 1, plogis(qlogis(Pz)))
  x <- rbinom(N, 1, plogis(Bzx * z + qlogis(Px)))

  y.var.epsilon <- (var(Bzy * z) + var(Bxy * x) + 2 * cov(Bzy * z, Bxy * x)) * ((1 - y_explained_variance) / y_explained_variance)
  y.epsilon <- rnorm(N, sd = sqrt(y.var.epsilon))
  y <- Bzy * z + Bxy * x + y.epsilon

  df <- data.table(x = x, y = y, z = z)

  if (m < N) {
    df <- df[sample(nrow(df), m), ":="(x.obs = x, z.obs = z)]
  } else {
    df <- df[, ":="(x.obs = x, z.obs = z)]
  }

  resids <- resid(lm(y ~ x + z))
  odds.x1 <- qlogis(prediction_accuracy) + y_bias * qlogis(pnorm(resids[x == 1])) + z_bias * qlogis(pnorm(z[x == 1], sd(z)))
  odds.x0 <- qlogis(prediction_accuracy, lower.tail = F) + y_bias * qlogis(pnorm(resids[x == 0])) + z_bias * qlogis(pnorm(z[x == 0], sd(z)))

  ## acc.x0 <- p.correct[df[,x==0]]
  ## acc.x1 <- p.correct[df[,x==1]]

  df[x == 0, w := plogis(rlogis(.N, odds.x0))]
  df[x == 1, w := plogis(rlogis(.N, odds.x1))]
  df[, w_pred := as.integer(w > 0.5)]

  ## now let's use another classifier for z
  odds.z1 <- qlogis(prediction_accuracy) + y_bias * qlogis(pnorm(resids[z == 1])) + x_bias * qlogis(pnorm(x[z == 1], sd(x)))
  odds.z0 <- qlogis(prediction_accuracy, lower.tail = F) + y_bias * qlogis(pnorm(resids[z == 0])) + x_bias * qlogis(pnorm(x[z == 0], sd(z)))

  df[z == 0, a := plogis(rlogis(.N, odds.z0))]
  df[z == 1, a := plogis(rlogis(.N, odds.z1))]
  df[, a_pred := as.integer(a > 0.5)]


  print(mean(df$w_pred == df$x))
  print(mean(df$w_pred == df$x))
  print(mean(df[y >= 0]$w_pred == df[y >= 0]$x))
  print(mean(df[y <= 0]$w_pred == df[y <= 0]$x))
  return(df)
}

parser <- arg_parser("Simulate data and fit corrected models with two independent variables that are classified")
parser <- add_argument(parser, "--N", default=5000, help="number of observations of w")
parser <- add_argument(parser, "--m", default=500, help="m the number of ground truth observations")
parser <- add_argument(parser, "--seed", default=51, help='seed for the rng')
parser <- add_argument(parser, "--outfile", help='output file', default='example_2.feather')
parser <- add_argument(parser, "--y_explained_variance", help='what proportion of the variance of y can be explained?', default=0.1)
parser <- add_argument(parser, "--prediction_accuracy", help='how accurate is the predictive model?', default=0.75)
parser <- add_argument(parser, "--accuracy_imbalance_difference", help='how much more accurate is the predictive model for one class than the other?', default=0.3)
parser <- add_argument(parser, "--Bzx", help='Effect of z on x', default=0.3)
parser <- add_argument(parser, "--Bzy", help='Effect of z on y', default=-0.3)
parser <- add_argument(parser, "--Bxy", help='Effect of z on y', default=0.3)
parser <- add_argument(parser, "--outcome_formula", help='formula for the outcome variable', default="y~x+z")
parser <- add_argument(parser, "--proxy_formula", help='formula for the proxy variable', default="w_pred~y*z*x")
parser <- add_argument(parser, "--y_bias", help='coefficient of y on the probability a classification is correct', default=-0.5)
parser <- add_argument(parser, "--z_bias", help='coefficient of z on the probability a classification is correct', default=0)
parser <- add_argument(parser, "--x_bias", help='coefficient of x on the probability a classification is correct', default=0)
parser <- add_argument(parser, "--truth_formula", help='formula for the true variable', default="x~z")
parser <- add_argument(parser, "--Px", help='base rate of x', default=0.5)
parser <- add_argument(parser, "--confint_method", help='method for approximating confidence intervals', default='quad')
args <- parse_args(parser)

B0 <- 0
Px <- args$Px
Bxy <- args$Bxy
Bzy <- args$Bzy
Bzx <- args$Bzx

if(args$m < args$N){

    df <- simulate_data(args$N, args$m, B0, Bxy, Bzx, Bzy, args$seed, args$y_explained_variance, args$prediction_accuracy, y_bias=args$y_bias)

    ## df.pc <- df[,.(x,y,z,w_pred,w)]
    ##                                     #    df.pc <- df.pc[,err:=x-w_pred]
    ## pc.df <- pc(suffStat=list(C=cor(df.pc),n=nrow(df.pc)),indepTest=gaussCItest,labels=names(df.pc),alpha=0.05)
    ## plot(pc.df)

    result <- list('N'=args$N,'m'=args$m,'B0'=B0,'Bxy'=Bxy, 'Bzx'=args$Bzx, 'Bzy'=Bzy, 'Px'=Px, 'seed'=args$seed, 'y_explained_variance'=args$y_explained_variance, 'prediction_accuracy'=args$prediction_accuracy, 'accuracy_imbalance_difference'=args$accuracy_imbalance_difference, 'y_bias'=args$y_bias,'outcome_formula'=args$outcome_formula, 'proxy_formula'=args$proxy_formula,truth_formula=args$truth_formula, confint_method=args$confint_method, error='')

  research_data <- df[is.na(x.obs), .(w = w_pred, a = a_pred, y = y)] |> as.data.frame()

  val_data <- df[!is.na(x.obs), .(w = w_pred, a = a_pred, z = z.obs, y = y, x = x.obs)] |> as.data.frame()

  outline <- run_simulation(y ~ x || w + z ||
    a, w ~ x + z + y, a ~ x + z + y, data=research_data,data2=val_data, result = result)
    
   
 outfile_lock <- lock(paste0(args$outfile, '_lock'),exclusive=TRUE)
    if(file.exists(args$outfile)){
        logdata <- read_feather(args$outfile)
        logdata <- rbind(logdata,as.data.table(outline), fill=TRUE)
    } else {
        logdata <- as.data.table(outline)
    }

    print(outline)
    write_feather(logdata, args$outfile)
    unlock(outfile_lock)
}
