.inter_prepare_model_data <- function(formula, survey, grid) {
  variables <- all.vars(formula[[3L]])
  missing_grid <- setdiff(variables, names(grid))
  if (length(missing_grid)) {
    stop("Every model predictor must be represented in the census margins. Missing: ",
         paste(missing_grid, collapse = ", "), call. = FALSE)
  }
  survey_model <- survey
  grid_model <- grid
  for (variable in variables) {
    survey_values <- survey[[variable]]
    grid_values <- grid[[variable]]
    if (is.numeric(survey_values) || is.integer(survey_values)) {
      converted <- suppressWarnings(as.numeric(as.character(grid_values)))
      if (anyNA(converted)) {
        stop("Numeric survey predictor `", variable,
             "` has non-numeric census categories.", call. = FALSE)
      }
      grid_model[[variable]] <- converted
    } else {
      levels_all <- unique(c(as.character(survey_values),
                             as.character(grid_values)))
      observed <- unique(as.character(survey_values[!is.na(survey_values)]))
      census_only <- setdiff(unique(as.character(grid_values)), observed)
      if (length(census_only)) {
        stop("Predictor `", variable,
             "` contains census levels absent from the entire survey: ",
             paste(census_only, collapse = ", "), call. = FALSE)
      }
      survey_model[[variable]] <- factor(as.character(survey_values),
                                         levels = levels_all)
      grid_model[[variable]] <- factor(as.character(grid_values),
                                       levels = levels_all)
    }
  }
  list(survey = survey_model, grid = grid_model, variables = variables)
}

.inter_calibrated_weights <- function(design, population) {
  dims <- setdiff(names(population$grid),
                  c(".cell", ".population", ".seed", ".prediction"))
  grid_key <- .inter_key(population$grid, dims)
  survey_key <- .inter_key(design$data, dims)
  cell <- match(survey_key, grid_key)
  original <- design$data[[design$weights]]
  sample_total <- numeric(nrow(population$grid))
  by_cell <- tapply(original[!is.na(cell)], cell[!is.na(cell)], sum,
                    na.rm = TRUE)
  if (length(by_cell)) {
    sample_total[as.integer(names(by_cell))] <- as.numeric(by_cell)
  }
  factor <- numeric(nrow(population$grid))
  supported <- sample_total > 0
  factor[supported] <- population$grid$.population[supported] /
    sample_total[supported]
  adjusted <- rep.int(0, length(original))
  adjusted[!is.na(cell)] <- original[!is.na(cell)] * factor[cell[!is.na(cell)]]
  adjusted
}

.inter_domain_estimates <- function(grid, survey, original_weights,
                                    calibrated_weights, outcome,
                                    prediction_survey, domain,
                                    residual_correction, family_label) {
  cell_key <- .inter_key(grid, domain)
  domain_keys <- unique(cell_key)
  representative <- grid[match(domain_keys, cell_key), domain, drop = FALSE]
  rownames(representative) <- NULL

  population <- as.numeric(.inter_group_sum(grid$.population, cell_key,
                                            domain_keys))
  synthetic_total <- as.numeric(.inter_group_sum(
    grid$.population * grid$.prediction, cell_key, domain_keys
  ))
  synthetic <- synthetic_total / population

  survey_key <- .inter_key(survey, domain)
  y <- survey[[outcome]]
  valid_model <- is.finite(y) & is.finite(prediction_survey) &
    is.finite(calibrated_weights)
  correction_total <- rep.int(0, length(domain_keys))
  if (isTRUE(residual_correction) && any(valid_model)) {
    correction_total <- as.numeric(.inter_group_sum(
      calibrated_weights[valid_model] *
        (y[valid_model] - prediction_survey[valid_model]),
      survey_key[valid_model], domain_keys
    ))
  }
  total <- synthetic_total + correction_total
  estimate <- total / population
  if (identical(family_label, "binomial")) {
    estimate <- pmin(1, pmax(0, estimate))
    total <- estimate * population
  }

  valid_direct <- is.finite(y) & is.finite(original_weights) &
    original_weights > 0
  direct_n <- as.numeric(table(factor(survey_key[valid_direct],
                                      levels = domain_keys)))
  direct_weight <- as.numeric(.inter_group_sum(
    original_weights[valid_direct], survey_key[valid_direct], domain_keys
  ))
  direct_total <- as.numeric(.inter_group_sum(
    original_weights[valid_direct] * y[valid_direct],
    survey_key[valid_direct], domain_keys
  ))
  direct <- ifelse(direct_weight > 0, direct_total / direct_weight, NA_real_)

  result <- cbind(
    representative,
    data.frame(
      estimate = estimate,
      total = total,
      population = population,
      synthetic = synthetic,
      residual_correction = correction_total / population,
      direct = direct,
      direct_n = direct_n,
      stringsAsFactors = FALSE
    )
  )
  result[is.finite(result$population) & result$population > 0, , drop = FALSE]
}

#' Fit joint margin-constrained intersectional estimates
#'
#' @param formula Model formula. The response is observed in the survey and
#'   every predictor must occur among the census-margin dimensions.
#' @param design An object returned by [inter_design()].
#' @param margins An object returned by [census_margins()].
#' @param family Either `"gaussian"`, `"binomial"`, or an R family object.
#' @param domain Character vector defining output domains. Defaults to area.
#' @param prior Positive seed added to every feasible latent cell.
#' @param max_iter Maximum generalized-IPF iterations.
#' @param tol Relative margin tolerance.
#' @param residual_correction Add a calibrated survey-residual correction.
#' @param max_cells Safety limit on the latent joint table.
#' @param model_args Optional named list passed to [stats::glm()].
#' @return An object of class `inter_sae`.
#' @export
fit_inter_sae <- function(formula, design, margins,
                          family = c("gaussian", "binomial"),
                          domain = NULL, prior = 0.5,
                          max_iter = 1000L, tol = 1e-8,
                          residual_correction = TRUE,
                          max_cells = 1e6, model_args = list()) {
  if (!inherits(design, "inter_design")) {
    stop("`design` must be created by `inter_design()`.", call. = FALSE)
  }
  if (!inherits(margins, "census_margins")) {
    stop("`margins` must be created by `census_margins()`.", call. = FALSE)
  }
  if (!inherits(formula, "formula") || length(formula) != 3L ||
      !is.symbol(formula[[2L]])) {
    stop("`formula` must contain one response and predictors.", call. = FALSE)
  }
  outcome <- all.vars(formula[[2L]])
  if (length(outcome) != 1L || !outcome %in% names(design$data)) {
    stop("The formula response must be one survey-data column.", call. = FALSE)
  }
  if (is.null(domain)) domain <- design$area
  if (!is.character(domain) || !length(domain) ||
      !all(domain %in% margins$dimensions)) {
    stop("`domain` must name one or more census dimensions.", call. = FALSE)
  }
  domain <- unique(domain)
  if (!is.numeric(max_iter) || length(max_iter) != 1L || max_iter < 1 ||
      !is.numeric(tol) || length(tol) != 1L || !is.finite(tol) || tol <= 0 ||
      !is.numeric(max_cells) || length(max_cells) != 1L || max_cells < 1) {
    stop("`max_iter` and `max_cells` must be positive; `tol` must be positive and finite.",
         call. = FALSE)
  }
  if (!is.list(model_args) ||
      any(names(model_args) %in% c("formula", "family", "data", "weights",
                                   "na.action", "model"))) {
    stop("`model_args` must be a list and may not override core model arguments.",
         call. = FALSE)
  }
  family_object <- if (missing(family)) {
    stats::gaussian()
  } else if (is.character(family)) {
    .inter_family(match.arg(family))
  } else {
    .inter_family(family)
  }
  family_label <- if (family_object$family %in% c("binomial", "quasibinomial")) {
    "binomial"
  } else {
    "gaussian"
  }
  y <- design$data[[outcome]]
  if (family_label == "binomial" &&
      any(!is.na(y) & !(y %in% c(0, 1)))) {
    stop("A binomial response must contain only 0, 1, or missing values.",
         call. = FALSE)
  }

  population <- .inter_reconstruct(
    design, margins, prior = prior, max_iter = as.integer(max_iter),
    tol = tol, max_cells = max_cells
  )
  prepared <- .inter_prepare_model_data(formula, design$data,
                                        population$grid)
  fitting_family <- if (family_label == "binomial") {
    stats::quasibinomial(link = family_object$link)
  } else {
    family_object
  }
  glm_call <- c(
    list(
      formula = formula,
      family = fitting_family,
      data = prepared$survey,
      weights = prepared$survey[[design$weights]],
      na.action = stats::na.exclude,
      model = TRUE
    ),
    model_args
  )
  model <- do.call(stats::glm, glm_call)
  prediction_cell <- as.numeric(stats::predict(
    model, newdata = prepared$grid, type = "response"
  ))
  if (any(!is.finite(prediction_cell))) {
    stop(paste0(
      "The outcome model cannot predict every census cell. ",
      "Simplify aliased interactions or collapse unsupported categories."
    ),
         call. = FALSE)
  }
  prediction_survey <- as.numeric(stats::predict(
    model, newdata = prepared$survey, type = "response"
  ))
  population$grid$.prediction <- prediction_cell
  calibrated_weights <- .inter_calibrated_weights(design, population)
  estimates <- .inter_domain_estimates(
    population$grid,
    prepared$survey,
    design$data[[design$weights]],
    calibrated_weights,
    outcome,
    prediction_survey,
    domain,
    residual_correction,
    family_label
  )

  structure(
    list(
      call = match.call(),
      formula = formula,
      outcome = outcome,
      family = family_object,
      family_label = family_label,
      domain = domain,
      design = design,
      margins = margins,
      population = population,
      model = model,
      estimates = estimates,
      calibrated_weights = calibrated_weights,
      control = list(
        prior = prior,
        max_iter = max_iter,
        tol = tol,
        residual_correction = residual_correction,
        max_cells = max_cells,
        model_args = model_args
      )
    ),
    class = "inter_sae"
  )
}

#' Extract intersectional estimates
#'
#' @param object A fitted `inter_sae` object.
#' @param statistic Return domain means/proportions, totals, or both.
#' @return A data frame.
#' @export
inter_estimates <- function(object, statistic = c("mean", "total", "both")) {
  if (!inherits(object, "inter_sae")) {
    stop("`object` must be a fitted `inter_sae` object.", call. = FALSE)
  }
  statistic <- match.arg(statistic)
  keep <- c(object$domain, "population", "direct", "direct_n")
  if (statistic %in% c("mean", "both")) {
    keep <- c(object$domain, "estimate", "synthetic", "residual_correction",
              "population", "direct", "direct_n")
  }
  if (statistic == "total") {
    keep <- c(object$domain, "total", "population", "direct_n")
  }
  if (statistic == "both") {
    keep <- c(object$domain, "estimate", "total", "synthetic",
              "residual_correction", "population", "direct", "direct_n")
  }
  object$estimates[keep]
}
