import multiprocessing as mp
from queue import Queue
import threading

{% for stage_num in range(1, stage_count+1) %}
def stage{{stage_num}}_function(in_q, out_q, num_workers):
    """Stage {{stage_num}} with multiple workers."""
    def worker():
        while True:
            item = in_q.get()
            if item is None:
                in_q.put(None)
                break
{% for line in stage_bodies[stage_num-1] %}
            {{ line }}
{% endfor %}
            out_q.put(item)

    # Spawn worker threads
    threads = []
    for _ in range(num_workers):
        t = threading.Thread(target=worker)
        t.daemon = True
        t.start()
        threads.append(t)

    for t in threads:
        t.join()
    out_q.put(None)
{% endfor %}

def run_pipeline_tip(input_data, stage_workers=None):
    """Run the pipeline with temporal instruction partitioning."""
    if stage_workers is None:
        stage_workers = [mp.cpu_count() for _ in range({{stage_count}})]

    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}}], stage_workers[{{stage_num-1}}])
    )
    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_tip(data))
