import multiprocessing as mp
from queue import Queue

{% for stage_num in range(1, stage_count+1) %}
def stage{{stage_num}}_function(in_q, out_q, batch_size):
    """Stage {{stage_num}} with temporal batching."""
    while True:
        batch = []
        for _ in range(batch_size):
            item = in_q.get()
            if item is None:
                break
            batch.append(item)

        if not batch and item is None:
            out_q.put(None)
            break

        # Process each batch_item
        for batch_item in batch:
            result = {{ stage_exprs[stage_num-1] | replace('item','batch_item') }}
            out_q.put(result)
        {% endfor %}

def run_pipeline_tdp(input_data, batch_size=None):
    """Run the pipeline with temporal domain partitioning."""
    if batch_size is None:
        batch_size = max(1, len(input_data) // (mp.cpu_count() * 2))

    queues = [mp.Queue() for _ in range({{stage_count}} + 1)]
    procs = []

    {% for stage_num in range(1, stage_count+1) %}
    p = mp.Process(
        target=stage{{stage_num}}_function,
        args=(queues[{{stage_num-1}}], queues[{{stage_num}}], batch_size)
    )
    p.start()
    procs.append(p)
    {% endfor %}

    for item in input_data:
        queues[0].put(item)
    queues[0].put(None)

    results = []
    while True:
        out = queues[-1].get()
        if out is None:
            break
        results.append(out)

    for p in procs:
        p.join()
    return results

if __name__ == '__main__':
    data = {{input_data}}
    print(run_pipeline_tdp(data))
