R/255_zzz_R_specific_visualize.R

In CVXR: Disciplined Convex Optimization

#####
## DO NOT EDIT THIS FILE!! EDIT THE SOURCE INSTEAD: rsrc_tree/zzz_R_specific/visualize.R
#####

## R-SPECIFIC: visualize() -- Smith form visualization of CVXR problems
##
## Entry point: visualize(problem, output = "text")
## Tree walker: .viz_walk_tree(expr) -> nested list (VisualizationNode)
## Model builder: .viz_build_model(problem) -> full JSON-serializable data model
##
## Output formats:
##   "text"  -- console display with cli
##   "json"  -- JSON dump (for HTML/Python interop)
##   "html"  -- interactive D3+KaTeX (Phase 2)
##   "latex" -- align* environments (Phase 3)
##   "tikz"  -- forest tree diagrams (Phase 3)


# ==========================================================================
# DCP diagnosis
# ==========================================================================

## Monotonicity label for display
.mono_label <- function(incr, decr) {
  if (incr && decr) return("affine-monotone")
  if (incr) return("increasing")
  if (decr) return("decreasing")
  "non-monotone"
}

## Diagnose WHY an expression fails DCP.
## Returns NULL if the expression IS DCP (fast path).
## For non-DCP atoms, returns a list with:
##   $atom_convex, $atom_concave  -- logical (intrinsic atom properties)
##   $reason         -- human-readable string
##   $arg_details    -- list of per-arg diagnostics (curvature, monotonicity, pass/fail)
.dcp_diagnose_expr <- function(expr) {
  ## Leaves are always DCP
  if (S7_inherits(expr, Leaf)) return(NULL)

  ## Constant expressions are always DCP
  if (is_constant(expr)) return(NULL)

  ## If already DCP, nothing to diagnose
  if (is_convex(expr) || is_concave(expr)) return(NULL)

  ## -- Atom-level properties -------------------------
  atom_cvx <- tryCatch(is_atom_convex(expr), error = function(e) FALSE)
  atom_ccv <- tryCatch(is_atom_concave(expr), error = function(e) FALSE)
  cls_name <- .clean_class_name(expr)

  ## -- Per-argument diagnostics ----------------------
  arg_details <- list()
  for (idx in seq_along(expr@args)) {
    arg <- expr@args[[idx]]
    arg_curv <- .latex_curvature(arg)
    arg_aff  <- is_affine(arg)
    arg_cvx  <- is_convex(arg)
    arg_ccv  <- is_concave(arg)
    incr     <- tryCatch(is_incr(expr, idx), error = function(e) FALSE)
    decr     <- tryCatch(is_decr(expr, idx), error = function(e) FALSE)

    ## Check convex composition for this arg
    cvx_ok <- arg_aff || (arg_cvx && incr) || (arg_ccv && decr)
    ## Check concave composition for this arg
    ccv_ok <- arg_aff || (arg_ccv && incr) || (arg_cvx && decr)

    arg_details[[idx]] <- list(
      index       = idx,
      curvature   = arg_curv,
      is_affine   = arg_aff,
      is_convex   = arg_cvx,
      is_concave  = arg_ccv,
      monotonicity = .mono_label(incr, decr),
      convex_rule_ok  = cvx_ok,
      concave_rule_ok = ccv_ok
    )
  }

  ## -- Build reason string ---------------------------
  if (!atom_cvx && !atom_ccv) {
    reason <- sprintf(
      "%s is neither convex nor concave (atom-level property).",
      cls_name
    )
  } else if (atom_cvx && !atom_ccv) {
    ## Atom is convex but composition rule fails
    bad <- vapply(arg_details, function(a) !a$convex_rule_ok, logical(1))
    bad_idx <- which(bad)
    reason <- sprintf(
      "%s is convex, but DCP convex composition fails at arg%s %s.",
      cls_name,
      if (length(bad_idx) > 1L) "s" else "",
      paste(bad_idx, collapse = ", ")
    )
  } else if (!atom_cvx && atom_ccv) {
    ## Atom is concave but composition rule fails
    bad <- vapply(arg_details, function(a) !a$concave_rule_ok, logical(1))
    bad_idx <- which(bad)
    reason <- sprintf(
      "%s is concave, but DCP concave composition fails at arg%s %s.",
      cls_name,
      if (length(bad_idx) > 1L) "s" else "",
      paste(bad_idx, collapse = ", ")
    )
  } else {
    ## Both convex and concave (affine atom), but args aren't affine
    bad <- vapply(arg_details, function(a) !a$is_affine, logical(1))
    bad_idx <- which(bad)
    reason <- sprintf(
      "%s is affine, but arg%s %s %s not affine.",
      cls_name,
      if (length(bad_idx) > 1L) "s" else "",
      paste(bad_idx, collapse = ", "),
      if (length(bad_idx) > 1L) "are" else "is"
    )
  }

  list(
    atom_convex  = atom_cvx,
    atom_concave = atom_ccv,
    reason       = reason,
    arg_details  = arg_details,
    fix_hints    = .dcp_fix_hints(expr)
  )
}

## Diagnose a constraint for DCP compliance.
## Returns NULL if the constraint IS DCP, or a list with:
##   $requirement -- what DCP requires (e.g., "convex" or "affine")
##   $actual      -- what the expression actually is
##   $reason      -- human-readable string
.dcp_diagnose_constraint <- function(con) {
  if (is_dcp(con)) return(NULL)
  con_type <- .clean_class_name(con)

  ## Determine what DCP requires for this constraint type
  if (con_type %in% c("Zero", "Equality")) {
    requirement <- "affine"
    expr <- con@args[[1L]]
    actual <- .latex_curvature(expr)
    reason <- sprintf(
      "%s constraint requires affine expression, but got %s.",
      con_type, actual
    )
  } else if (con_type %in% c("NonPos", "Inequality")) {
    requirement <- "convex"
    expr <- con@args[[1L]]
    actual <- .latex_curvature(expr)
    reason <- sprintf(
      "%s constraint requires convex expression, but got %s.",
      con_type, actual
    )
  } else if (con_type == "NonNeg") {
    requirement <- "concave"
    expr <- con@args[[1L]]
    actual <- .latex_curvature(expr)
    reason <- sprintf(
      "%s constraint requires concave expression, but got %s.",
      con_type, actual
    )
  } else if (con_type == "PSD") {
    requirement <- "affine"
    expr <- con@args[[1L]]
    actual <- .latex_curvature(expr)
    reason <- sprintf(
      "PSD constraint requires affine matrix expression, but got %s.",
      actual
    )
  } else {
    requirement <- "DCP-compliant"
    actual <- "unknown"
    reason <- sprintf("%s constraint is not DCP.", con_type)
  }

  list(
    type        = con_type,
    requirement = requirement,
    actual      = actual,
    reason      = reason
  )
}


## Generate constructive fix hints for non-DCP expressions.
## Returns a character vector of suggestions, or NULL if no hints available.
## Two sources: (A) curvature-dependent atoms with known fixes,
## (B) domain constraints via atom_domain() introspection.
.dcp_fix_hints <- function(expr) {
  if (S7_inherits(expr, Leaf)) return(NULL)
  cls <- .clean_class_name(expr)
  hints <- character(0L)

  ## -- (A) Curvature-dependent atoms ------------------
  if (cls == "QuadForm") {
    P <- expr@args[[2L]]
    if (!is_constant(P)) {
      hints <- c(hints, "P (second argument) must be constant (use a Parameter or numeric matrix)")
    } else if (!is_psd(P) && !is_nsd(P)) {
      hints <- c(hints, "P must be PSD for convexity, or NSD for concavity")
    } else if (is_psd(P)) {
      hints <- c(hints, "P is PSD, so quad_form(x, P) is convex -- check that x is affine")
    } else if (is_nsd(P)) {
      hints <- c(hints, "P is NSD, so quad_form(x, P) is concave -- check that x is affine")
    }
  } else if (cls == "SymbolicQuadForm") {
    orig <- expr@original_expression
    orig_hints <- .dcp_fix_hints(orig)
    if (!is.null(orig_hints)) hints <- c(hints, orig_hints)
  } else if (cls == "Perspective") {
    if (!is_nonneg(expr@args[[1L]])) {
      hints <- c(hints, "The scalar argument s must be nonneg for Perspective to be DCP")
    }
    f_expr <- expr@.f
    if (!is_convex(f_expr) && !is_concave(f_expr)) {
      hints <- c(hints, "The function f must be convex or concave")
    }
  } else if (cls == "Xexp") {
    if (!is_nonneg(expr@args[[1L]])) {
      hints <- c(hints, "x must be nonneg for xexp(x) to be convex")
    }
  }

  ## -- (B) Domain constraints via atom_domain() -------
  dom <- tryCatch(atom_domain(expr), error = function(e) list())
  for (con in dom) {
    con_cls <- .clean_class_name(con)
    ## Try to name the argument involved: find variable names in constraint
    arg_desc <- tryCatch({
      vars <- variables(con)
      if (length(vars) > 0L) {
        paste(vapply(vars, expr_name, character(1)), collapse = ", ")
      } else {
        "argument"
      }
    }, error = function(e) "argument")
    if (con_cls == "PSD") {
      hint <- sprintf("%s must be PSD (positive semidefinite)", arg_desc)
    } else if (con_cls %in% c("NonNeg", "Inequality")) {
      hint <- sprintf("%s must be nonneg (>= 0)", arg_desc)
    } else if (con_cls %in% c("NonPos")) {
      hint <- sprintf("%s must be nonpos (<= 0)", arg_desc)
    } else if (con_cls %in% c("Zero", "Equality")) {
      hint <- sprintf("%s must satisfy an equality constraint", arg_desc)
    } else {
      next
    }
    ## Avoid duplicating hints already generated by (A)
    if (!any(grepl(arg_desc, hints, fixed = TRUE))) {
      hints <- c(hints, hint)
    }
  }

  if (length(hints) == 0L) NULL else unique(hints)
}


# ==========================================================================
# Tree walker
# ==========================================================================

## Walk an expression tree, assigning auxiliary variables t_1, t_2, ... to
## each internal node. Returns a nested list (VisualizationNode) with all
## annotation data for every stage of the Smith form pipeline.
##
## @param expr An Expression (the subtree root).
## @param counter An environment with $i (integer counter for t_i naming).
## @returns A list with class "VizNode".

## Format a scalar value for display: integers shown without decimals,
## non-integers use `digits` significant figures (default 4, user-settable).
.viz_format_scalar <- function(val, digits = 4L) {
  v <- val[1L]
  if (is.finite(v) && v == round(v)) {
    ## Integer-valued: show without decimals (1, 0, -3)
    as.character(as.integer(v))
  } else {
    formatC(v, format = "g", digits = digits)
  }
}

.viz_walk_tree <- function(expr, counter) {
  ## -- Leaf nodes --------------------------------------------------

  if (S7_inherits(expr, Leaf)) {
    ## Leaves get their own name, not a t_i
    if (S7_inherits(expr, Variable)) {
      leaf_name <- expr_name(expr)
      ## Use custom latex_name if provided, otherwise bold variable name
      lnm <- expr@.latex_name
      leaf_latex <- if (nzchar(lnm)) lnm else sprintf("\\mathbf{%s}", leaf_name)
    } else if (S7_inherits(expr, Parameter)) {
      leaf_name <- expr_name(expr)
      lnm <- expr@.latex_name
      leaf_latex <- if (nzchar(lnm)) lnm else sprintf("\\theta_{\\text{%s}}", leaf_name)
    } else {
      ## Constant: try to show value
      val <- tryCatch(value(expr), error = function(e) NULL)
      if (!is.null(val) && length(val) == 1L) {
        leaf_name <- .viz_format_scalar(val, digits = counter$digits)
        leaf_latex <- leaf_name
      } else {
        leaf_name <- sprintf("C[%s]", paste(expr@shape, collapse = "x"))
        leaf_latex <- sprintf("C_{%s}", paste(expr@shape, collapse = "\\times "))
      }
    }

    ann <- smith_annotation(expr, aux_var = leaf_latex, child_vars = character(0))

    node <- list(
      id            = leaf_name,
      id_latex      = leaf_latex,
      expr_latex    = leaf_latex,   ## recursive expression = leaf name
      atom_class    = .clean_class_name(expr),
      is_leaf       = TRUE,
      shape         = as.integer(expr@shape),
      curvature     = .latex_curvature(expr),
      sign          = .latex_sign(expr),
      children      = list(),
      annotation    = ann
    )
    class(node) <- "VizNode"
    return(node)
  }

  ## -- Internal nodes (Atoms) -------------------------------------
  ## Recurse into children first (bottom-up)
  child_nodes <- lapply(expr@args, .viz_walk_tree, counter = counter)

  ## Assign auxiliary variable to this node
  counter$i <- counter$i + 1L
  aux_var <- sprintf("t_{%d}", counter$i)

  ## Collect child variable names for annotation.
  ## Brace-wrap each so that annotations can safely append subscripts
  ## (e.g. %s_i) without creating invalid double-subscripts like t_{1}_i.
  child_vars <- vapply(child_nodes, function(ch) {
    sprintf("{%s}", ch$id_latex)
  }, character(1))

  ## Get annotation from the atom (using t_i names for Smith form)
  ann <- smith_annotation(expr, aux_var = aux_var, child_vars = child_vars)

  ## Recursive expression: call annotation again with child expr_latex
  ## to get latex_definition in terms of the original variables
  child_exprs <- vapply(child_nodes, function(ch) {
    sprintf("{%s}", ch$expr_latex)
  }, character(1))
  expr_ann <- smith_annotation(expr, aux_var = aux_var, child_vars = child_exprs)
  expr_latex <- expr_ann$latex_definition

  node <- list(
    id            = sprintf("t_%d", counter$i),
    id_latex      = aux_var,
    expr_latex    = expr_latex,    ## full recursive expression
    atom_class    = .clean_class_name(expr),
    is_leaf       = FALSE,
    shape         = as.integer(expr@shape),
    curvature     = .latex_curvature(expr),
    sign          = .latex_sign(expr),
    children      = child_nodes,
    annotation    = ann
  )
  class(node) <- "VizNode"
  node
}


# ==========================================================================
# Model builder
# ==========================================================================

## Build the full visualization data model from a Problem.
## Returns a list with: problem metadata, objective tree, constraint trees,
## and flattened stage data (smith_form, relaxed_smith_form).

.viz_build_model <- function(problem, digits = 4L) {
  counter <- new.env(parent = emptyenv())
  counter$i <- 0L
  counter$digits <- digits

  ## Walk objective
  obj_expr <- problem@objective@args[[1L]]  # The expression inside Minimize/Maximize
  obj_tree <- .viz_walk_tree(obj_expr, counter)

  ## Walk each constraint
  constr_trees <- lapply(problem@constraints, function(con) {
    ## Constraints have args: for Inequality/Equality, walk the expression
    ## For cone constraints (SOC, PSD, etc.), walk args
    if (length(con@args) > 0L) {
      subtrees <- lapply(con@args, .viz_walk_tree, counter = counter)
    } else {
      subtrees <- list()
    }
    ## Build constraint expression in math form
    con_type <- .clean_class_name(con)
    sub_exprs <- vapply(subtrees, function(st) st$expr_latex, character(1))
    con_expr <- if (con_type == "Inequality" && length(sub_exprs) == 2L) {
      sprintf("%s \\leq %s", sub_exprs[[1L]], sub_exprs[[2L]])
    } else if (con_type == "Equality" && length(sub_exprs) == 2L) {
      sprintf("%s = %s", sub_exprs[[1L]], sub_exprs[[2L]])
    } else if (con_type == "Zero" && length(sub_exprs) >= 1L) {
      sprintf("%s = 0", sub_exprs[[1L]])
    } else if (con_type == "NonPos" && length(sub_exprs) >= 1L) {
      sprintf("%s \\leq 0", sub_exprs[[1L]])
    } else if (con_type == "NonNeg" && length(sub_exprs) >= 1L) {
      sprintf("%s \\geq 0", sub_exprs[[1L]])
    } else if (con_type == "PSD" && length(sub_exprs) >= 1L) {
      sprintf("%s \\succeq 0", sub_exprs[[1L]])
    } else if (con_type == "SOC" && length(sub_exprs) >= 2L) {
      sprintf("\\lVert %s \\rVert_2 \\leq %s", sub_exprs[[2L]], sub_exprs[[1L]])
    } else {
      paste(sub_exprs, collapse = ",\\; ")
    }

    list(
      type       = con_type,
      label      = if (nzchar(con@.label)) con@.label else NULL,
      subtrees   = subtrees,
      shape      = as.integer(con@shape),
      expr_latex = con_expr
    )
  })

  ## -- Flatten: collect all nodes in DFS order ----------------------
  all_nodes <- list()
  .collect_nodes <- function(node) {
    for (ch in node$children) .collect_nodes(ch)
    all_nodes[[length(all_nodes) + 1L]] <<- node
  }
  .collect_nodes(obj_tree)
  for (ct in constr_trees) {
    for (st in ct$subtrees) .collect_nodes(st)
  }

  ## -- Build stage data ---------------------------------------------
  sense <- if (S7_inherits(problem@objective, Minimize)) "minimize" else "maximize"

  ## Smith form: list of equations (one per non-leaf node)
  smith_eqs <- vapply(all_nodes, function(nd) {
    if (nd$is_leaf) return("")
    nd$annotation$smith
  }, character(1))
  smith_eqs <- smith_eqs[nzchar(smith_eqs)]

  ## Relaxed Smith form
  relaxed_eqs <- vapply(all_nodes, function(nd) {
    if (nd$is_leaf) return("")
    nd$annotation$relaxed
  }, character(1))
  relaxed_eqs <- relaxed_eqs[nzchar(relaxed_eqs)]

  ## Conic form: collect from nodes that have non-NULL conic
  conic_eqs <- list()
  for (nd in all_nodes) {
    if (nd$is_leaf) next
    ann_conic <- nd$annotation$conic
    if (!is.null(ann_conic)) {
      conic_eqs <- c(conic_eqs, ann_conic)
    } else {
      ## Stub: show relaxed form as placeholder
      ## Strip $ delimiters from relaxed to keep pure LaTeX
      relaxed_clean <- gsub("^\\$+|\\$+$", "", nd$annotation$relaxed)
      conic_eqs <- c(conic_eqs, list(
        sprintf("%s \\quad \\text{[conic form: see canonicalizer]}", relaxed_clean)
      ))
    }
  }

  list(
    problem = list(
      sense       = sense,
      n_vars      = length(variables(problem)),
      n_constrs   = length(problem@constraints)
    ),
    objective_tree = obj_tree,
    constraint_trees = constr_trees,
    stages = list(
      smith_form         = smith_eqs,
      relaxed_smith_form = relaxed_eqs,
      conic_form         = conic_eqs
    ),
    all_nodes = all_nodes
  )
}


# ==========================================================================
# Text renderer
# ==========================================================================

## Map curvature string to cli ANSI color for text-mode display
.viz_curv_color <- function(curv_str) {
  curv <- tolower(curv_str)
  if (grepl("convex", curv) && !grepl("concave", curv)) {
    cli::col_blue(curv_str)
  } else if (grepl("concave", curv)) {
    cli::col_red(curv_str)
  } else if (grepl("affine", curv)) {
    cli::col_green(curv_str)
  } else if (grepl("constant", curv)) {
    cli::col_grey(curv_str)
  } else {
    ## unknown / nondcp
    cli::col_yellow(curv_str)
  }
}

## Print a tree to console with indentation (box-drawing style)
.viz_print_tree <- function(node, prefix = "", child_prefix = "", is_root = TRUE) {
  ## Node label with curvature coloring
  colored_curv <- .viz_curv_color(node$curvature)
  if (node$is_leaf) {
    label <- paste0(node$id, "  [", colored_curv, ", ",
                    paste(node$shape, collapse = "x"), "]")
  } else {
    label <- paste0(node$id, " = ", node$atom_class, "(...)  [",
                    colored_curv, ", ", node$sign, ", ",
                    paste(node$shape, collapse = "x"), "]")
  }
  cat(paste0(prefix, label, "\n"))

  ## Recurse into children
  n_children <- length(node$children)
  for (i in seq_len(n_children)) {
    is_last_child <- (i == n_children)
    if (is_last_child) {
      .viz_print_tree(node$children[[i]],
                      prefix       = paste0(child_prefix, "\\-- "),
                      child_prefix = paste0(child_prefix, "    "),
                      is_root = FALSE)
    } else {
      .viz_print_tree(node$children[[i]],
                      prefix       = paste0(child_prefix, "|-- "),
                      child_prefix = paste0(child_prefix, "|   "),
                      is_root = FALSE)
    }
  }
}

## Print stage equations
.viz_print_stage <- function(title, eqs) {
  cli_rule(title)
  for (eq in eqs) {
    ## Strip $ delimiters for text display, show as-is otherwise
    clean <- gsub("\\$", "", eq)
    ## Simplify some LaTeX for text
    clean <- gsub("\\\\geq", ">=", clean)
    clean <- gsub("\\\\leq", "<=", clean)
    clean <- gsub("\\\\mathbf\\{([^}]*)\\}", "\\1", clean)
    clean <- gsub("\\\\mathcal\\{Q\\}", "Q", clean)
    clean <- gsub("\\\\mathcal\\{K\\}", "K", clean)
    clean <- gsub("\\\\lVert", "||", clean)
    clean <- gsub("\\\\rVert", "||", clean)
    clean <- gsub("\\\\varphi\\^\\{([^}]*)\\}", "phi^{\\1}", clean)
    clean <- gsub("\\\\text\\{([^}]*)\\}", "\\1", clean)
    clean <- gsub("\\\\texttt\\{([^}]*)\\}", "\\1", clean)
    clean <- gsub("\\\\operatorname\\{([^}]*)\\}", "\\1", clean)
    clean <- gsub("\\\\in", "in", clean)
    clean <- gsub("\\\\quad", "  ", clean)
    clean <- gsub("\\\\;", " ", clean)
    clean <- gsub("\\\\,", " ", clean)
    clean <- gsub("\\\\top", "'", clean)
    clean <- gsub("\\\\cdot", ".", clean)
    clean <- gsub("\\\\ldots", "...", clean)
    clean <- gsub("\\\\left|\\\\right", "", clean)
    clean <- gsub("\\\\left\\(", "(", clean)
    clean <- gsub("\\\\right\\)", ")", clean)
    clean <- gsub("\\\\frac\\{([^}]*)\\}\\{([^}]*)\\}", "(\\1)/(\\2)", clean)
    cat(sprintf("  %s\n", clean))
  }
}


## -- On-demand DCP analysis overlay --------------------------------------
## Computes DCP analysis as a SEPARATE pass -- never called from
## .viz_walk_tree() or .viz_build_model(). This ensures zero overhead
## on the normal visualization path. The cost of re-walking the expression
## tree is acceptable since this is only done when the user requests it
## (i.e., the problem is NOT DCP and we want to explain why).

## Walk the expression tree and VizNode tree in parallel, building a map
## from VizNode ID -> DCP violation info.
.viz_dcp_walk <- function(expr, viz_node, violations) {
  if (S7_inherits(expr, Leaf)) return()

  ## Diagnose this expression (returns NULL if DCP)
  diag <- .dcp_diagnose_expr(expr)
  if (!is.null(diag)) {
    violations[[viz_node$id]] <- list(
      id         = viz_node$id,
      atom_class = viz_node$atom_class,
      curvature  = viz_node$curvature,
      shape      = viz_node$shape,
      violation  = diag
    )
  }

  ## Recurse into children in parallel
  for (i in seq_along(expr@args)) {
    .viz_dcp_walk(expr@args[[i]], viz_node$children[[i]], violations)
  }
}

## Build the complete DCP overlay for a problem. Returns a list with:
##   $is_dcp         -- logical
##   $objective_ok   -- logical
##   $objective_issue -- string or NULL
##   $constraint_issues -- list of per-constraint diagnostics
##   $violations     -- named list: VizNode ID -> violation info
.viz_dcp_overlay <- function(problem, model) {
  prob_is_dcp <- is_dcp(problem)
  result <- list(is_dcp = prob_is_dcp)
  if (prob_is_dcp) return(result)

  sense <- model$problem$sense
  obj_expr <- problem@objective@args[[1L]]

  ## Objective analysis
  obj_is_dcp <- is_dcp(problem@objective)
  result$objective_ok <- obj_is_dcp
  if (!obj_is_dcp) {
    if (sense == "minimize") {
      result$objective_issue <- sprintf(
        "Minimize requires a convex objective, but expression is %s.",
        .latex_curvature(obj_expr)
      )
    } else {
      result$objective_issue <- sprintf(
        "Maximize requires a concave objective, but expression is %s.",
        .latex_curvature(obj_expr)
      )
    }
  }

  ## Constraint analysis
  constr_issues <- list()
  for (ci in seq_along(problem@constraints)) {
    diag <- .dcp_diagnose_constraint(problem@constraints[[ci]])
    if (!is.null(diag)) {
      diag$index <- ci
      constr_issues <- c(constr_issues, list(diag))
    }
  }
  result$constraint_issues <- constr_issues

  ## Per-node violations: walk expression tree + VizNode tree in parallel
  violations <- new.env(hash = TRUE, parent = emptyenv())
  .viz_dcp_walk(obj_expr, model$objective_tree, violations)
  for (ci in seq_along(problem@constraints)) {
    con <- problem@constraints[[ci]]
    ct <- model$constraint_trees[[ci]]
    for (ai in seq_along(con@args)) {
      .viz_dcp_walk(con@args[[ai]], ct$subtrees[[ai]], violations)
    }
  }
  result$violations <- as.list(violations)

  result
}


## Print DCP analysis section (text mode)
.viz_print_dcp_analysis <- function(problem, model) {
  dcp <- .viz_dcp_overlay(problem, model)
  cli_rule("DCP ANALYSIS")
  cat("  \u2717 Problem is NOT DCP compliant.\n")

  ## -- Objective ------------------------------------
  sense <- model$problem$sense
  if (!isTRUE(dcp$objective_ok)) {
    req <- if (sense == "minimize") "convex" else "concave"
    cat(sprintf("  Objective (%s): requires %s expression\n",
                tools::toTitleCase(sense), req))
    ## Print violations in the objective tree
    .viz_print_dcp_violations(model$objective_tree, dcp$violations, indent = "    ")
  } else {
    cat(sprintf("  Objective (%s): \u2713 DCP OK\n",
                tools::toTitleCase(sense)))
  }

  ## -- Constraint issues ----------------------------
  issues <- dcp$constraint_issues
  if (length(issues) > 0L) {
    for (ci in issues) {
      cat(sprintf("  Constraint [%d] (%s): requires %s expression\n",
                  ci$index, ci$type, ci$requirement))
      ct <- model$constraint_trees[[ci$index]]
      for (st in ct$subtrees) {
        .viz_print_dcp_violations(st, dcp$violations, indent = "    ")
      }
    }
  }
}

## Print DCP violations for a VizNode subtree, using the violations map
.viz_print_dcp_violations <- function(node, violations, indent = "  ") {
  if (node$is_leaf) return()

  viol <- violations[[node$id]]
  if (!is.null(viol)) {
    v <- viol$violation
    cat(sprintf("%s\u2717 %s: %s  [%s]\n",
                indent, node$atom_class, node$curvature,
                paste(node$shape, collapse = "x")))

    ## Per-arg details
    for (ad in v$arg_details) {
      pass_cvx <- if (ad$convex_rule_ok) "\u2713" else "\u2717"
      pass_ccv <- if (ad$concave_rule_ok) "\u2713" else "\u2717"
      cat(sprintf("%s  arg %d: %s, %s  (cvx-rule %s, ccv-rule %s)\n",
                  indent, ad$index, ad$curvature, ad$monotonicity,
                  pass_cvx, pass_ccv))
    }

    ## Reason
    cat(sprintf("%s  => %s\n", indent, v$reason))

    ## Fix hints (if any)
    if (!is.null(v$fix_hints)) {
      for (hint in v$fix_hints) {
        cat(sprintf("%s  Hint: %s\n", indent, hint))
      }
    }
  }

  ## Recurse into children
  for (ch in node$children) {
    .viz_print_dcp_violations(ch, violations, indent = paste0(indent, "  "))
  }
}


# ==========================================================================
# JSON renderer
# ==========================================================================

## Convert a VizNode tree to a JSON-compatible list (strip R classes)
.viz_node_to_list <- function(node) {
  children <- lapply(node$children, .viz_node_to_list)
  list(
    id         = node$id,
    id_latex   = node$id_latex,
    expr_latex = node$expr_latex,
    atom_class = node$atom_class,
    is_leaf    = node$is_leaf,
    shape      = node$shape,
    curvature  = node$curvature,
    sign       = node$sign,
    children   = children,
    annotation = node$annotation
  )
}

.viz_require_jsonlite <- function() {
  if (!requireNamespace("jsonlite", quietly = TRUE)) {
    cli_abort(c(
      "The {.pkg jsonlite} package is required for {.code json} and {.code html} output.",
      "i" = 'Install it with {.run install.packages("jsonlite")}.'
    ))
  }
}

.viz_model_to_json <- function(model, dcp_overlay = NULL) {
  .viz_require_jsonlite()
  json_model <- list(
    problem          = model$problem,
    objective_tree   = .viz_node_to_list(model$objective_tree),
    constraint_trees = lapply(model$constraint_trees, function(ct) {
      list(
        type       = ct$type,
        label      = ct$label,
        shape      = ct$shape,
        expr_latex = ct$expr_latex,
        subtrees   = lapply(ct$subtrees, .viz_node_to_list)
      )
    }),
    stages = model$stages
  )
  ## Inject DCP overlay if provided (only for non-DCP problems)
  if (!is.null(dcp_overlay)) {
    json_model$dcp_analysis <- dcp_overlay
  }
  ## Inject solver data if provided (Stages 4-5)
  if (!is.null(model$solver_data)) {
    json_model$solver_data <- model$solver_data
  }
  jsonlite::toJSON(json_model, auto_unbox = TRUE, pretty = TRUE, null = "null")
}


# ==========================================================================
# Solver data extraction (Stages 4-5: Standard Form + Solver Data)
# ==========================================================================
## These helpers are only called when visualize(solver = ...) is specified.
## They run entirely off the hot path -- never from solve() or problem_data()
## in normal usage.

## Build solver data summary for visualization.
## Calls problem_data() once and extracts a JSON-serializable summary.
##
## @param problem A Problem object.
## @param solver NULL (default solver) or character solver name.
## @param digits Significant digits for numeric display.
## @returns A list with solver_name, interface, dims, cone_product,
##   var_map, blocks, matrices.
.viz_build_solver_data <- function(problem, solver = NULL, digits = 4L) {
  pd <- problem_data(problem, solver = solver)
  data <- pd$data
  chain <- pd$chain
  inv_data <- pd$inverse_data

  ## Determine solver interface (conic vs QP)
  solver_reduction <- chain@reductions[[length(chain@reductions)]]
  interface <- if (S7_inherits(solver_reduction, QpSolver)) "qp" else "conic"
  slv_name <- solver_name(solver_reduction)

  ## Find the ConeMatrixStuffing inverse data (second-to-last reduction)
  cms_idx <- length(inv_data) - 1L
  cms_inv <- if (cms_idx >= 1L) inv_data[[cms_idx]] else NULL

  ## Build variable mapping table
  var_map <- .viz_var_mapping(cms_inv, variables(problem))

  ## Build dimension summary
  dims <- .viz_dims_summary(data, interface)

  ## Build cone product LaTeX
  cone_product <- .viz_cone_product_latex(data[["dims"]])

  ## Build block structure description
  blocks <- .viz_block_structure(data[["dims"]])

  ## Matrix summaries
  matrices <- .viz_matrix_summaries(data, interface, digits)

  list(
    solver_name  = slv_name,
    interface    = interface,
    dims         = dims,
    cone_product = cone_product,
    var_map      = var_map,
    blocks       = blocks,
    matrices     = matrices
  )
}

## Map solver variable indices back to user variables.
## Returns a list of lists, each with: name, offset, size, shape, origin.
.viz_var_mapping <- function(cms_inv, orig_vars) {
  if (is.null(cms_inv)) return(list())

  ## Build set of user variable IDs
  user_ids <- vapply(orig_vars, function(v) as.character(v@id), character(1))
  user_names <- vapply(orig_vars, expr_name, character(1))
  names(user_names) <- user_ids

  var_offsets <- cms_inv@var_offsets
  var_shapes <- cms_inv@var_shapes
  all_ids <- names(var_offsets)

  result <- vector("list", length(all_ids))
  for (i in seq_along(all_ids)) {
    vid <- all_ids[[i]]
    offset <- var_offsets[[vid]]
    shape <- var_shapes[[vid]]
    sz <- prod(shape)
    is_user <- vid %in% user_ids
    nm <- if (is_user) user_names[[vid]] else sprintf("aux_%s", vid)
    origin <- if (is_user) {
      sprintf("user variable (%s)", paste(shape, collapse = "x"))
    } else {
      "auxiliary (canonicalization)"
    }
    result[[i]] <- list(
      name   = nm,
      offset = offset,
      size   = sz,
      shape  = as.integer(shape),
      origin = origin
    )
  }
  result
}

## Build dimension summary from solver data.
.viz_dims_summary <- function(data, interface) {
  A <- data[["A"]]
  if (is.null(A)) {
    ## QP interface may use different field names
    A <- data[["A_eq"]]
  }

  n_vars <- length(data[["c"]])
  if (n_vars == 0L) n_vars <- length(data[["q"]])
  n_constraints <- if (!is.null(A)) nrow(A) else 0L
  nnz_A <- if (!is.null(A)) Matrix::nnzero(A) else 0L
  total <- as.double(n_vars) * as.double(n_constraints)
  sparsity <- if (total > 0) nnz_A / total else 0

  result <- list(
    n_vars        = n_vars,
    n_constraints = n_constraints,
    nnz_A         = nnz_A,
    sparsity      = sparsity
  )

  ## QP: also report nnz(P)
  P <- data[["P"]]
  if (interface == "qp" && !is.null(P)) {
    result$nnz_P <- Matrix::nnzero(P)
  }
  result
}

## Build LaTeX cone product string.
.viz_cone_product_latex <- function(cone_dims) {
  if (is.null(cone_dims)) return("")
  parts <- character(0)

  if (cone_dims@zero > 0L)
    parts <- c(parts, sprintf("\\{0\\}^{%d}", cone_dims@zero))
  if (cone_dims@nonneg > 0L)
    parts <- c(parts, sprintf("\\mathbb{R}_+^{%d}", cone_dims@nonneg))
  for (s in cone_dims@soc)
    parts <- c(parts, sprintf("\\mathcal{Q}^{%d}", s))
  for (s in cone_dims@psd)
    parts <- c(parts, sprintf("\\mathcal{S}_+^{%d}", s))
  if (cone_dims@exp > 0L)
    parts <- c(parts, sprintf("\\mathcal{K}_{\\exp}^{%d}", cone_dims@exp))
  for (alpha in cone_dims@p3d)
    parts <- c(parts, sprintf("\\mathcal{K}_{\\mathrm{pow}}^{%.2g}", alpha))
  for (alphas in cone_dims@pnd) {
    a_str <- paste(sprintf("%.2g", alphas), collapse = ",")
    parts <- c(parts, sprintf("\\mathcal{K}_{\\mathrm{pow}}^{(%s)}", a_str))
  }

  if (length(parts) == 0L) return("\\mathbb{R}^{0}")
  paste(parts, collapse = " \\times ")
}

## Build block structure description from cone dimensions.
## Returns a list of lists, each with: type, label, n_rows, latex.
.viz_block_structure <- function(cone_dims) {
  if (is.null(cone_dims)) return(list())
  blocks <- list()
  row <- 0L

  if (cone_dims@zero > 0L) {
    nr <- cone_dims@zero
    blocks <- c(blocks, list(list(
      type = "zero", label = "equalities",
      start = row + 1L, n_rows = nr,
      latex = sprintf("\\{0\\}^{%d}", nr)
    )))
    row <- row + nr
  }

  if (cone_dims@nonneg > 0L) {
    nr <- cone_dims@nonneg
    blocks <- c(blocks, list(list(
      type = "nonneg", label = "inequalities",
      start = row + 1L, n_rows = nr,
      latex = sprintf("\\mathbb{R}_+^{%d}", nr)
    )))
    row <- row + nr
  }

  for (i in seq_along(cone_dims@soc)) {
    nr <- cone_dims@soc[i]
    blocks <- c(blocks, list(list(
      type = "soc", label = sprintf("SOC(%d)", nr),
      start = row + 1L, n_rows = nr,
      latex = sprintf("\\mathcal{Q}^{%d}", nr)
    )))
    row <- row + nr
  }

  for (i in seq_along(cone_dims@psd)) {
    n <- cone_dims@psd[i]
    nr <- as.integer(n * (n + 1L) / 2L)
    blocks <- c(blocks, list(list(
      type = "psd", label = sprintf("PSD(%d)", n),
      start = row + 1L, n_rows = nr,
      latex = sprintf("\\mathcal{S}_+^{%d}", n)
    )))
    row <- row + nr
  }

  if (cone_dims@exp > 0L) {
    nr <- 3L * cone_dims@exp
    blocks <- c(blocks, list(list(
      type = "exp", label = sprintf("ExpCone x%d", cone_dims@exp),
      start = row + 1L, n_rows = nr,
      latex = sprintf("\\mathcal{K}_{\\exp}^{%d}", cone_dims@exp)
    )))
    row <- row + nr
  }

  if (length(cone_dims@p3d) > 0L) {
    nr <- 3L * length(cone_dims@p3d)
    blocks <- c(blocks, list(list(
      type = "pow3d", label = sprintf("PowCone3D x%d", length(cone_dims@p3d)),
      start = row + 1L, n_rows = nr,
      latex = sprintf("\\mathcal{K}_{\\mathrm{pow}}^{%d}", length(cone_dims@p3d))
    )))
    row <- row + nr
  }

  blocks
}

## Summarize each matrix/vector in the solver data.
.viz_matrix_summaries <- function(data, interface, digits) {
  .summarize_mat <- function(M, name) {
    if (is.null(M)) return(NULL)
    if (is.numeric(M) && !is.matrix(M)) {
      ## Vector
      list(
        name    = name,
        type    = "vector",
        length  = length(M),
        nnz     = sum(M != 0),
        range   = if (length(M) > 0L) c(min(M), max(M)) else c(0, 0)
      )
    } else {
      ## Matrix (possibly sparse)
      nr <- nrow(M)
      nc <- ncol(M)
      nnz <- if (inherits(M, "sparseMatrix")) Matrix::nnzero(M) else sum(M != 0)
      total <- as.double(nr) * as.double(nc)
      density <- if (total > 0) nnz / total else 0

      res <- list(
        name    = name,
        type    = "matrix",
        nrow    = nr,
        ncol    = nc,
        nnz     = nnz,
        density = density
      )

      ## For small problems, include sparse triplet for sparsity pattern
      if (total > 0 && total < 2000 && inherits(M, "sparseMatrix")) {
        trip <- as(M, "TsparseMatrix")
        res$triplet <- list(
          i = as.integer(trip@i + 1L),  # 1-based
          j = as.integer(trip@j + 1L),
          x = round(trip@x, digits)
        )
      }
      res
    }
  }

  result <- list()
  if (interface == "conic") {
    result$A <- .summarize_mat(data[["A"]], "A")
    result$b <- .summarize_mat(data[["b"]], "b")
    result$c <- .summarize_mat(data[["c"]], "c")
    if (!is.null(data[["P"]])) result$P <- .summarize_mat(data[["P"]], "P")
  } else {
    ## QP interface
    result$P <- .summarize_mat(data[["P"]], "P")
    result$q <- .summarize_mat(data[["q"]], "q")
    result$A_eq <- .summarize_mat(data[["A_eq"]], "A_eq")
    result$b_eq <- .summarize_mat(data[["b_eq"]], "b_eq")
    result$F_ineq <- .summarize_mat(data[["F_ineq"]], "F_ineq")
    result$g_ineq <- .summarize_mat(data[["g_ineq"]], "g_ineq")
  }
  result
}


# ==========================================================================
# Text renderer for Stages 4-5
# ==========================================================================

## Print solver data for text output
.viz_print_solver_data <- function(sd) {
  cli_rule("STANDARD CONE FORM ({sd$solver_name})")

  if (sd$interface == "conic") {
    cat("  minimize  c'x\n")
    cat("  s.t.      Ax + s = b,  s in K\n\n")
  } else {
    cat("  minimize  (1/2)x'Px + q'x\n")
    cat("  s.t.      A_eq x = b_eq\n")
    cat("            F x <= g\n\n")
  }

  ## Dimension table
  d <- sd$dims
  cat(sprintf("  Variables (n):     %d\n", d$n_vars))
  cat(sprintf("  Constraints (m):   %d\n", d$n_constraints))
  cat(sprintf("  nnz(A):           %d\n", d$nnz_A))
  cat(sprintf("  Density:          %.1f%%\n", d$sparsity * 100))
  if (!is.null(d$nnz_P))
    cat(sprintf("  nnz(P):           %d\n", d$nnz_P))

  ## Cone product (LaTeX-to-text)
  if (sd$interface == "conic" && nzchar(sd$cone_product)) {
    kstr <- sd$cone_product
    kstr <- gsub("\\\\times", "x", kstr)
    kstr <- gsub("\\\\\\{0\\\\\\}", "{0}", kstr)
    kstr <- gsub("\\\\mathbb\\{R\\}_\\+", "R+", kstr)
    kstr <- gsub("\\\\mathcal\\{Q\\}", "Q", kstr)
    kstr <- gsub("\\\\mathcal\\{S\\}_\\+", "S+", kstr)
    kstr <- gsub("\\\\mathcal\\{K\\}_\\{\\\\exp\\}", "K_exp", kstr)
    kstr <- gsub("\\\\mathcal\\{K\\}_\\{\\\\mathrm\\{pow\\}\\}", "K_pow", kstr)
    kstr <- gsub("\\^\\{([^}]*)\\}", "^\\1", kstr)
    cat(sprintf("\n  K = %s\n", kstr))
  }

  ## Variable mapping
  if (length(sd$var_map) > 0L) {
    cat("\n  Variable Mapping:\n")
    for (vm in sd$var_map) {
      idx_lo <- vm$offset + 1L
      idx_hi <- vm$offset + vm$size
      cat(sprintf("    x[%d:%d]  %-12s  %s\n",
                  idx_lo, idx_hi, vm$name, vm$origin))
    }
  }

  ## Block structure
  if (length(sd$blocks) > 0L) {
    cat("\n  Block Structure of A:\n")
    for (blk in sd$blocks) {
      cat(sprintf("    rows %d-%d  (%d rows)  %s\n",
                  blk$start, blk$start + blk$n_rows - 1L,
                  blk$n_rows, blk$label))
    }
  }

  ## Matrix summaries
  if (length(sd$matrices) > 0L) {
    cli_rule("SOLVER DATA MATRICES")
    for (ms in sd$matrices) {
      if (is.null(ms)) next
      if (ms$type == "vector") {
        cat(sprintf("  %s: vector, length %d, nnz %d, range [%s, %s]\n",
                    ms$name, ms$length, ms$nnz,
                    formatC(ms$range[1], format = "g", digits = 4L),
                    formatC(ms$range[2], format = "g", digits = 4L)))
      } else {
        cat(sprintf("  %s: %d x %d, nnz %d (%.1f%%)\n",
                    ms$name, ms$nrow, ms$ncol,
                    ms$nnz, ms$density * 100))
      }
    }
  }
}


# ==========================================================================
# Entry point: visualize()
# ==========================================================================

#' Visualize the Canonicalization Pipeline of a CVXR Problem
#'
#' Displays the Smith form decomposition of a convex optimization problem,
#' showing each stage of the DCP canonicalization pipeline:
#' expression tree, Smith form, relaxed Smith form, conic form, and
#' (optionally) standard cone form and solver data.
#'
#' @param problem A [Problem] object.
#' @param output Character: output format.
#'   \describe{
#'     \item{\code{"text"}}{Console display (default).}
#'     \item{\code{"json"}}{JSON data model (for interop with HTML/Python).}
#'     \item{\code{"html"}}{Interactive D3+KaTeX HTML (Phase 2).}
#'     \item{\code{"latex"}}{LaTeX align* environments (Phase 3).}
#'     \item{\code{"tikz"}}{TikZ forest tree diagrams (Phase 3).}
#'   }
#' @param solver Solver specification for matrix stuffing stages (4-5).
#'   \code{NULL} (default) shows only Stages 0-3 with zero overhead.
#'   \code{TRUE} uses the default solver (same as [psolve()]).
#'   A character string (e.g., \code{"Clarabel"}) uses that specific solver.
#' @param digits Integer: significant digits for displaying scalar constants.
#'   Integer-valued constants (0, 1, -3) always display without decimals
#'   regardless of this setting.  Defaults to 4.
#' @param file Character: path for HTML output file.  If \code{NULL}
#'   (default), a temporary file is used.
#' @param open Logical: whether to open the HTML file in a browser.
#'   Defaults to \code{TRUE} in interactive sessions.
#' @param doc_base Character: base URL for atom documentation links.
#'   Defaults to the CVXR pkgdown site.
#' @returns For \code{"text"}: invisible model list.
#'   For \code{"json"}: a JSON string (or list if jsonlite not available).
#'   For \code{"html"}: the file path (invisibly).
#'   For other formats: the rendered output (Phase 2+).
#'
#' @examples
#' \dontrun{
#' x <- Variable(3, name = "x")
#' prob <- Problem(Minimize(p_norm(x, 2)), list(x >= 1))
#' visualize(prob)                          # Stages 0-3 only
#' visualize(prob, solver = TRUE)           # Stages 0-5, default solver
#' visualize(prob, solver = "Clarabel")     # Stages 0-5, specific solver
#' visualize(prob, output = "html", solver = TRUE)
#' }
#'
#' @export
visualize <- function(problem,
                      output = c("text", "json", "html", "latex", "tikz"),
                      solver = NULL,
                      digits = 4L,
                      file = NULL,
                      open = interactive(),
                      doc_base = "https://cvxr.rbind.io/reference/") {
  if (!S7_inherits(problem, Problem)) {
    cli_abort("{.fn visualize} requires a {.cls Problem} object.")
  }
  output <- match.arg(output)

  ## Build the data model (Stages 0-3)
  model <- .viz_build_model(problem, digits = as.integer(digits))

  ## Solver data (Stages 4-5) -- only when solver is specified
  if (!is.null(solver)) {
    if (!is_dcp(problem)) {
      cli_warn("Solver data requires a DCP-compliant problem. Skipping Stages 4-5.")
    } else {
      effective_solver <- if (isTRUE(solver)) NULL else solver
      model$solver_data <- tryCatch(
        .viz_build_solver_data(problem, solver = effective_solver,
                               digits = as.integer(digits)),
        error = function(e) {
          cli_warn("Failed to extract solver data: {conditionMessage(e)}")
          NULL
        }
      )
    }
  }

  switch(output,
    text = {
      sense <- model$problem$sense

      ## Expression tree
      cli_rule("Expression Tree ({toupper(sense)})")
      .viz_print_tree(model$objective_tree)

      if (length(model$constraint_trees) > 0L) {
        cli_rule("Constraints")
        for (i in seq_along(model$constraint_trees)) {
          ct <- model$constraint_trees[[i]]
          con <- problem@constraints[[i]]
          con_type <- sub("^CVXR::", "", ct$type)
          lbl <- if (!is.null(ct$label)) paste0(" (", ct$label, ")") else ""
          dcp_mark <- if (is_dcp(con)) cli::col_green("\u2713") else cli::col_red("\u2717")
          cli_text("  {dcp_mark} [{i}] {con_type} {paste(ct$shape, collapse = 'x')}{lbl}")
          for (st in ct$subtrees) {
            .viz_print_tree(st, prefix = "      ", child_prefix = "      ")
          }
        }
      }

      ## DCP analysis (only when problem is NOT DCP -- computed on demand)
      if (!is_dcp(problem)) {
        .viz_print_dcp_analysis(problem, model)
      }

      ## Smith form
      .viz_print_stage("SMITH FORM", model$stages$smith_form)

      ## Relaxed Smith form
      .viz_print_stage("RELAXED SMITH FORM", model$stages$relaxed_smith_form)

      ## Conic form
      .viz_print_stage("CONIC FORM", model$stages$conic_form)

      ## Stages 4-5: solver data (only when solver specified)
      if (!is.null(model$solver_data)) {
        .viz_print_solver_data(model$solver_data)
      }

      invisible(model)
    },

    json = {
      dcp_ov <- if (!is_dcp(problem)) .viz_dcp_overlay(problem, model)
      .viz_model_to_json(model, dcp_overlay = dcp_ov)
    },

    html = {
      dcp_ov <- if (!is_dcp(problem)) .viz_dcp_overlay(problem, model)
      .viz_render_html(model, file = file, open = open, doc_base = doc_base,
                       dcp_overlay = dcp_ov)
    },

    latex = {
      cli_inform("LaTeX output will be available in Phase 3.")
      invisible(model)
    },

    tikz = {
      cli_inform("TikZ output will be available in Phase 3.")
      invisible(model)
    }
  )
}