# noqa: E501
import time
import advchksys
from concurrent.futures import ThreadPoolExecutor


def parallel_benchmark_2d(manager, num_chunks, chunk_size, max_workers=None):
    print(
        f"\n[2D] Parallel Benchmark: {num_chunks} chunks of size {chunk_size}x{chunk_size} (max_workers={max_workers})"
    )
    t0 = time.perf_counter()

    def create_and_fill(i):
        x, y = i % 100, i // 100
        chunk = manager.LoadOrCreateChunk(x, y, chunk_size, chunk_size)
        for cx in range(chunk_size):
            for cy in range(chunk_size):
                chunk[cx, cy] = (cx + cy) % 256

    # Parallel chunk creation and filling
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        list(executor.map(create_and_fill, range(num_chunks)))

    t1 = time.perf_counter()
    print(
        f"  Created and filled {num_chunks} chunks in {t1-t0:.3f} seconds (parallel)"
    )

    # Parallel chunk access (sum all cells)
    t2 = time.perf_counter()

    def sum_chunk(i):
        x, y = i % 100, i // 100
        chunk = manager.GetChunk(x, y)
        subtotal = 0
        if chunk is not None:
            for cx in range(chunk_size):
                for cy in range(chunk_size):
                    subtotal += chunk[cx, cy]
        return subtotal

    total = 0
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        for subtotal in executor.map(sum_chunk, range(num_chunks)):
            total += subtotal

    t3 = time.perf_counter()
    print(
        f"  Accessed {num_chunks} chunks in {t3-t2:.3f} seconds (sum: {total}, parallel)"
    )

    # Parallel chunk unload
    t4 = time.perf_counter()

    def unload_chunk(i):
        x, y = i % 100, i // 100
        manager.UnloadChunk(x, y)

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        list(executor.map(unload_chunk, range(num_chunks)))

    t5 = time.perf_counter()
    print(f"  Unloaded {num_chunks} chunks in {t5-t4:.3f} seconds (parallel)")


def parallel_benchmark_3d(
    manager,
    num_chunks,
    chunk_size_x,
    chunk_size_y,
    chunk_size_z,
    max_workers=None,
):
    print(
        f"\n[3D] Parallel Benchmark: {num_chunks} chunks of size {chunk_size_x}x{chunk_size_y}x{chunk_size_z} (max_workers={max_workers})"
    )
    t0 = time.perf_counter()

    def create_and_fill(i):
        x, y, z = i % 10, (i // 10) % 10, i // 100
        chunk = manager.LoadOrCreateChunk(
            x, y, z, chunk_size_x, chunk_size_y, chunk_size_z
        )
        for cx in range(chunk_size_x):
            for cy in range(chunk_size_y):
                for cz in range(chunk_size_z):
                    chunk[cx, cy, cz] = (cx + cy + cz) % 256

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        list(executor.map(create_and_fill, range(num_chunks)))

    t1 = time.perf_counter()
    print(
        f"  Created and filled {num_chunks} chunks in {t1-t0:.3f} seconds (parallel)"
    )

    # Parallel chunk access (sum all cells)
    t2 = time.perf_counter()

    def sum_chunk(i):
        x, y, z = i % 10, (i // 10) % 10, i // 100
        chunk = manager.GetChunk(x, y, z)
        subtotal = 0
        if chunk is not None:
            for cx in range(chunk_size_x):
                for cy in range(chunk_size_y):
                    for cz in range(chunk_size_z):
                        subtotal += chunk[cx, cy, cz]
        return subtotal

    total = 0
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        for subtotal in executor.map(sum_chunk, range(num_chunks)):
            total += subtotal

    t3 = time.perf_counter()
    print(
        f"  Accessed {num_chunks} chunks in {t3-t2:.3f} seconds (sum: {total}, parallel)"
    )

    # Parallel chunk unload
    t4 = time.perf_counter()

    def unload_chunk(i):
        x, y, z = i % 10, (i // 10) % 10, i // 100
        manager.UnloadChunk(x, y, z)

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        list(executor.map(unload_chunk, range(num_chunks)))

    t5 = time.perf_counter()
    print(f"  Unloaded {num_chunks} chunks in {t5-t4:.3f} seconds (parallel)")


def main():
    print("AdvChkSys Python Parallel Benchmark")
    print("Library version:", advchksys.get_version())

    # 2D Parallel Benchmark
    constraints2d = advchksys.create_constraints(
        min_x=0, max_x=99, min_y=0, max_y=200, max_loaded=10000000
    )
    manager2d = advchksys.create_2d_manager(constraints2d)
    parallel_benchmark_2d(
        manager2d, num_chunks=20000, chunk_size=32, max_workers=4
    )

    # 3D Parallel Benchmark
    constraints3d = advchksys.create_constraints(
        min_x=0, max_x=9, min_y=0, max_y=9, max_loaded=10000
    )
    manager3d = advchksys.create_3d_manager(constraints3d)
    parallel_benchmark_3d(
        manager3d,
        num_chunks=100,
        chunk_size_x=16,
        chunk_size_y=16,
        chunk_size_z=32,
        max_workers=4,
    )


if __name__ == "__main__":
    main()
