#!/usr/bin/env python
# coding: utf-8

# DistArray Julia Set
# ===================
# 
# The Julia set, for a given complex number $c$, is the set of points $z$
# such that $|z_{i}|$ remains bounded where $z_{i+1} = z_{i}^2 + c$.
# 
# This can be plotted by counting how many iterations are required for $|z_{i}|$ to exceed a cutoff.
# 
# Depending on the value of $c$, the Julia set may be connected and contain
# a lot of points, or it could be disconnected and contain fewer points.
# The points in the set will require the maximum iteration count, so
# the connected sets will usually take longer to compute.

# First, some imports.

# In[1]:


from __future__ import print_function

import numpy
import matplotlib.pyplot as plt

# DistArray imports
from distarray.globalapi import Context, Distribution
from distarray.globalapi.distarray import DistArray

# inline plots
get_ipython().run_line_magic('matplotlib', 'inline')

# bigger figures
plt.rcParams.update({'figure.figsize': (10, 10)})


# Julia set kernels
# -----------------
# To avoid round-trips between the client and engines, our strategy will be to push a function to each engine to compute that engine's local section of the Julia set.  Here we have two options for this "kernel": a version that avoids NumPy fancy indexing, and one that uses it.

# In[2]:


def numpy_julia_calc(z, c, z_max, n_max):
    """Calculate the Julia set using NumPy.

    Parameters
    ----------
    z : NumPy array
        array of complex values whose iterations we will count.
    c : complex
        Complex number to add at each iteration.
    z_max : float
        Magnitude of complex value that we assume goes to infinity.
    n_max : int
        Maximum number of iterations.
    """
    z = numpy.asarray(z)
    counts = numpy.zeros_like(z, dtype=numpy.int32)
    hits = numpy.zeros_like(z, dtype=numpy.bool)
    mask = numpy.zeros_like(z, dtype=numpy.bool)
    n = 0

    while not numpy.all(hits) and n < n_max:
        z = z * z + c
        mask = (abs(z) > z_max) & (~hits)
        counts[mask] = n
        hits |= mask
        z[hits] = 0
        n += 1
    counts[~hits] = n_max
    return counts


def fancy_numpy_julia_calc(z, c, z_max, n_max):
    """Calculate the Julia set using NumPy fancy indexing.

    Parameters
    ----------
    z : NumPy array
        array of complex values whose iterations we will count.
    c : complex
        Complex number to add at each iteration.
    z_max : float
        Magnitude of complex value that we assume goes to infinity.
    n_max : int
        Maximum number of iterations.
    """
    z = numpy.asarray(z)
    counts = numpy.zeros_like(z, dtype=numpy.int32)
    hits = numpy.zeros_like(z, dtype=numpy.bool)
    mask = numpy.zeros_like(z, dtype=numpy.bool)
    n = 0

    while not numpy.all(hits) and n < n_max:
        z[~hits] = z[~hits] * z[~hits] + c
        mask = (abs(z) > z_max) & (~hits)
        counts[mask] = n
        hits |= mask
        n += 1
    counts[~hits] = n_max
    return counts


# Coordinating functions
# ----------------------
# Here we have functions that create a DistArray representing the complex plane and functions that coordinate applying the kernel on each distributed section of the DistArray.

# In[3]:


def create_complex_plane(context, resolution, dist, re_ax, im_ax):
    """Create a DistArray containing points on the complex plane.

    Parameters
    ----------
    context : DistArray Context
    resolution : 2-tuple
        The number of points along Re and Im axes.
    dist : 2-element sequence or dict
        dist_type for of the DistArray Distribution.
    re_ax : 2-tuple
        The (lower, upper) range of the Real axis.
    im_ax : 2-tuple
        The (lower, upper) range of the Imaginary axis.
    """

    import numpy as np

    def fill_complex_plane(arr, re_ax, im_ax, resolution):
        """Fill in points on the complex coordinate plane."""
        re_step = float(re_ax[1] - re_ax[0]) / resolution[0]
        im_step = float(im_ax[1] - im_ax[0]) / resolution[1]
        for i in arr.distribution[0].global_iter:
            for j in arr.distribution[1].global_iter:
                arr.global_index[i, j] = complex(re_ax[0] + re_step * i,
                                                 im_ax[0] + im_step * j)

    # Create an empty distributed array.
    distribution = Distribution(context, (resolution[0], resolution[1]),
                                dist=dist)
    complex_plane = context.empty(distribution, dtype=np.complex64)
    context.apply(fill_complex_plane,
                  (complex_plane.key, re_ax, im_ax, resolution))
    return complex_plane


# In[4]:


def local_julia_calc(la, c, z_max, n_max, kernel):
    """Calculate the number of iterations for the point to escape.

    Parameters
    ----------
    la : LocalArray
        LocalArray of complex values whose iterations we will count.
    c : complex
        Complex number to add at each iteration.
    z_max : float
        Magnitude of complex value that we assume goes to infinity.
    n_max : int
        Maximum number of iterations.
    kernel : function
        Kernel to use for computation of the Julia set.  Options are 'fancy',
        'numpy', or 'cython'.
    """
    from distarray.localapi import LocalArray
    counts = kernel(la, c, z_max, n_max)
    res = LocalArray(la.distribution, buf=counts)
    return proxyize(res)  # noqa


def distributed_julia_calc(distarray, c, z_max, n_max,
                           kernel=fancy_numpy_julia_calc):
    """Calculate the Julia set for an array of points in the complex plane.

    Parameters
    ----------
    distarray : DistArray
        DistArray of complex values whose iterations we will count.
    c : complex
        Complex number to add at each iteration.
    z_max : float
        Magnitude of complex value that we assume goes to infinity.
    n_max : int
        Maximum number of iterations.
    kernel: function
        Kernel to use for computation of the Julia set.  Options are 'fancy',
        'numpy', or 'cython'.
    """
    context = distarray.context
    iters_key = context.apply(local_julia_calc,
                              (distarray.key, c, z_max, n_max),
                              {'kernel': kernel})
    iters_da = DistArray.from_localarrays(iters_key[0], context=context,
                                          dtype=numpy.int32)
    return iters_da


# Using DistArray to explore the Julia set
# ----------------------------------------

# In[5]:


context = Context()  # handles communicating with the engines


# In[6]:


# create the complex plane
resolution = (512, 512)
dist = 'cc'  # different distributions of the data among the engines
             # one character per dimension
             # 'c' stands for 'cyclic', 'b' for block
re_ax = (-1.5, 1.5)
im_ax = (-1.5, 1.5)
complex_plane = create_complex_plane(context, resolution, dist, re_ax, im_ax)

# calculate the Julia set
z_max = 2.0
n_max = 100
c = complex(0.285, 0.01)
iters_da = distributed_julia_calc(complex_plane, c, z_max, n_max, kernel=numpy_julia_calc)

# plot it!
plt.matshow(iters_da.toarray(), cmap='hot')


# In[7]:


context.close()