Source: utils/pythonGenerator.js

/**
 * Python Code Generator for AutoRA Workflows
 *
 * Generates executable Python code from workflow graph state.
 *
 * @module utils/pythonGenerator
 */

/**
 * One execution block produced by getExecutionOrder: a `once` block runs its
 * nodes a single time, while a `loop` block runs its child blocks inside
 * `for cycle_N in range(maxCounter)` (nested loops become nested block trees).
 *
 * @typedef {{type: 'once', nodes: Object[]} | {type: 'loop', maxCounter: number, children: Block[]}} Block
 */

export const CONTROL_NODE_TYPES = ['start_point', 'end_point', 'filter_point']

// Python code templates
export const TEMPLATES = {
  header: (date) => `"""
AutoRA Workflow

Generated by AutoRA Workflow Editor
Generated on: ${date}
"""`,

  standardImports: `from autora.state import on_state, Delta, estimator_on_state, StandardState`,

  dataImports: `import pandas as pd`,

  defaultVariables: `    # Initialize variables - customize this based on your experiment
    # You may need to define your own variables or get them from an experiment runner
    variables = VariableCollection(
        independent_variables=[
            Variable(name="x", allowed_values=np.linspace(-1, 1, 100))
        ],
        dependent_variables=[
            Variable(name="y")
        ]
    )`,

  initState: `    # Initialize state
    state = StandardState(variables=variables)`,

  // Service-account credential entries for firebase runners. All fields are
  // placeholders the user must replace with values from their Firebase
  // service-account JSON.
  firebaseCredentialEntries: `            "type": "service_account",
            "project_id": "project_id",
            "private_key_id": "private_key_id",
            "private_key": "-----BEGIN PRIVATE KEY-----\\n...\\n-----END PRIVATE KEY-----\\n",
            "client_email": "xxx@iam.gserviceaccount.com",
            "client_id": "001",
            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
            "token_uri": "https://oauth2.googleapis.com/token",
            "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
            "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/...",
            "universe_domain": "googleapis.com"`
}

/**
 * Simple code builder for cleaner Python generation
 */
export class CodeBuilder {
  constructor() {
    this.lines = []
  }

  /**
   * Append a single line of text.
   *
   * @param {string} [text=''] - Line contents.
   * @returns {CodeBuilder} This builder, for chaining.
   */
  line(text = '') {
    this.lines.push(text)
    return this
  }

  /**
   * Append a line prefixed with the given number of 4-space indent levels.
   *
   * @param {string} text - Line contents.
   * @param {number} [level=1] - Number of indent levels.
   * @returns {CodeBuilder} This builder, for chaining.
   */
  indent(text, level = 1) {
    this.lines.push('    '.repeat(level) + text)
    return this
  }

  /**
   * Append a `#`-prefixed Python comment at the given indent level.
   *
   * @param {string} text - Comment text (without the leading `#`).
   * @param {number} [level=0] - Number of indent levels.
   * @returns {CodeBuilder} This builder, for chaining.
   */
  comment(text, level = 0) {
    return this.indent(`# ${text}`, level)
  }

  /**
   * Append a blank line.
   *
   * @returns {CodeBuilder} This builder, for chaining.
   */
  blank() {
    return this.line('')
  }

  /**
   * Append a multi-line string as separate lines.
   *
   * @param {string} text - Newline-separated text.
   * @returns {CodeBuilder} This builder, for chaining.
   */
  multiline(text) {
    text.split('\n').forEach(l => this.line(l))
    return this
  }

  /**
   * Render all accumulated lines as a single newline-joined string.
   *
   * @returns {string} The built source code.
   */
  toString() {
    return this.lines.join('\n')
  }
}

/**
 * Fold a set of filter intervals over an ordered component list into a nested
 * block tree spanning `[lo, hi]`. Intervals passed in are all within that range.
 * Each range's *top-level* intervals (those not strictly contained in another)
 * become loop blocks; their inner intervals recurse into child loops; and any
 * component not covered by a top-level interval accumulates into a `once` block.
 *
 * @param {Object[]} orderedNodes - Components in execution order.
 * @param {number} lo - First component index of this range (inclusive).
 * @param {number} hi - Last component index of this range (inclusive).
 * @param {Array<{start: number, end: number, maxCounter: number}>} intervals - Filter intervals within `[lo, hi]`.
 * @returns {Block[]} Ordered child blocks covering `[lo, hi]`.
 */
function buildBlockTree(orderedNodes, lo, hi, intervals) {
  const topLevel = intervals
    .filter(iv => !intervals.some(o =>
      o !== iv && o.start <= iv.start && iv.end <= o.end &&
      (o.start < iv.start || iv.end < o.end)))
    .sort((a, b) => a.start - b.start)

  const blocks = []
  let onceRun = []
  const flushOnce = () => {
    if (onceRun.length) { blocks.push({ type: 'once', nodes: onceRun }); onceRun = [] }
  }

  let pos = lo
  let ti = 0
  while (pos <= hi) {
    const iv = topLevel[ti]
    if (iv && iv.start === pos) {
      flushOnce()
      const inner = intervals.filter(o => o !== iv && o.start >= iv.start && o.end <= iv.end)
      blocks.push({
        type: 'loop',
        maxCounter: iv.maxCounter,
        children: buildBlockTree(orderedNodes, iv.start, iv.end, inner)
      })
      pos = iv.end + 1
      ti++
      while (topLevel[ti] && topLevel[ti].start < pos) ti++
    } else {
      onceRun.push(orderedNodes[pos])
      pos++
    }
  }
  flushOnce()
  return blocks
}

function assertNestedOrDisjointIntervals(intervals) {
  const sorted = [...intervals].sort((a, b) => {
    if (a.start !== b.start) return a.start - b.start
    return b.end - a.end
  })

  const stack = []
  sorted.forEach(interval => {
    while (stack.length && interval.start > stack[stack.length - 1].end) stack.pop()
    const parent = stack[stack.length - 1]
    if (parent && interval.end > parent.end) {
      throw new Error(
        'Filter loops must be either disjoint or fully nested. ' +
        'Partially overlapping loops are not supported.'
      )
    }
    if (parent && interval.start === parent.start && interval.end === parent.end) {
      throw new Error(
        'Two Filter loops enclose exactly the same components. ' +
        'Remove one, or place a component between them to nest.'
      )
    }
    stack.push(interval)
  })
}

/**
 * Flatten a (possibly nested) block tree into its components in execution order.
 *
 * @param {Block[]} blocks - Execution blocks from getExecutionOrder.
 * @returns {Object[]} All component nodes, in order.
 */
export function flattenBlockNodes(blocks) {
  return blocks.flatMap(b => (b.type === 'loop' ? flattenBlockNodes(b.children) : b.nodes))
}

/**
 * Traverse the workflow graph and split it into a (possibly nested) tree of
 * execution blocks, supporting any number of loops — including nested loops
 * (each Filter node defines one loop).
 *
 * The graph is walked forward from the Start node. At each Filter the traversal
 * follows the *exit* output (toward the end) and records the *loop-back* output,
 * which points to an already-visited node marking where that loop's body began.
 * Each filter therefore spans an interval over the ordered components (loop-back
 * target → last component before the filter). Intervals that contain one another
 * become nested loops, disjoint intervals become sibling loops, and components
 * covered by no interval run once.
 *
 * @param {Object[]} nodes - Graph nodes, each with `id`, `type` and optional `filterParams`.
 * @param {Object[]} connections - Graph edges, each with `sourceId` and `targetId`.
 * @returns {{blocks: Block[]}} A tree of ordered execution blocks, where a
 *   `Block` is either `{type: 'once', nodes: Object[]}` (its nodes run a single
 *   time) or `{type: 'loop', maxCounter: number, children: Block[]}` (its child
 *   blocks run inside `for cycle_N in range(maxCounter)`).
 */
export function getExecutionOrder(nodes, connections) {
  const startNode = nodes.find(n => n.type === 'start_point')
  if (!startNode) {
    throw new Error('Workflow must have a Start node')
  }

  if (!nodes.some(n => n.type === 'end_point')) {
    throw new Error('Workflow must have an End node')
  }

  const nodeById = new Map(nodes.map(n => [n.id, n]))

  // Build adjacency map (source -> [targets])
  const adjacency = {}
  connections.forEach(conn => {
    if (!adjacency[conn.sourceId]) adjacency[conn.sourceId] = []
    adjacency[conn.sourceId].push(conn.targetId)
  })

  // Forward traversal from Start. Follow each node's (single) forward edge,
  // taking the exit branch at filters. Record components in order and filters
  // with the loop-back target that closes their loop.
  const sequence = []   // items: { node } | { loopBackId, maxCounter }
  const seen = []       // ids of components visited so far (loop-back detection)
  const visited = new Set()
  let current = startNode.id

  while (current != null && !visited.has(current)) {
    visited.add(current)
    const node = nodeById.get(current)
    if (!node || node.type === 'end_point') break

    if (node.type === 'start_point') {
      current = (adjacency[current] || [])[0]
      continue
    }

    if (node.type === 'filter_point') {
      const targets = adjacency[current] || []
      // The loop-back branch points to a node already visited (an ancestor);
      // the other branch is the exit that continues toward the end.
      const loopBackId = targets.find(id => seen.includes(id)) ?? null
      if (!loopBackId) {
        throw new Error(
          'The Filter node is missing its loop-back output connection that closes ' +
          'the experiment loop. Please connect the Filter back to a component to ' +
          'form the loop before generating code.'
        )
      }
      sequence.push({ loopBackId, maxCounter: node.filterParams?.maxCounter ?? 1 })
      current = targets.find(id => id !== loopBackId) ?? null
      continue
    }

    // Regular component
    seen.push(current)
    sequence.push({ node })
    current = (adjacency[current] || [])[0]
  }

  // Turn the linear sequence into ordered components plus one interval per
  // filter (loop-back target index → last component index before the filter),
  // then fold those intervals into a nested block tree.
  const orderedNodes = []
  const indexById = new Map()
  const intervals = []
  for (const item of sequence) {
    if (item.node) {
      indexById.set(item.node.id, orderedNodes.length)
      orderedNodes.push(item.node)
    } else {
      intervals.push({
        start: indexById.get(item.loopBackId),
        end: orderedNodes.length - 1,
        maxCounter: item.maxCounter
      })
    }
  }

  assertNestedOrDisjointIntervals(intervals)

  // A filter-less workflow has no interval — every component simply runs once.
  const blocks = buildBlockTree(orderedNodes, 0, orderedNodes.length - 1, intervals)

  return { blocks }
}

/**
 * Generate a valid Python variable name from a component name.
 * Parenthesized qualifiers like "(Synthetic, Economics)" are dropped.
 *
 * @param {string} name - Human-readable component name.
 * @returns {string} A lowercase, underscore-separated identifier.
 */
export function toPythonName(name) {
  return name
    .replace(/\([^)]*\)/g, ' ')
    .toLowerCase()
    .replace(/[^a-z0-9]+/g, '_')
    .replace(/^_+|_+$/g, '')
    .replace(/_+/g, '_')
}

/**
 * Format a JavaScript value as Python literal
 *
 * @param {*} value - Value to format (null/undefined, boolean, string, array or number).
 * @returns {string} The Python literal representation.
 */
function formatPythonValue(value) {
  if (value === null || value === undefined) return 'None'
  if (typeof value === 'boolean') return value ? 'True' : 'False'
  if (typeof value === 'string') {
    // A string that already holds a Python container literal (dict or list),
    // e.g. a "dict"-typed field like fixed_effects, is emitted verbatim so it
    // is not wrapped in quotes and stays a real dict/list in the generated code.
    const trimmed = value.trim()
    if (/^\{[\s\S]*\}$/.test(trimmed) || /^\[[\s\S]*\]$/.test(trimmed)) return trimmed
    return `"${value.replace(/"/g, '\\"')}"`
  }
  if (Array.isArray(value)) return `[${value.map(formatPythonValue).join(', ')}]`
  return String(value)
}

/**
 * Build parameter string from params object
 *
 * @param {Object} params - Map of parameter name to value.
 * @param {string[]} [exclude=[]] - Parameter names to omit.
 * @returns {string} Comma-separated `name=value` keyword arguments (nulls skipped).
 */
function buildParamString(params, exclude = []) {
  return Object.entries(params)
    .filter(([k, v]) => v !== null && v !== undefined && !exclude.includes(k))
    .map(([k, v]) => `${k}=${formatPythonValue(v)}`)
    .join(', ')
}

// The import that a `sympify`-flagged param needs. Which params are SymPy
// expressions is declared per-component in the JSON (a param with
// `"sympify": true`); this is only the fixed fact of where `sympify` lives.
const SYMPIFY_IMPORT = { module: 'sympy', name: 'sympify' }

/**
 * Whether a param's value is a blank (empty/whitespace-only) string. Such a
 * value counts as "unset" — the user cleared the input — so the param is omitted
 * and the runner's own default applies rather than emitting an invalid literal
 * (e.g. `sympify("")`, which raises at runtime).
 *
 * @param {*} value - The param value.
 * @returns {boolean}
 */
function isBlankString(value) {
  return typeof value === 'string' && value.trim() === ''
}

/**
 * Build a factory/constructor parameter string, wrapping any params named in
 * `sympifyNames` (declared `"sympify": true` in the component JSON) in
 * `sympify(...)` so their string value is parsed into a SymPy expression. A
 * sympify param left blank is treated as unset and omitted (so the runner's
 * default applies) rather than emitting `sympify("")`, which would raise.
 *
 * @param {Object} params - Map of parameter name to value.
 * @param {string[]} [sympifyNames=[]] - Names of params to wrap in `sympify(...)`.
 * @param {string[]} [exclude=[]] - Parameter names to omit.
 * @returns {string} Comma-separated `name=value` keyword arguments (nulls skipped).
 */
function buildFactoryParamString(params, sympifyNames = [], exclude = []) {
  return Object.entries(params)
    .filter(([k, v]) => v !== null && v !== undefined && !exclude.includes(k) &&
      !(sympifyNames.includes(k) && isBlankString(v)))
    .map(([k, v]) => sympifyNames.includes(k)
      ? `${k}=sympify(${formatPythonValue(String(v))})`
      : `${k}=${formatPythonValue(v)}`)
    .join(', ')
}

/**
 * Whether a runner takes X/y (IV/DV) Variable objects on its factory call.
 * These are declared as parameters with datatype "IV"/"DV" in the component's
 * JSON and collected into `meta.xyParams` by prepareWorkflow.
 *
 * @param {Object} meta - Component metadata.
 * @returns {boolean}
 */
function needsXYVariables(meta) {
  return isSyntheticRunner(meta) && (meta?.xyParams?.length > 0)
}

/**
 * Module-level variable name holding a synthetic runner's built object. Derived
 * from the (already-unique) wrapper `varName` so identical wrappers share it
 * (keeping wrapper-dedup intact) while two distinct synthetic runners get
 * distinct names and never overwrite one another's global `runner`.
 *
 * @param {Object} meta - Component metadata with an assigned `varName`.
 * @returns {string} A unique Python identifier for this runner's object.
 */
function runnerVarName(meta) {
  return `${meta.varName.replace(/_on_state(_\d+)?$/, '$1')}_runner`
}

/**
 * Build the multi-line `<runner> = <name>(...)` call for a runner that takes X/y
 * (IV/DV) Variable objects. The IV/DV literals are emitted verbatim from the
 * component JSON's IV/DV parameter definitions (the node value, or the declared
 * default); a TODO comment prompts the user to adjust the names and ranges.
 * Emitted with a 4-space base indent for use inside a function body.
 *
 * @param {Object} meta - Runner metadata with `pythonName`, `params`, `runParamNames`, `xyParams`, `varName`.
 * @param {number} [baseSpaces=4] - Leading indentation for the first line.
 * @returns {string} The indented, newline-joined call.
 */
function buildXYRunnerCall(meta, baseSpaces = 4) {
  const { pythonName, params, runParamNames = [], xyParams = [], sympifyParams = [] } = meta
  const xyNames = xyParams.map(p => p.name)
  // Regular factory params (sympify-flagged ones are wrapped by
  // buildFactoryParamString); the IV/DV params are emitted verbatim below, under
  // the TODO comment, so exclude them (and the run params) from the ordinary
  // keyword arguments.
  const factory = buildFactoryParamString(params, sympifyParams, [...runParamNames, ...xyNames])
  const pad = ' '.repeat(baseSpaces)
  const pad2 = ' '.repeat(baseSpaces + 4)

  const lines = [`${pad}${runnerVarName(meta)} = ${pythonName}(`]
  if (factory) lines.push(`${pad2}${factory},`)
  lines.push(`${pad2}# TODO: adjust the variable names and ranges below for your experiment`)
  xyParams.forEach(p => {
    // Indent any newlines inside multi-line values so the generated Python stays
    // syntactically valid regardless of how the user formatted the value.
    const indentedValue = String(p.value).replace(/\n/g, `\n${pad2}`)
    lines.push(`${pad2}${p.name}=${indentedValue},`)
  })
  // Close the factory call on the final argument line.
  lines[lines.length - 1] = lines[lines.length - 1].replace(/,$/, ')')
  return lines.join('\n')
}

/**
 * Generate wrapper function for a theorist component
 *
 * @param {CodeBuilder} code - Builder to append the wrapper to.
 * @param {Object} meta - Component metadata with `pythonName`, `params`, `varName`, `nodeName` and `runParamNames` (non-constructor params to exclude from instantiation).
 * @returns {void}
 */
function generateTheoristWrapper(code, meta) {
  const { pythonName, params, varName, nodeName, runParamNames = [] } = meta
  // Only pass constructor (__init__) params to instantiation; params belonging
  // to other methods (e.g. fit) are excluded so instantiation does not fail.
  const paramStr = buildParamString(params, runParamNames)

  code.comment(nodeName)
  code.line(`${varName} = estimator_on_state(${pythonName}(${paramStr}))`)
  code.blank()
}

/**
 * Whether a runner component is a synthetic experiment runner. Synthetic
 * runners (under `autora.experiment_runner.synthetic.*`) return an object that
 * exposes `.variables` and `.run(conditions)`. Real data-collection runners
 * (e.g. firebase) return a plain callable that takes conditions directly and
 * has no `.variables`.
 *
 * @param {Object} meta - Component metadata with `protocolType` and `importPath`.
 * @returns {boolean} True for synthetic experiment runners.
 */
function isSyntheticRunner(meta) {
  return meta?.protocolType === 'experiment_runner' &&
    (meta.importPath || '').includes('.synthetic.')
}

/**
 * Generate wrapper function for an experiment runner component
 *
 * @param {CodeBuilder} code - Builder to append the wrapper to.
 * @param {Object} meta - Component metadata with `pythonName`, `params`, `varName`, `nodeName`, `runParamNames`, `importPath` and `usesFirebaseCredentials`.
 * @returns {void}
 */
function generateRunnerWrapper(code, meta) {
  const { pythonName, params, varName, nodeName, runParamNames = [], sympifyParams = [] } = meta

  code.comment(nodeName)

  if (isSyntheticRunner(meta)) {
    const runParamStr = buildParamString(
      Object.fromEntries(Object.entries(params).filter(([name]) => runParamNames.includes(name)))
    )
    const runArgs = ['conditions=conditions', runParamStr].filter(Boolean).join(', ')
    // Build the runner once at module scope, under a name unique to this wrapper,
    // so distinct synthetic runners don't overwrite one another's global and the
    // variables setup (see generateVariablesSetup) can reuse it — no rebuild.
    const runVar = runnerVarName(meta)
    if (needsXYVariables(meta)) {
      // Synthesize the required X (IVs) and y (DV); the workflow does not carry them.
      code.multiline(buildXYRunnerCall(meta, 0))
    } else {
      const factoryParamStr = buildFactoryParamString(params, sympifyParams, runParamNames)
      code.line(`${runVar} = ${pythonName}(${factoryParamStr})`)
    }
    code.blank()
    code.line('@on_state()')
    code.line(`def ${varName}(conditions: pd.DataFrame) -> Delta:`)
    if (meta.runReturnsDV) {
      // This runner's `.run()` returns the DV values (one per condition), not a
      // full experiment_data frame (e.g. bandit/Q-learning). Assemble one keyed
      // by the runner's own IV/DV variable names.
      code.indent(`dv_values = ${runVar}.run(${runArgs})`)
      // Some runners return a tuple of outputs when extra results are requested
      // (e.g. Q-learning's return_choice_probabilities=True yields
      // (choices, probabilities)). Keep the first element as the DV values; any
      // additional outputs are not part of the declared dependent variable.
      code.indent('if isinstance(dv_values, tuple):')
      code.indent('dv_values = dv_values[0]', 2)
      code.indent('experiment_data = pd.DataFrame({')
      code.indent(`${runVar}.variables.independent_variables[0].name: list(conditions.iloc[:, 0]),`, 2)
      code.indent(`${runVar}.variables.dependent_variables[0].name: dv_values,`, 2)
      code.indent('})')
      code.indent('return Delta(experiment_data=experiment_data)')
    } else {
      code.indent(`return Delta(experiment_data=${runVar}.run(${runArgs}))`)
    }
    code.blank()
    return
  }

  code.line('@on_state()')
  code.line(`def ${varName}(conditions: pd.DataFrame) -> Delta:`)

  if (meta.usesFirebaseCredentials) {
    // Firebase runners require a service-account credentials dict; emit a
    // filled-in template (placeholders for the user to replace) so the runner
    // does not fail on a missing `firebase_credentials`.
    const otherParams = buildParamString(params, ['firebase_credentials'])
    code.indent(`runner = ${pythonName}(`)
    code.indent('# TODO: replace the placeholders below with your own Firebase service-account credentials', 2)
    code.indent('firebase_credentials={', 2)
    code.multiline(TEMPLATES.firebaseCredentialEntries)
    code.indent(otherParams ? `}, ${otherParams})` : '})', 2)
    code.indent('return Delta(experiment_data=runner(conditions))')
  } else {
    // Real runners return a plain callable invoked directly with conditions;
    // every param configures the factory (there is no separate run() call).
    const factoryParamStr = buildParamString(params)
    code.indent(`runner = ${pythonName}(${factoryParamStr})`)
    code.indent('return Delta(experiment_data=runner(conditions))')
  }
  code.blank()
}

/**
 * Find the variable name of a data type spec, descending through list wrappers
 *
 * @param {Object} dataType - Data type spec with a `name` or a nested `variable`.
 * @returns {string|null} The resolved data type name, or null if none.
 */
function dataTypeName(dataType) {
  if (!dataType) return null
  if (dataType.name) return dataType.name
  if (dataType.variable) return dataTypeName(dataType.variable)
  return null
}

/**
 * Generate wrapper function for an experimentalist component (pooler/sampler)
 *
 * @param {CodeBuilder} code - Builder to append the wrapper to.
 * @param {Object} meta - Component metadata with `pythonName`, `params`, `varName`, `nodeName`, `inputDataType` and `outputDataType`.
 * @returns {void}
 */
function generateExperimentalistWrapper(code, meta) {
  const { pythonName, params, varName, nodeName, inputDataType, outputDataType } = meta
  const isPooler = pythonName.includes('pool')
  const isSampler = pythonName.includes('sample') || pythonName.includes('sampler')

  code.comment(nodeName)
  code.line('@on_state()')

  if (inputDataType === null) {
    // Component consumes nothing from state (e.g. bandit random pooler):
    // all of its inputs are plain parameters. Its raw output (nested lists)
    // is wrapped in a DataFrame column so StandardState can coerce it.
    const paramStr = buildParamString(params)
    const call = `${pythonName}(${paramStr})`
    const outName = dataTypeName(outputDataType)
    const value = outName ? `pd.DataFrame({"${outName}": ${call}})` : call
    code.line(`def ${varName}() -> Delta:`)
    code.indent(`return Delta(conditions=${value})`)
  } else if (isPooler) {
    const paramStr = buildParamString(params)
    code.line(`def ${varName}(variables: VariableCollection) -> Delta:`)
    const call = paramStr
      ? `${pythonName}(variables, ${paramStr})`
      : `${pythonName}(variables)`
    code.indent(`return Delta(conditions=${call})`)
  } else if (isSampler) {
    const numSamples = params.num_samples ?? 1
    const otherParamStr = buildParamString(params, ['num_samples'])
    // Some samplers (e.g. the LHS sampler) also require `reference_conditions`,
    // which is not a state field. Derive it from the IV columns of the
    // already-collected experiment_data (empty on the first cycle);
    // experiment_data and variables are auto-injected by @on_state().
    const inputVarNames = Array.isArray(inputDataType?.variables)
      ? inputDataType.variables.map(v => v.name)
      : []
    const needsReference = inputVarNames.includes('reference_conditions')

    if (needsReference) {
      code.line(`def ${varName}(conditions: pd.DataFrame, experiment_data: pd.DataFrame, variables: VariableCollection, num_samples: int = ${numSamples}) -> Delta:`)
      code.indent('reference_conditions = experiment_data[[v.name for v in variables.independent_variables]] if experiment_data is not None else conditions.iloc[0:0]')
      const refCall = otherParamStr
        ? `${pythonName}(conditions=conditions, reference_conditions=reference_conditions, num_samples=num_samples, ${otherParamStr})`
        : `${pythonName}(conditions=conditions, reference_conditions=reference_conditions, num_samples=num_samples)`
      code.indent(`return Delta(conditions=${refCall})`)
    } else {
      code.line(`def ${varName}(conditions: pd.DataFrame, num_samples: int = ${numSamples}) -> Delta:`)
      const call = otherParamStr
        ? `${pythonName}(conditions=conditions, num_samples=num_samples, ${otherParamStr})`
        : `${pythonName}(conditions=conditions, num_samples=num_samples)`
      code.indent(`return Delta(conditions=${call})`)
    }
  } else {
    const paramStr = buildParamString(params)
    code.line(`def ${varName}(conditions: pd.DataFrame) -> Delta:`)
    const call = paramStr
      ? `${pythonName}(conditions, ${paramStr})`
      : `${pythonName}(conditions)`
    code.indent(`return Delta(conditions=${call})`)
  }
  code.blank()
}

/**
 * Dispatch wrapper generation based on the component's protocol type.
 *
 * @param {CodeBuilder} code - Builder to append the wrapper to.
 * @param {Object} meta - Component metadata; `protocolType` selects the wrapper style.
 * @returns {void}
 */
export function generateWrapper(code, meta) {
  if (meta.protocolType === 'theorist') {
    generateTheoristWrapper(code, meta)
  } else if (meta.protocolType === 'experiment_runner') {
    generateRunnerWrapper(code, meta)
  } else {
    generateExperimentalistWrapper(code, meta)
  }
}

/**
 * Collect execution order, imports and per-component metadata from the
 * workflow state. Shared by the Python file and Jupyter notebook generators.
 *
 * @param {Object} state - Editor state with `nodes`, `connections` and `components`.
 * @returns {Object} `{ blocks, imports, componentMeta, derivesVariablesFromRunner, needsEquationVariables }`.
 */
export function prepareWorkflow(state) {
  const { nodes, connections, components } = state
  const allComponents = components ? Object.values(components).flat() : []
  const { blocks } = getExecutionOrder(nodes, connections)

  const allPathNodes = flattenBlockNodes(blocks)
  if (allPathNodes.length === 0) {
    throw new Error('No components found in workflow. Add components and connect them.')
  }

  // Collect imports and component metadata
  const imports = new Map()
  const componentMeta = new Map()

  // Track varName assignments to ensure uniqueness per distinct wrapper signature:
  // - same base varName + same signature → reuse (identical wrapper body, safe to share)
  // - same base varName + different signature → add numeric suffix to disambiguate
  const varNameToSignature = new Map()  // varName → JSON-serialized signature
  const varNameSuffix = new Map()       // base varName → next available numeric suffix

  allPathNodes.forEach(node => {
    const protocol = allComponents.find(c => c.uuid === node.protocolUuid)
    if (!protocol) {
      console.warn(`Protocol not found for node: ${node.name}`)
      return
    }

    const { importPath, pythonName, protocolType, file, inputDataType } = protocol

    // Experimentalist pythonNames are bare function names (pool, sample, ...)
    // shared across modules, so import them aliased to the JSON file name
    let alias = null
    if (protocolType === 'experimentalist' && file) {
      const fileAlias = toPythonName(file.replace(/\.json$/, ''))
      if (fileAlias && fileAlias !== pythonName) alias = fileAlias
    }

    if (!imports.has(importPath)) imports.set(importPath, new Set())
    imports.get(importPath).add(alias ? `${pythonName} as ${alias}` : pythonName)

    // Parameters are grouped by function in the JSON. The instantiation group
    // configures the constructor/factory: theorist classes use an "__init__"
    // group, runner factories use a group named after the function (pythonName).
    // Any other group (e.g. "fit" for a theorist, "run" for a runner) is applied
    // elsewhere, not at instantiation, so its params must be excluded from it.
    const initGroup = protocol.parameters?.[pythonName] || protocol.parameters?.['__init__'] || []
    const factoryParamNames = new Set(initGroup.map(p => p.name))
    const declaredParamNames = Object.values(protocol.parameters || {})
      .flatMap(groupParams => (Array.isArray(groupParams) ? groupParams : []).map(p => p.name))
    const runParamNames = declaredParamNames.filter(name => !factoryParamNames.has(name))
    // Runners that require Firebase service-account credentials get a filled-in
    // credentials template emitted automatically (see generateRunnerWrapper).
    const usesFirebaseCredentials = declaredParamNames.includes('firebase_credentials')

    // IV/DV factory arguments (datatype "IV"/"DV" in the JSON, e.g. a synthetic
    // runner's X and y). Their literal is taken from the node value if set,
    // otherwise the declared default, and emitted verbatim (see buildXYRunnerCall).
    const xyParams = initGroup
      .filter(p => p.datatype === 'IV' || p.datatype === 'DV')
      .map(p => ({
        name: p.name,
        datatype: p.datatype,
        value: String(String((node.parameters || {})[p.name] ?? '').trim() || (p.default ?? ''))
      }))

    // Factory params declared `"sympify": true` in the JSON: their string value
    // is a SymPy expression and is wrapped in `sympify(...)` when emitted (see
    // buildFactoryParamString). Pull in the `sympify` import only when such a
    // param is actually set to a non-blank value (a blank one is omitted, so the
    // import would be unused).
    const sympifyParams = initGroup.filter(p => p.sympify === true).map(p => p.name)
    if (sympifyParams.some(name => !isBlankString((node.parameters || {})[name] ?? ''))) {
      if (!imports.has(SYMPIFY_IMPORT.module)) imports.set(SYMPIFY_IMPORT.module, new Set())
      imports.get(SYMPIFY_IMPORT.module).add(SYMPIFY_IMPORT.name)
    }

    componentMeta.set(node.id, {
      importPath,
      // Name as imported into the generated file (alias when aliased)
      pythonName: alias || pythonName,
      inputDataType,
      outputDataType: protocol.outputDataType,
      runParamNames,
      xyParams,
      sympifyParams,
      // Runner whose `.run()` returns the DV values rather than a full
      // experiment_data frame (see generateRunnerWrapper).
      runReturnsDV: protocol.runReturnsDV === true,
      usesFirebaseCredentials,
      protocolType,
      params: node.parameters || {},
      varName: null, // assigned below
      nodeName: node.name
    })

    // Assign varName: reuse when base name and signature both match; add a
    // numeric suffix when the same base varName is claimed with a different
    // wrapper signature so the two generated `def` blocks don't collide.
    const meta = componentMeta.get(node.id)
    const signature = JSON.stringify({ pythonName: meta.pythonName, params: meta.params })
    const baseName = `${toPythonName(node.name)}_on_state`

    if (varNameToSignature.has(baseName) && varNameToSignature.get(baseName) === signature) {
      // Identical signature for this base name: safe to reuse the same wrapper.
      meta.varName = baseName
    } else if (!varNameToSignature.has(baseName)) {
      // First time this base name is seen: claim it.
      varNameToSignature.set(baseName, signature)
      meta.varName = baseName
    } else {
      // Same base name but different signature: find or create a suffixed variant.
      let suffix = varNameSuffix.get(baseName) || 1
      let candidate
      do {
        candidate = `${baseName}_${suffix}`
        suffix++
      } while (varNameToSignature.has(candidate) && varNameToSignature.get(candidate) !== signature)
      varNameSuffix.set(baseName, suffix)
      varNameToSignature.set(candidate, signature)
      meta.varName = candidate
    }
  })

  // Only synthetic runners provide `.variables`; with such a runner the
  // variables are derived from it, otherwise placeholder variables are emitted.
  const derivesVariablesFromRunner = [...componentMeta.values()].some(isSyntheticRunner)
  // Runners taking IV/DV factory arguments require IV/DV and numpy imports.
  const needsEquationVariables = [...componentMeta.values()].some(needsXYVariables)

  return { blocks, imports, componentMeta, derivesVariablesFromRunner, needsEquationVariables }
}

/**
 * Build the variables-initialization block (indented for a function body).
 * When the workflow contains a synthetic experiment runner, the variables are
 * taken from the `runner` already built in that runner's component definition
 * (see generateRunnerWrapper) rather than rebuilding it; otherwise fall back to
 * the placeholder template (real runners such as firebase have no `.variables`).
 *
 * @param {Map} componentMeta - Map of node id to component metadata.
 * @param {Object[]} orderedNodes - Nodes in execution order, used to find any runner.
 * @returns {string} Indented Python source for the variables setup block.
 */
export function generateVariablesSetup(componentMeta, orderedNodes) {
  const runnerMeta = orderedNodes
    .map(node => componentMeta.get(node.id))
    .find(isSyntheticRunner)

  if (!runnerMeta) return TEMPLATES.defaultVariables

  // Reuse this runner's object defined in the component section — do not rebuild
  // it. Use its unique variable name so we bind the intended runner even when the
  // workflow defines several distinct synthetic runners.
  const runVar = runnerVarName(runnerMeta)
  return [
    '    # Variables are governed by the experiment runner defined above',
    `    assert ${runVar}.variables is not None`,
    `    variables = ${runVar}.variables`
  ].join('\n')
}

/**
 * Emit the import block (standard + component + data imports) into a builder.
 * `Variable` and numpy are only needed by the placeholder variables template,
 * so they are skipped when the variables come from an experiment runner.
 *
 * @param {CodeBuilder} code - Builder to append the imports to.
 * @param {Map} imports - Map of import path to a Set of imported (possibly aliased) names.
 * @param {Object} [options] - Options object.
 * @param {boolean} [options.usesPlaceholderVariables=true] - Whether to also import `Variable` and numpy for the placeholder template.
 * @param {boolean} [options.usesEquationVariables=false] - Whether to also import `IV`, `DV` and numpy for a synthesized equation runner X/y.
 * @returns {void}
 */
export function generateImports(code, imports, { usesPlaceholderVariables = true, usesEquationVariables = false } = {}) {
  code.multiline(TEMPLATES.standardImports)
  const variableNames = ['VariableCollection']
  if (usesPlaceholderVariables) variableNames.push('Variable')
  if (usesEquationVariables) variableNames.push('IV', 'DV')
  code.line(`from autora.variable import ${variableNames.join(', ')}`)

  imports.forEach((names, importPath) => {
    code.line(`from ${importPath} import ${Array.from(names).join(', ')}`)
  })

  code.blank()
  code.multiline(TEMPLATES.dataImports)
  if (usesPlaceholderVariables || usesEquationVariables) code.line('import numpy as np')
}

/**
 * Generate Python code from workflow state
 *
 * @param {Object} state - Editor state with `nodes`, `connections` and `components`.
 * @returns {string} A complete, runnable Python script as a string.
 */
export function generatePythonCode(state) {
  const { blocks, imports, componentMeta, derivesVariablesFromRunner, needsEquationVariables } = prepareWorkflow(state)

  // Build Python code
  const code = new CodeBuilder()

  // Header and imports
  code.multiline(TEMPLATES.header(new Date().toISOString()))
  code.blank()
  generateImports(code, imports, { usesPlaceholderVariables: !derivesVariablesFromRunner, usesEquationVariables: needsEquationVariables })
  code.blank().blank()

  // Generate wrapper functions. Components that produce an identical definition
  // (same function name and parameters) are emitted only once and reused by
  // every call site, so no duplicate `def`s appear.
  const seenWrappers = new Set()
  componentMeta.forEach((meta) => {
    const wrapper = new CodeBuilder()
    generateWrapper(wrapper, meta)
    const text = wrapper.toString()
    if (seenWrappers.has(text)) return
    seenWrappers.add(text)
    code.multiline(text)
  })

  code.blank()

  // Generate main function
  code.line('def main():')
  code.multiline(generateVariablesSetup(componentMeta, flattenBlockNodes(blocks)))
  code.indent('')
  code.multiline(TEMPLATES.initState)
  code.indent('')

  // Emit a single `state = fn(state[, num_samples=…])` call at the given level.
  const addComponentCall = (node, level) => {
    const meta = componentMeta.get(node.id)
    if (!meta) return

    const { varName, nodeName, pythonName } = meta
    const isSampler = pythonName.includes('sample') || pythonName.includes('sampler')

    code.indent(`# ${nodeName}`, level)
    if (isSampler && node.parameters?.num_samples != null) {
      code.indent(`state = ${varName}(state, num_samples=${node.parameters.num_samples})`, level)
    } else {
      code.indent(`state = ${varName}(state)`, level)
    }
    code.indent('', level)
  }

  // Emit the block tree in order: `once` blocks run their nodes a single time,
  // `loop` blocks wrap their children in a for-loop and recurse (nested loops
  // become nested for-loops). A loop prints its cycle only when it directly runs
  // components; a loop that only holds nested loops emits no cycle print.
  const renderBlocks = (blks, level, loopDepth = 0) => {
    blks.forEach(block => {
      if (block.type === 'loop') {
        const loopVar = `cycle_${loopDepth}`
        code.indent(`# Experiment loop (${block.maxCounter} cycles)`, level)
        code.indent(`for ${loopVar} in range(${block.maxCounter}):`, level)
        if (block.children.some(c => c.type === 'once')) {
          code.indent(`print(f'Cycle {${loopVar}}')`, level + 1)
          code.indent('', level + 1)
        }
        renderBlocks(block.children, level + 1, loopDepth + 1)
      } else {
        block.nodes.forEach(node => addComponentCall(node, level))
      }
    })
  }
  renderBlocks(blocks, 1)

  code.indent('')
  code.indent('print("Workflow completed!")')
  code.indent('return state')
  code.blank().blank()
  code.line("if __name__ == '__main__':")
  code.indent('main()')
  code.blank()

  return code.toString()
}

/**
 * Collect the set of pip packages required by the workflow's components.
 *
 * @param {Object} state - Editor state with `nodes` and `components`.
 * @returns {string[]} Deduplicated list of pip package specifiers.
 */
export function collectPipPackages(state) {
  const { nodes, components } = state
  const allComponents = components ? Object.values(components).flat() : []
  const pipPackages = new Set()

  nodes.forEach(node => {
    if (CONTROL_NODE_TYPES.includes(node.type)) return
    const protocol = allComponents.find(c => c.uuid === node.protocolUuid)
    if (protocol?.pipInstall) pipPackages.add(protocol.pipInstall)
  })

  return Array.from(pipPackages)
}

/**
 * Generate pip install commands for all required packages
 *
 * @param {Object} state - Editor state with `nodes` and `components`.
 * @returns {string} A `pip install ...` command, or a comment when none are required.
 */
export function generatePipInstalls(state) {
  const pipPackages = collectPipPackages(state)

  if (pipPackages.length === 0) return '# No additional packages required'
  return `pip install ${pipPackages.join(' ')}`
}