W <- function(...) {
  vars <- as.character(substitute(list(...)))[-1]
  structure(vars, class = "W_terms")
}

Z <- function(...) {
  vars <- as.character(substitute(list(...)))[-1]
  structure(vars, class = "Z_terms")
}

extract_vars <- function(expr) {
  if (is.symbol(expr)) {
    return(as.character(expr))
  }
  if (is.call(expr) && as.character(expr[[1]]) == "+") {
    return(c(
      extract_vars(expr[[2]]),
      extract_vars(expr[[3]])
    ))
  }
  stop("Unsupported expression inside W() or Z()")
}


parse_eiv_formula <- function(formula) {

  f <- as.formula(formula)
  lhs <- as.character(f[[2]])
  rhs <- f[[3]]

  if (!is.call(rhs)) {
    stop("Right-hand side must be a formula")
  }

  W_vars <- character(0)
  Z_vars <- character(0)

  ## Walk through RHS terms
  terms <- if (rhs[[1]] == "+") as.list(rhs)[-1] else list(rhs)

  for (term in terms) {

    if (is.call(term) && term[[1]] == as.name("W")) {
      W_vars <- extract_vars(term[[2]])
    }

    if (is.call(term) && term[[1]] == as.name("Z")) {
      Z_vars <- extract_vars(term[[2]])
    }
  }

  if (length(W_vars) == 0) {
    stop("No W(...) term found in formula")
  }

  list(
    y = lhs,
    W = W_vars,
    Z = Z_vars
  )
}

prepare_eiv_data <- function(data, y_name, W_names, Z_names) {

  if (!"unit" %in% names(data)) {
    stop("data must contain a column named 'unit' identifying statistical units")
  }

  units <- unique(data$unit)
  n     <- length(units)
  p     <- length(W_names)
  J     <- max(table(data$unit))

  ## Build W array: n × p × J
  W <- array(NA_real_, dim = c(n, p, J),
             dimnames = list(units, W_names, NULL))

  for (i in seq_along(units)) {
    rows <- data$unit == units[i]
    Wi   <- as.matrix(data[rows, W_names])
    W[i, , seq_len(nrow(Wi))] <- t(Wi)
  }

  ## One row per unit
  data_u <- data[!duplicated(data$unit), ]

  y <- data_u[[y_name]]

  if (length(Z_names) > 0) {

    Z_formula <- reformulate(Z_names, intercept = TRUE)

    terms_Z   <- terms(Z_formula, data = data_u)

    Zmat <- model.matrix(terms_Z, data_u)

  } else {

    Zmat <- matrix(1, nrow = nrow(data_u), ncol = 1)

    colnames(Zmat) <- "(Intercept)"

    terms_Z <- NULL

  }

  Wbar <- apply(W, c(1, 2), mean, na.rm = TRUE)

  list(
    y = y,
    Zmat = Zmat,
    W = W,
    Wbar = Wbar,
    terms_Z = terms_Z,
    contrasts_Z = attr(Zmat, "contrasts")
  )
}


