import aner.parallel

// Replace this policy with parallel.cluster(config: "private.repeat.cluster")
// to run both stages on enrolled workers. Each stage has separate worker sessions.
// Same-host workers can map a shared-memory-backed history directly in stage two.
let execution = parallel.local(workers: 1)

fn simulate(parameters: List, iteration: Int64, seed: Int64,
            output: OutputRegion<Float32>) -> Unit {
    let input = cast<FrozenArray<Float32>>(parameters.get(0))
    var state = input[iteration % input.len()]
    for step in Range(8) {
        state += Float32(0.25)
        output[step, 0] = state
        output[step, 1] = Float32(iteration)
    }
}

fn final_state(history: OutputResults<Float32>, iteration: Int64) -> Float32 {
    return history[iteration, 7, 0]
}

fn summarize(parameters: List, iteration: Int64, seed: Int64,
             output: OutputRegion<Float32>) -> Unit {
    let history = cast<OutputResults<Float32>>(parameters.get(0))
    output[0] = final_state(history, iteration)
    output[1] = history[iteration, 7, 1]
}

let input = Array<Float32>(4, Float32(1.0)).snapshot()
let history = parallel.repeat(iterations: 1000, task: simulate,
    parameters: List(input), output_shape: List(8, 2),
    execution: execution, seed: 2026)

// Submit the immutable result directly. Local tasks retain its original native
// buffer. Compatible coordinator-host cluster workers map shared history directly,
// without restaging or numerical transport; other hosts receive their own backing.
let summary = parallel.repeat(iterations: 1000, task: summarize,
    parameters: List(history), output_shape: List(2),
    execution: execution, seed: 2026)

print(summary.shape())          // List(1000, 2)
print(summary[999, 0])          // 3.0
print(summary[999, 1])          // 999.0
print(history[999, 7, 0])       // 3.0: the first result is still readable
print(summary.summary())       // Shared input payload and full retained storage