#!/usr/bin/env python3
"""Demonstrate dynamic worker scaling with pp.

Website: https://www.parallelpython.com
License: Apache-2.0 (see LICENSE and NOTICE)

Calculates the partial alternating-harmonic sum
1 - 1/2 + 1/3 - 1/4 + ... across subtasks, re-sizing the worker pool
on each iteration to show how results vary with pool size.

Usage:
    python dynamic_ncpus.py
"""

import math
import time

import pp


def part_sum(start, end):
    """Calculate a partial alternating harmonic sum over [start, end)."""
    s = 0.0
    for x in range(start, end):
        if x % 2 == 0:
            s -= 1.0 / x
        else:
            s += 1.0 / x
    return s


def main():
    print("Usage: python dynamic_ncpus.py")
    print()

    start = 1
    end = 2000000

    parts = 64
    step = (end - start) // parts + 1

    # Create jobserver
    job_server = pp.Server()

    # Run the same task with different numbers of active workers
    for ncpus in (1, 2, 4, 8, 16, 1):
        job_server.set_ncpus(ncpus)
        jobs = []
        start_time = time.time()
        print("Starting ", job_server.get_ncpus(), " workers")
        for index in range(parts):
            starti = start + index * step
            endi = min(start + (index + 1) * step, end)
            jobs.append(job_server.submit(part_sum, (starti, endi)))

        total = sum(job() for job in jobs)
        print(f"Partial sum is {total:.6f} | diff = {math.log(2) - total:.6e}")
        print(f"Time elapsed: {time.time() - start_time:.3f}s\n")

    job_server.print_stats()
    job_server.destroy()


if __name__ == "__main__":
    main()
