"""
Test parallel performance with vector operations using shared memory
"""
from cpu_parallel_driver import CPUDriver
import numpy as np
import multiprocessing as mp

def create_test_program(size=1000):
    # Create vector data in shared memory
    shared_vectors = mp.RawArray('f', size * 256)
    vectors_np = np.frombuffer(shared_vectors, dtype=np.float32).reshape(size, 256)
    vectors_np[:] = np.random.rand(size, 256)
    
    program = []
    # Create parallel vector operations
    for i in range(size):
        # Distribute across cores and threads
        core_id = i % mp.cpu_count()
        thread_id = (i // mp.cpu_count()) % 2
        
        program.append({
            'type': 'VADD',
            'core_id': core_id,
            'thread_id': thread_id,
            'vector_offset': i * 256,
            'dest': i % 1024,  # Keep result indices within shared memory bounds
            'size': 256
        })
        
    return program, shared_vectors

def main():
    # Initialize CPU driver with all available cores
    driver = CPUDriver(num_cores=mp.cpu_count())
    print(f"\nInitialized CPU driver with {mp.cpu_count()} cores")
    
    # Create test program
    size = 10000  # 10K vector operations
    program, shared_vectors = create_test_program(size)
    print(f"Created {size} parallel vector operations")
    
    # Run parallel execution
    print("\nStarting parallel execution...")
    start = time.time()
    
    # Execute all operations in parallel
    results = driver.execute_batch(program)
    
    # Wait for all operations
    for result in results:
        result.wait()
        
    end = time.time()
    duration = end - start
    
    # Calculate performance metrics
    ops_per_sec = size / duration
    ops_per_core = ops_per_sec / mp.cpu_count()
    
    print(f"\nExecution Results:")
    print(f"- Total operations: {size}")
    print(f"- Execution time: {duration:.3f} seconds")
    print(f"- Operations/second: {ops_per_sec:.2f}")
    print(f"- Operations/core/second: {ops_per_core:.2f}")
    
if __name__ == '__main__':
    import time
    main()