# LC Instrumental Variable Probit Model
# Only one continuous endogenous variable
# No panel data
# By               : Mauricio Sarrias
# Created          : February 8, 2019
# Last modification: March 24, 2020



LcIv <- function(formula, 
                 data,
                 subset, 
                 na.action,
                 method = "bfgs", 
                 model  =  c("lc", "lciv"),
                 estimation = c("fiml", "twostep"), 
                 Q = 2, 
                 gradient = TRUE,
                 init.value = c("lc", "ivp"),
                 firststep = c("ols", "lc"), 
                 print.init = FALSE, 
                 messages  = TRUE,
                 get.init =  FALSE, 
                 start =  NULL, 
                 ...){
  # Match call
  require("Formula", quietly = TRUE)
  require("Rchoice", quietly = TRUE)
  require("maxLik", quietly = TRUE)
  callT <- match.call(expand.dots = TRUE)
  nframe <- length(sys.calls())
  model <- match.arg(model)
  estimation <- match.arg(estimation)
  firststep <- match.arg(firststep)
  
  # some checks
  if (Q < 2) stop("Classes cannot be lower than 2")
  
  # model frame
  mf <- match.call(expand.dots = FALSE)
  m <- match(c("formula", "data", "subset", "na.action"), names(mf), 0)
  mf <- mf[c(1, m)]
  F1 <- Formula(formula)
  mf[[1]] <- as.name("model.frame")
  mf$formula <- F1
  mf <- eval(mf, parent.frame())
  
  # Extract variables for each part
  if (model == "lciv") {
    y1 <- model.part(F1, data = mf, lhs = 1)[, 1]
    y2 <- model.part(F1, data = mf, lhs = 1)[, 2]
    X  <- model.matrix(F1, data = mf, rhs = 1)
    Z  <- model.matrix(F1, data = mf, rhs = 2)
    
  } else {
    y1 <- model.response(mf)
    X <- model.matrix(F1, mf)
  }
  H  <- model.matrix(F1, data = mf, rhs = 3)
  
  # stop if IV and model is underidentified
  if (model == "lciv" && (ncol(Z) < ncol(X))) stop("model underidentified")
  if (messages) {
    if (model == "lciv" && (ncol(Z) == ncol(X))) cat("Note: model just identified", fill = TRUE)
    if (model == "lciv" && (ncol(Z) > ncol(X)))  cat("Note: overindentified model", fill = TRUE)
  }
  ## Create matrix H------------------
  indata <- mf
  cldata <- indata[rep(seq_len(nrow(indata)), each = Q), ] #expand H
  class   <- factor(rep(1:Q, nrow(X)))
  N <- length(y1)  
  cldata   <- cbind(cldata, class, row.names = NULL)
  class.var <- formula(F1, rhs = 3, lhs = 0) 
  has.int  <- attr(terms(class.var), "intercept") == 1L
  if (has.int) intercept.char <- "factor(class)" else intercept.char <- NULL 
  if (!has.int) {
    class.var <- update(class.var, ~. + 1)
    class.var <- update(class.var, ~.)
    class.var.char <- as.character(class.var)[2]
    has.xclass <- as.character(class.var)[2]
    class.var.var <- colnames(model.matrix(update(class.var, ~. + 1), cldata))[-1]
    class.var.char <- paste("(", class.var.var, "):class", sep = "")
    if (class.var.char == "1") class.var.char <- class.var.var <- NULL
  } else {
    has.xclass <- as.character(class.var)[2]
    if (has.xclass == "1") {
      class.var.char <- NULL
    } else {
      class.var.var <- colnames(model.matrix(update(class.var, ~. + 1), cldata))[-1]
      class.var.char <- paste("(", class.var.var, "):class", sep = "")
    } 
  }
  resp.name <- as.character(attr(F1, "lhs"))
  form.char <- paste(c(intercept.char, class.var.char), collapse = "+")
  nformula <- as.formula(paste(resp.name, " ~ ", form.char))
  H <- model.matrix(nformula, cldata)[, -1, drop = F]
  
  lev1 <- levels(class)[1]
  lev1 <- paste("class", lev1, sep = "")
  if (has.xclass != "1") {
    toremove <- unlist(lapply(as.list(class.var.var), function(x) paste(lev1, x, sep = ":")))
    revtoremove <- unlist(lapply(as.list(class.var.var), function(x) paste(x, lev1, sep = ":")))
    toremove <- colnames(H) %in% c(toremove, revtoremove)
    H <- H[, !toremove, drop = FALSE]
  } 
  namesH <- colnames(H)
  for (i in 1:length(namesH)) namesH[i] <- sub('factor', '', namesH[i])
  colnames(H) <- namesH
  
  Hl <- vector(length = Q, mode = "list")
  names(Hl) <- levels(class)
  for (i in levels(class)) {
    Hl[[i]] <- H[class == i, , drop = FALSE]
  }
  
  # Initial values:-----
  if (is.null(start)){
    s_class <- rep(0, ncol(H))
    names(s_class) <- colnames(H)
    if (model == "lc") {
      fp         <- formula(F1, rhs = 1, lhs = 1)
      probit     <- Rchoice(fp, data = mf, family = binomial('probit'))
      beta    <- coef(probit)
      init.shift <- seq(-0.05, 0.05, length.out = Q)
      lc.beta    <- c()
      lc.nbeta   <- c()
      for (i in 1:Q) {
        lc.beta  <- c(lc.beta,  beta + init.shift[i])
        lc.nbeta <- c(lc.nbeta , paste('class', i, names(beta) , sep = '.'))
      }
      names(lc.beta)  <- lc.nbeta
      start <- c(lc.beta, s_class)
    } else {
      if (estimation == "fiml") {
        # Two options IVP and LC
        if (init.value == "ivp") {
          if(messages) cat("Estimating an IV probit for initial values", fill = TRUE)
          fiv     <- formula(F1, rhs = 1:2, lhs = 1)
          iv      <- iv_probit(fiv, data = mf)
          beta    <- coef(iv)[1:ncol(X)]
          delta   <- coef(iv)[(ncol(X) + 1):(ncol(X) + ncol(Z))]
          lnsigma <- tail(coef(iv), n = 2L)[1]
          athrho  <- tail(coef(iv), n = 1L)
          init.shift <- seq(-0.02, 0.02, length.out = Q)
          lc.beta    <- c()
          lc.delta   <- c()
          lc.lnsigma <- c()
          lc.rho     <- c()
          lc.nbeta <- c()
          lc.ndelta <- c()
          lc.nlnsigma <- c()
          lc.nrho <- c()
          for (i in 1:Q) {
            lc.beta    <- c(lc.beta,  beta + init.shift[i])
            lc.delta   <- c(lc.delta, delta + init.shift[i])
            lc.lnsigma <- c(lc.lnsigma, lnsigma + init.shift[i])
            lc.rho <- c(lc.rho, athrho + init.shift[i])
            lc.nbeta    <- c(lc.nbeta , paste('class', i, names(beta) , sep = '.'))
            lc.ndelta   <- c(lc.ndelta, paste('class', i, names(delta), sep = '.'))
            lc.nlnsigma <- c(lc.nlnsigma, paste('class', i, "lnsigma", sep = '.'))
            lc.nrho     <- c(lc.nrho, paste('class', i, "athro", sep = '.'))
          }
          names(lc.beta)    <- lc.nbeta
          names(lc.delta)   <- lc.ndelta
          names(lc.lnsigma) <- lc.nlnsigma
          names(lc.rho) <- lc.nrho
          start <- c(lc.beta, lc.delta, s_class, lc.lnsigma, lc.rho)
        } else {
          # Obtain intial values from lc model
          if(messages) cat("Estimating a LC on the first equation for initial values", fill = TRUE)
          probit     <- glm(y1 ~ X - 1, family = binomial('probit'))
          beta       <- coef(probit)
          init.shift <- seq(-0.05, 0.05, length.out = Q)
          lc.beta    <- c()
          lc.nbeta   <- c()
          for (i in 1:Q) {
            lc.beta  <- c(lc.beta,  beta + init.shift[i])
            lc.nbeta <- c(lc.nbeta , paste('class', i, colnames(X) , sep = '.'))
          }
          names(lc.beta)  <- lc.nbeta
          start <- c(lc.beta, s_class)
          lcp <- maxLik(ml_bin_lc, start = start, y = y1, X = X, H = Hl, Q = Q, method = method, gradient = TRUE)
          print(summary(lcp))
          beta.lc <- coef(lcp)[1:(ncol(X)*Q)]
          s_class <- coef(lcp)[-(1:(ncol(X)*Q))]
          if (messages) cat("Estimating a LC on the second equation for initial values", fill = TRUE)
          ols          <- lm(y2 ~ Z - 1, data = mf)
          sigma2   <- sum(resid(ols) ^ 2 ) / ols$df.residual
          names(sigma2) <- c("sigma2")
          beta         <- coef(ols)
          names(beta)  <- colnames(Z)
          init.shift <- seq(-0.05, 0.05, length.out = Q)
          lc.beta    <- c()
          lc.sigma   <- c()
          lc.nbeta   <- c()
          lc.athanr  <- c()
          rho <- 0 
          for (i in 1:Q) {
            lc.sigma <- c(lc.sigma,  sigma2 + init.shift[i])
            lc.beta  <- c(lc.beta,  beta + init.shift[i])
            lc.athanr  <- c(lc.athanr,  rho + init.shift[i])
            lc.nbeta <- c(lc.nbeta , paste('class', i, names(beta) , sep = '.'))
          }
          names(lc.sigma) <- paste('class', 1:Q , 'lnsigma', sep = '.')
          names(lc.beta)  <- lc.nbeta
          start <- c(lc.sigma, lc.beta, s_class)
          lin_clas <- maxLik(ml_lin_lc, start = start, y = y2, X = Z, H = Hl, Q = Q, method = method)
          print(summary(lin_clas))
          #Obtain residuals
          sigma2 <- coef(lin_clas)[1:Q]
          delta.lc <-  coef(lin_clas)[-c(1:Q)]
          names(lc.athanr) <- paste('class', 1:Q , 'athro', sep = '.')
          start <- c(beta.lc, delta.lc, log(sqrt(sigma2)), lc.athanr)
        }
      } else {
        if (firststep == "lc") {
          ols          <- lm(y2 ~ Z - 1, data = mf)
          sigma2       <- sigma(ols) ^ 2
          names(sigma2) <- c("sigma2")
          beta         <- coef(ols)
          names(beta)  <- colnames(Z)
          init.shift <- seq(-0.05, 0.05, length.out = Q)
          lc.beta    <- c()
          lc.sigma   <- c()
          lc.nbeta   <- c()
          for (i in 1:Q) {
            lc.sigma <- c(lc.sigma,  sigma2 + init.shift[i])
            lc.beta  <- c(lc.beta,  beta + init.shift[i])
            lc.nbeta <- c(lc.nbeta , paste('class', i, names(beta) , sep = '.'))
          }
          names(lc.sigma) <- paste('class.sigma', 1:Q , sep = '.')
          names(lc.beta)  <- lc.nbeta
          start <- c(lc.sigma, lc.beta, s_class)
          if (messages) cat("Estimating a LC IV first step", "\n")
          lin_clas <- maxLik(ml_lin_lc, start = start, y = y2, X = Z, H = Hl, Q = Q, method = method)
          print(summary(lin_clas))
          delta_hat <- coef(lin_clas)
          #Obtain residuals
          sigma2 <- delta_hat[1:Q]
          betas <-  delta_hat[-c(1:Q)]
          beta <- matrix(betas[1L:(nrow(Z) * Q)], nrow = ncol(Z), ncol = Q)
          rownames(beta) <- colnames(Z); colnames(beta) <- paste("class", 1:Q, sep = ":")
          gamma <- betas[-c(1L:(ncol(Z) * Q))]
          # Make weights from a multinomial logit model
          ew <- lapply(Hl, function(x) exp(crossprod(t(x), gamma)))
          sew <- suml(ew)
          Wnq <- lapply(ew, function(x){ v <- x / sew;
          v[is.na(v)] <- 0;
          as.vector(v)})
          Wnq <- Reduce(cbind, Wnq) # W
          # Make Probit Probability
          #y2_hat <- rowSums()
          #res <- (y2 - y2_hat) / repRows(sqrt(sigma2), nrow(Z))
          res <- rowSums((y2 - Wnq * tcrossprod(Z, t(beta)))/ repRows(sqrt(sigma2), n = nrow(Z)))
          Xaux <- cbind(X, res)
          probit     <- glm(y1 ~ Xaux - 1, family = binomial(link = 'probit'))
          beta    <- coef(probit)
          names(beta) <- c(colnames(X), "res")
          init.shift <- seq(-0.05, 0.05, length.out = Q)
          lc.beta    <- c()
          lc.nbeta   <- c()
          for (i in 1:Q) {
            lc.beta  <- c(lc.beta,  beta + init.shift[i])
            lc.nbeta <- c(lc.nbeta , paste('class', i, names(beta) , sep = '.'))
          }
          names(lc.beta)  <- lc.nbeta
          colnames(Xaux) <- c(colnames(X), "res")
          start <- c(lc.beta, s_class)
        } else {
          ols          <- lm(y2 ~ Z - 1, data = mf)
          sigma       <- sigma(ols)
          res <- ols$residuals / sigma
          Xaux <- cbind(X, res)
          probit     <- glm(y1 ~ Xaux - 1, family = binomial(link = 'probit'))
          beta    <- coef(probit)
          names(beta) <- c(colnames(X), "res")
          init.shift <- seq(-0.05, 0.05, length.out = Q)
          lc.beta    <- c()
          lc.nbeta   <- c()
          for (i in 1:Q) {
            lc.beta  <- c(lc.beta,  beta + init.shift[i])
            lc.nbeta <- c(lc.nbeta , paste('class', i, names(beta) , sep = '.'))
          }
          names(lc.beta)  <- lc.nbeta
          colnames(Xaux) <- c(colnames(X), "res")
          start <- c(lc.beta, s_class)
        }
      }
    }
  } else {
    start <- start 
  }
  
  
  if (print.init) {
    cat("Initial values:", "\n")
    print(start)
  }
  if (get.init) return(start)
  
  # Estimate the model ----
  opt <- callT
  opt$start <- start
  opt$gradient <- as.name('gradient')
  m <- match(c('method', 'print.level', 'iterlim',
               'start','tol', 'ftol', 'steptol', 'fixed', 'constraints', 
               'control'),
             names(opt), 0L)
  opt <- opt[c(1L, m)]
  opt[[1]] <- as.name("maxLik")
  if (model  == "lc") {
    if(messages) cat("Estimating a LC Probit model", "\n")
    opt$logLik <-  as.name("ml_bin_lc")
    opt[c("y", "X", "H", "Q")] <- list(as.name("y1"), 
                                       as.name("X"), 
                                       as.name("Hl"),
                                       as.name("Q"))
  } else {
    if (estimation == "fiml") {
      if(messages) cat("Estimating an IVLC-Probit model", "\n")
      opt$logLik <- as.name("ml_lciv")
      opt[c("y1", "y2", "X", "Z", "W", "Q")] <- list(as.name("y1"), 
                                                     as.name("y2"), 
                                                     as.name("X"), 
                                                     as.name("Z"), 
                                                     as.name("Hl"), 
                                                     as.name("Q")) 
    } else {
      if(messages) cat("Estimating an LC IV secon step", "\n")
      opt$logLik <-  as.name("ml_bin_lc")
      opt[c("y", "X", "H", "Q")] <- list(as.name("y1"), 
                                         as.name("Xaux"), 
                                         as.name("Hl"),
                                         as.name("Q"))
    }
    
  }
  out <- eval(opt, sys.frame(which = nframe))
  
  # Obtain post estimates
  if (model == "lciv") {
    opt[[1]] <- as.name('ml_lciv')
    names(opt)[[2]] <- 'theta'
    opt[[2]] <- coef(out)
    opt$gradient <- FALSE
    opt$get.post <- TRUE
    opt$fixed <- opt$steptol <- opt$iterlim <- opt$method <- opt$print.level <- opt$tol <- opt$ftol <- opt$logLik <- opt$start <- NULL
    opt$constraints <- NULL


    again  <- eval(opt, sys.frame(which = nframe))
    out$P1 <- attr(again, 'P1') 
    out$index1 <- attr(again, 'index1') 
    out$index2 <- attr(again, 'index2') 
    out$Piq    <- attr(again, 'Piq')    
    out$Pi     <- attr(again, 'Pi')     
    out$Wiq    <- attr(again, 'Wiq')    
    out$ai     <- attr(again, 'ai') 
    out$y1     <- y1
    out$y2     <- y2
    out$X      <- X
    out$Z      <- Z
    out$W      <- Hl
  }
  
  
  if (model == "lc") {
    opt[[1]] <- as.name('ml_bin_lc')
    names(opt)[[2]] <- 'theta'
    opt[[2]] <- coef(out)
    opt$gradient <- FALSE
    opt$get.post <- TRUE
    opt$fixed <- opt$steptol <- opt$iterlim <- opt$method <- opt$print.level <- opt$tol <- opt$ftol <- opt$logLik <- opt$start <- NULL
    opt$constraints <- NULL
    
    
    again  <- eval(opt, sys.frame(which = nframe))
    out$Piq    <- attr(again, 'Piq')    
    out$Pi     <- attr(again, 'Pi')     
    out$Wiq    <- attr(again, 'Wiq')    
  }
  
  return(out)
}  



# Maximum Likelihood
ml_lciv <- function(theta, y1, y2, X, Z, W, Q, gradient = TRUE, get.post = FALSE){
  #param: bq, dq, lq, lnsigmaq, athrhoq
  K <- ncol(X)
  P <- ncol(Z)
  J <- ncol(W[[1]])
  N <- nrow(X)
  
  beta    <- matrix(theta[1L:(K * Q)], nrow = K, ncol = Q)                        # Matrix of K*Q
  delta   <- matrix(theta[((K * Q) + 1):((P * Q) + (K * Q))], nrow = P, ncol = Q) # Matrix of P*Q
  lambda  <- theta[((P * Q) + (K * Q) + 1):(J + (P * Q) + (K * Q))]     # vector of (Q - 1) * J
  lnsigma <- repRows(tail(theta, n = 2 * Q)[1:Q], n = nrow(X))
  #print(tail(theta, n = 2 * Q)[1:Q])
  athrho  <- repRows(tail(theta, n = Q), n = nrow(X))
  sigma   <- exp(lnsigma)
  #rho     <- tanh(athrho)
  rho     <- (exp(2 * athrho) - 1) / (exp(2 * athrho) + 1)
  #if ((rho < -1) || (rho > 1)) return(NA) #Fixme
  
  rownames(beta)  <- colnames(X); colnames(beta)  <- paste("class", 1:Q, sep = ":")
  rownames(delta) <- colnames(Z); colnames(delta) <- paste("class", 1:Q, sep = ":")
  
  F  <- pnorm
  f  <- dnorm
  ff <- function(x) -x * dnorm(x)
  
  
  index1 <- tcrossprod(X, t(beta))    #N * Q
  #index1  <- pmin(index1, 100)
  index2 <- tcrossprod(Z, t(delta))   #N * Q
  #index2 <- pmin(index2, 100)
  
  y2 <- repCols(y2, n = Q)
  q  <- repCols(2 * y1 - 1, n = Q)
  ai <- q * (index1 + (rho / sigma) * (y2 - index2)) / sqrt(1 - rho ^ 2)
  #ai <- pmin(ai, 100)
  #ai <- pmax(ai, -100)
  bi <- (y2 - index2) / sigma
  #bi <- pmin(bi, 100)
  #bi <- pmax(bi, -100)
  
  # Make weights from a multinomial logit model
  ew <- lapply(W, function(x) exp(crossprod(t(x), lambda)))
  sew <- suml(ew)
  Wiq <- lapply(ew, function(x){ v <- x / sew;
  v[is.na(v)] <- 0;
  as.vector(v)})
  Wiq <- Reduce(cbind, Wiq) 
  
  P1  <- pmax(F(ai), .Machine$double.eps)
  P1[is.nan(P1)] <- 0
  P2  <- (1 / sigma) * f(bi)
  P2[is.nan(P2)] <- 0
  Piq <- Wiq * P1 * P2 
  Pi <- pmax(apply(Piq, 1, sum), .Machine$double.eps)
  Pi[is.na(Pi)] <- 0
  LL <- log(Pi)
  # print(LL)
  
  ##Gradient
  if (gradient) {
    mila <- f(ai) / P1 # This is critial, otherwise Nan
    milb <- ff(bi) / f(bi)
    #milb <- pmax(ff(bi) / f(bi), .Machine$double.eps)
    Qiq <-  (Wiq * P1 * P2 / Pi)
    dp_b <- mila * q  / sqrt(1 - rho ^ 2)
    eta <-  Qiq * dp_b  
    etar <- eta[, rep(1:Q, each = K)]
    Xg <- X[, rep(1:K, Q)] # N * (K * Q)
    grad.beta <- Xg * etar
    
    dp_d <- -(mila * q * (rho / sigma) / (sqrt(1 - rho ^ 2)) + milb * (1 / sigma))
    eta <- Qiq * dp_d
    etar <- eta[, rep(1:Q, each = P)]
    Zg <- Z[, rep(1:P, Q)] # N * (P * Q)
    grad.delta <- Zg * etar
    
    dp_ls <- -(mila * q * rho / sqrt(1 - rho ^ 2) + milb) * ((y2 - index2) * exp(-lnsigma)) - 1
    grad.lnsigma <- Qiq * dp_ls
    
    a <- index1; b <- sigma; c <- (y2 - index2)
    dp_t <- mila * q * (exp(-athrho) * ((c + a*b) * exp(2 * athrho)  + c - a*b)) / (2*b)
    grad.athrho <- Qiq * dp_t
    
    Wg <- vector(mode = "list", length = Q)
    IQ <- diag(Q)
    for (q in 1:Q) Wg[[q]] <- rowSums(Qiq * (repRows(IQ[q, ], N) - repCols(Wiq[, q], Q)))
    grad.lambda <- suml(mapply("*", W, Wg, SIMPLIFY = FALSE))
    grad <- cbind(grad.beta, grad.delta, grad.lambda, grad.lnsigma, grad.athrho)
    colnames(grad) <- names(theta)
    attr(LL,'gradient') <- grad
  }
  if (get.post) {
    attr(LL, 'P1')     <- P1
    attr(LL, 'index1') <- index1
    attr(LL, 'index2') <- index2
    attr(LL, 'Piq')    <- Piq
    attr(LL, 'Pi')     <- Pi
    attr(LL, 'Wiq')    <- Wiq
    attr(LL, 'ai')     <- (index1 + (rho / sigma) * (y2 - index2)) / sqrt(1 - rho ^ 2)
  }
  LL
}


ml_bin_lc <- function(theta, y, X, H, Q, id = NULL, weights = NULL,
                      gradient = TRUE, get.post = FALSE){
  # Parameters and globals
  K <- ncol(X)
  N <- nrow(X)
  panel <- !is.null(id)
  if (is.null(weights)) weights <- 1
  if (panel) if (length(weights == 1)) weights <- rep(weights, N)
  
  beta <- matrix(theta[1L:(K * Q)], nrow = K, ncol = Q)
  rownames(beta) <- colnames(X); colnames(beta) <- paste("class", 1:Q, sep = ":")
  gamma <- theta[-c(1L:(K * Q))]
  
  if (get.post) bi <- t(beta)
  
  # Make weights from a multinomial logit model
  ew <- lapply(H, function(x) exp(crossprod(t(x), gamma)))
  sew <- suml(ew)
  Wnq <- lapply(ew, function(x){ v <- x / sew;
  v[is.na(v)] <- 0;
  as.vector(v)})
  Wnq <- Reduce(cbind, Wnq) # W
  
  # Make Probit Probability
  pfun <- pnorm
  dfun <- dnorm
  mill <- function(x) dfun(x) / pmax(pfun(x), .Machine$double.eps)
  q <- 2 * y - 1
  
  index <- tcrossprod(X, t(beta))
  Pnq <- pfun(q * index) # N * Q
  if (panel) Pnq <- apply(Pnq, 2, tapply, id, prod)
  WPnq <- Wnq * Pnq 
  Ln <- apply(WPnq, 1, sum)
  Ln  <- pmax(Ln, .Machine$double.eps)
  Ln[is.nan(Ln)] <- 0
  if (get.post) Qnr <- WPnq / Ln
  lnL <- if (panel) sum(log(Ln) * weights[!duplicated(id)]) else sum(log(Ln) * weights)
  
  ## Gradient
  if (gradient) {
    lambda <- q * mill(q * index) # N * Q
    Qnr <- WPnq / Ln
    if (panel) Qnr <- Qnr[id, ]
    eta <- lambda * Qnr
    etar <- eta[, rep(1:Q, each = K)]
    Xg <- X[, rep(1:K, Q)] # N * (K * Q)
    grad.beta <- Xg * etar
    
    if (panel) {
      Wnq <- Wnq[id, ]
      H <- lapply(H, function(x) x[id, ])
    }
    Wg <- vector(mode = "list", length = Q)
    IQ <- diag(Q)
    for(q in 1:Q) Wg[[q]] <- rowSums(Qnr * (repRows(IQ[q, ], N) - repCols(Wnq[, q], Q)))
    grad.gamma <- suml(mapply("*", H, Wg, SIMPLIFY = FALSE))
    gari <- cbind(grad.beta, grad.gamma)
    colnames(gari) <- names(theta)
    attr(lnL, "gradient") <- gari * weights
  }
  if (get.post) {
    if (get.post) {
      attr(lnL, 'Piq')    <- Pnq
      attr(lnL, 'Pi')     <- Ln
      attr(lnL, 'Wiq')    <- Wnq
    }
  }
  lnL
}




ml_lin_lc <- function(theta, y, X, H, Q, id = NULL, weights = NULL,
                      gradient = TRUE, get.post = FALSE){
  # Parameters and globals
  K <- ncol(X)
  N <- nrow(X)
  panel <- !is.null(id)
  if (is.null(weights)) weights <- 1
  if (panel) if (length(weights == 1)) weights <- rep(weights, N)
  
  # Coefs
  sigma2 <- theta[1:Q]
  betas <-  theta[-c(1:Q)]
  
  beta <- matrix(betas[1L:(K * Q)], nrow = K, ncol = Q)
  rownames(beta) <- colnames(X); colnames(beta) <- paste("class", 1:Q, sep = ":")
  gamma <- betas[-c(1L:(K * Q))]
  
  if (get.post) bi <- t(beta)
  
  # Make weights from a multinomial logit model
  ew <- lapply(H, function(x) exp(crossprod(t(x), gamma)))
  sew <- suml(ew)
  Wnq <- lapply(ew, function(x){ v <- x / sew;
  v[is.na(v)] <- 0;
  as.vector(v)})
  Wnq <- Reduce(cbind, Wnq) # W
  
  # Make Probit Probability
  index <- tcrossprod(X, t(beta))
  Pnq <- (2 * pi * sigma2) ^ (-1 / 2) * exp(-(1 / (2 * sigma2)) * (y - index) ^ 2)
  if (panel) Pnq <- apply(Pnq, 2, tapply, id, prod)
  WPnq <- Wnq * Pnq 
  Ln <- apply(WPnq, 1, sum)
  Ln <- pmax(Ln, .Machine$double.eps)
  Ln[is.na(Ln)] <- 0
  if (get.post) Qnr <- WPnq / Ln
  lnL <- if (panel) sum(log(Ln) * weights[!duplicated(id)]) else sum(log(Ln) * weights)
  
  ## Gradient
  if (gradient) {
    lambda <- (y - index) / sigma2 # N * Q
    Qnr <- WPnq / Ln
    if (panel) Qnr <- Qnr[id, ]
    eta <- lambda * Qnr
    etar <- eta[, rep(1:Q, each = K)]
    Xg <- X[, rep(1:K, Q)] # N * (K * Q)
    grad.beta <- Xg * etar
    grad.sigma <- Qnr * (-1 / (2 * sigma2) + (y - index) ^ 2 / (2 * sigma2 ^ 2))
    
    if (panel) {
      Wnq <- Wnq[id, ]
      H <- lapply(H, function(x) x[id, ])
    }
    Wg <- vector(mode = "list", length = Q)
    IQ <- diag(Q)
    for(q in 1:Q) Wg[[q]] <- rowSums(Qnr * (repRows(IQ[q, ], N) - repCols(Wnq[, q], Q)))
    grad.gamma <- suml(mapply("*", H, Wg, SIMPLIFY = FALSE))
    gari <- cbind(grad.sigma, grad.beta, grad.gamma)
    colnames(gari) <- names(theta)
    attr(lnL, "gradient") <- gari * weights
  }
  if (get.post) {
    attr(lnL, "bi") <- bi
    attr(lnL, "Qir") <- Qnr
    attr(lnL, "Wnq") <- Wnq
  }
  lnL
}

## suml function from mlogit (Croissant, 2013)
suml <- function(x){
  n <- length(x)
  if (!is.null(dim(x[[1]]))) {
    d <- dim(x[[1]])
    s <- matrix(0,d[1], d[2])
    for (i in 1:n) {
      x[[i]][is.na(x[[i]])] <- 0
      s <- s + x[[i]]
    }
  }
  else{
    s <- rep(0,length(x[[n]]))
    for (i in 1:n) {
      x[[i]][is.na(x[[i]])] <- 0
      s <- s + x[[i]]
    }
  }
  s
}

iv_probit <- function(formula, data, method = "bfgs",  ...){
  require("Formula")
  mf <- match.call(expand.dots = FALSE)
  m <- match(c("formula", "data", "subset", "na.action"), names(mf), 0)
  mf <- mf[c(1, m)]
  
  F1 <- Formula(formula)
  mf[[1]] <- as.name("model.frame")
  mf$formula <- F1
  mf <- eval(mf, parent.frame())
  
  y1 <- model.part(F1, data = mf, lhs = 1)[, 1]
  y2 <- model.part(F1, data = mf, lhs = 1)[, 2]
  X  <- model.matrix(F1, data = mf, rhs = 1)
  Z  <- model.matrix(F1, data = mf, rhs = 2)
  if (ncol(Z) < ncol(X)) stop("model underidentified")
  
  #Initial values
  fp <- as.formula(paste(as.character(attr(F1, "lhs")[[1]][2]), "~", 
                         paste(as.character(attr(F1, "rhs")[[1]])[-1], 
                               collapse = " + ")))
  fl <- as.formula(paste(as.character(attr(F1, "lhs")[[1]][3]), "~",
                         paste(as.character(attr(F1, "rhs")[[2]])[-1], 
                               collapse = " + ")))
  
  probit   <- glm(fp, data = mf,  family = binomial("probit"))
  linear   <- lm(fl, data = mf)
  sigma2   <- sum(resid(linear) ^ 2) / linear$df.residual
  lnsigma  <- log(sqrt(sigma2))
  athrho   <- 0
  start    <- c(coef(probit), coef(linear), lnsigma, athrho)
  nam_prob <- paste("eq.1", colnames(X), sep = ".")
  nam_lin  <- paste("eq.2", colnames(Z), sep = ".")
  names(start) <- c(nam_prob, nam_lin,  "lnsigma", "athrho")
  #print(start)
  
  require("maxLik")
  out <- maxLik(logLik = lnbinary_iv, start = start, y1 = y1, y2 = y2, X = X, Z = Z, 
                method =  method, ...)
  
  # Get post estimates
  post <- lnbinary_iv(coef(out), y1 = y1, y2 = y2, X = X, Z = Z, get.post =  TRUE)
  out$P1     <- attr(post, 'P1') 
  out$index1 <- attr(post, 'index1') 
  out$index2 <- attr(post, 'index2') 
  out$Pi     <- attr(post, 'Pi')     
  out$ai     <- attr(post, 'ai')  
  out
}


lnbinary_iv <- function(param, y1, y2, X, Z, get.post =  FALSE){
  F  <- pnorm
  f  <- dnorm
  ff <- function(x) -x * dnorm(x)
  mills <- function(x) f(x) / F(x)
  K      <- ncol(X)
  P      <- ncol(Z)
  beta   <- param[1L:K]
  delta  <- param[(K + 1L):(K + P)]
  lnsigma <- param[(K + P + 1L)]
  athrho  <- tail(param, n = 1L)
  sigma  <- exp(lnsigma)
  rho    <- (exp(2 * athrho) - 1) / (exp(2 * athrho) + 1)
  #if (rho < -1 || rho > 1) {
  #  warning("rho out of range")
  #  return(NA)
  #}
  index1 <- crossprod(t(X), beta)
  index2 <- crossprod(t(Z), delta)
  q      <- 2 * y1 - 1
  ai <- q * (index1 + (rho / sigma) * (y2 - index2)) / sqrt(1 - rho ^ 2)
  P1 <- pmax(F(ai), .Machine$double.eps)
  P1[is.nan(P1)] <- 0
  bi <- (y2 - index2) / sigma
  P2  <- (1 / sigma) * f(bi)
  P2[is.nan(P2)] <- 0
  Pi <- P1 * P2
  Pi[is.na(Pi)] <- 0
  Li     <- log(Pi)
  Li
  

  gb <-  X * drop(mills(ai) * q  / sqrt(1 - rho ^ 2))
  gd <- -Z * drop(mills(ai) * q * (rho / sigma) / (sqrt(1 - rho ^ 2)) + (ff(bi) / f(bi)) * (1 / sigma))
  glnsigma <- -(mills(ai) * q * rho / sqrt(1 - rho ^ 2) + (ff(bi) / f(bi))) * ((y2 - index2) * exp(-lnsigma)) - 1
  a <- index1; b <- sigma; c <- (y2 - index2)
  gathrho <- mills(ai) * q * (exp(-athrho) * ((c + a*b) * exp(2 * athrho)  + c - a*b)) / (2*b) 
  attr(Li,'gradient') <- cbind(gb, gd, glnsigma, gathrho)
  if (get.post) {
      attr(Li, 'P1')     <- P1
      attr(Li, 'index1') <- index1
      attr(Li, 'index2') <- index2
      attr(Li, 'Pi')     <- Li
      attr(Li, 'ai')     <- (index1 + (rho / sigma) * (y2 - index2)) / sqrt(1 - rho ^ 2)
  }
  Li
}


repRows <- function(x, n){
  matrix(rep(x, each = n), nrow = n)
}

repCols <- function(x, n){
  matrix(rep(x, each = n), ncol = n, byrow = TRUE)
}