#!/usr/bin/env python3
"""Automatic differentiation demo for Parallel Python.

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

Uses class methods as parallel functions (available since pp 1.4).  The
program calculates the partial sums of f(x) = x - x**2/2 + x**3/3 - x**4/4 + ...
and first derivatives f'(x) using automatic differentiation.
In the limit f(x) = ln(x+1) and f'(x) = 1/(x+1).

Usage: python auto_diff.py [ncpus]
"""

import sys

import pp


# Partial implementation of an automatic differentiation class
class AD:
    def __init__(self, x, dx=0.0):
        self.x = float(x)
        self.dx = float(dx)

    def __pow__(self, val):
        if isinstance(val, int):
            return AD(self.x ** val, val * self.x ** (val - 1) * self.dx)
        raise TypeError("Second argument must be an integer")

    def __add__(self, val):
        if isinstance(val, AD):
            return AD(self.x + val.x, self.dx + val.dx)
        return AD(self.x + val, self.dx)

    def __radd__(self, val):
        return self + val

    def __mul__(self, val):
        if isinstance(val, AD):
            return AD(self.x * val.x, self.x * val.dx + val.x * self.dx)
        return AD(self.x * val, val * self.dx)

    def __rmul__(self, val):
        return self * val

    def __truediv__(self, val):
        if isinstance(val, AD):
            return self * AD(1 / val.x, -val.dx / val.x ** 2)
        return self * (1 / float(val))

    def __rtruediv__(self, val):
        return AD(val) / self

    def __sub__(self, val):
        if isinstance(val, AD):
            return AD(self.x - val.x, self.dx - val.dx)
        return AD(self.x - val, self.dx)

    def __repr__(self):
        return str((self.x, self.dx))


class PartialSum:
    """Contains methods which will be executed in parallel."""

    def __init__(self, n):
        self.n = n

    def t_log(self, x):
        """Truncated natural logarithm."""
        return self.partial_sum(x - 1)

    def partial_sum(self, x):
        """Partial sum for truncated natural logarithm."""
        return sum(float(i % 2 and 1 or -1) * x ** i / i for i in range(1, self.n))


def main():
    print("Usage: python auto_diff.py [ncpus]")
    print("    [ncpus] - the number of workers to run in parallel,")
    print("    if omitted it will be set to the number of processors in the system")

    # tuple of all parallel python servers to connect with
    # ppservers = ("*",)          # auto-discover
    # ppservers = ("10.0.0.1",)  # static IP
    ppservers = ()

    if len(sys.argv) > 1:
        ncpus = int(sys.argv[1])
        job_server = pp.Server(ncpus, ppservers=ppservers)
    else:
        job_server = pp.Server(ppservers=ppservers)

    print("Starting pp with", job_server.get_ncpus(), "workers")

    proc = PartialSum(20000)

    results = []
    for i in range(32):
        # Creates an object with x = float(i)/32 + 1 and dx = 1.0
        ad_x = AD(float(i) / 32 + 1, 1.0)
        # Submits a job of calculating proc.t_log(x).
        f = job_server.submit(proc.t_log, (ad_x,))
        results.append((ad_x.x, f))

    for x, f in results:
        # Retrieves the result of the calculation
        val = f()
        print("t_log(%.6f) = %.6f, t_log'(%f) = %f" % (x, val.x, x, val.dx))

    # Print execution statistics
    job_server.print_stats()
    job_server.wait()
    job_server.destroy()


if __name__ == "__main__":
    main()