import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import random

class JuliaSet:

    def __init__(self):
        """
        Constructor of the JuliaSet class

        :param size: size in pixels (for both width and height)
        :param dpi: dots per inch (default 300)
        """

        # Initialize image related parameters
        self.size = 256
        self.dpi = 300
        self.norm = True
        self.mirror = False

        # Initialize process related parameters
        self.escrad = 3
        self.niter = 250

    def param(self, **kwargs):
        """
        Get parameters from input dictionary and set attributes.

        :param kwargs: a dictionary in the form 
            `{'arg1':value, ..., 'argN': value}`
        """

        # Check if kwargs in not empty
        if kwargs is not None:
            # Image related parameters
            if 'size' in kwargs:
                self.size = kwargs.pop('size', 256)
            if 'dpi' in kwargs:
                self.dpi = kwargs.pop('dpi', 300)
            if 'norm' in kwargs:
                self.norm = kwargs.pop('norm', True)
            if 'mirror' in kwargs:
                self.mirror = kwargs.pop('mirror', False)

            # Process related parameters
            if 'escrad' in kwargs:
                self.escrad = kwargs.pop('escrad', 3)
            if 'niter' in kwargs:
                self.niter = kwargs.pop('niter', 250)

        # If kwargs is not empty there is some invalid keywords
        if kwargs:
            print("{} are invalid keyword arguments!".format(kwargs.keys()))

    def run(self, show=False, fname='juilaset-output'):
        """
        Run the Julia set generator

        :param mirror: if True the julia is mirrored horizontally and
            vertically; each mirror is concatenate with the original
            to produce a new image
        :param norm: if true the Julia set is normalized by its
            absolute maximum value.
        :param show: if show is `False` th eoutput image will be
            written as a PNG file named `fname`
        :param fname: Name of the output PNG file to write on disk
        """

        # Get a complex value among a list of best Julia sets 
        cpxNum = self.getComplexValue()

        # Get the target area
        # For more randomness, the target area is a random 
        # subset of a wide one defined with x[-1.5, 1.5] and
        # y[-1.5, 1.5]
        xrng, yrng = self.getTargetArea()

        # Process
        julia = self.processJulia(cpxNum, xrng, yrng)

        # Normalization
        if(self.norm):
            julia /= np.amax(np.abs(julia))

        # Mirroring
        if(self.mirror):
            # Horizontal mirroring and concatenate
            juliamirror = np.flip(julia, axis=1)
            julia = np.concatenate((julia, juliamirror), axis=1)
            # Vertical mirroring and concatenate
            juliamirror = np.flip(julia, axis=0)
            julia = np.concatenate((julia, juliamirror), axis=0)

        # Plot the output with a random colormap using matplotlib
        self.plotJuliaSet(julia, show=show, fname=fname)

    def getComplexValue(self):
        """
        Random choice in a list of best complex values for Julia 
        sets (real, imag).

        :return cpxNum: a semi-random complex value
        """

        # Define the list of best complex values 
        cpxList = [
            (-0.10, 0.650), (0.00, 0.80), (0.370, 0.100), 
            (0.355, 0.355), (-0.54, 0.54), (0.340, -0.05),
            (0.37, 0.10), (0.355, 0.355)
        ]

        # Randomly choose one
        cpxTmp =  random.choice(cpxList)

        # Manipulate the base value slightly to make it a little more unique
        cpxNum = self.twearkComplex(cpxTmp)

        return cpxNum

    def twearkComplex(self, cpxTmp):
        """
        Manipulate the base value slightly to make it a little more unique.

        :param cpxTmp: complex value to modify
        :param cpxNum: a slightly manipulate version of the input
        """

        # Get the signs for the imaginary parts
        isign = random.randrange(-1, 1, 2)

        # Get a value variation for for real and imaginary parts
        # The possible variation range is fixed at +/- 2% to stay
        # In the neightborhood of the initial value
        rsigma = random.uniform(0.98, 1.02)
        isigma = random.uniform(0.98, 1.02)

        # Apply modification and return the new complex value
        realPart = cpxTmp[0] * rsigma
        imagPart = cpxTmp[1] * isigma * isign

        return complex(realPart, imagPart)

    def getTargetArea(self):
        """
        For more randomness, the target area is a random 
        subset of a wide one defined with x[-1.5, 1.5] and
        y[-1.5, 1.5]

        :return xrng, yrng: tuples containing (xmin, xmax)
            and (ymin, ymax)
        """

        # Randomly choose the center of the target area
        # Possible values are in [-1.0, 1.0] to stay in an
        # area where there are always pieces of fractals
        xctr = random.uniform(-1.0,1.0)
        yctr = random.uniform(-1.0,1.0)

        # Extend around the center
        xrng = (xctr-0.5, xctr+0.5)
        yrng = (yctr-0.5, yctr+0.5)

        return xrng, yrng

    def processJulia(self, cpxNum, xrng, yrng):
        """
        Calculate the Julia set for the given input parameters.

        :param cpxNum: complex value acting as a seed for the Julia set
        :param xrng: range of values (min, max) for the x-axis
        :param yrng: range of values (min, max) for the y-axis
        :param escrad: escape radius
        :param niter: maximum number of iterations
        """

        # Initialize numpy array of dimensions (size, size) with zeros
        julia = np.ones((self.size, self.size), dtype=np.float32)

        # Calculate the width (equal to height) of the image since the
        # image is defined as a square
        width = xrng[1] - xrng[0]  # xmax - xmin = ymax - ymin

        # Randomly choose the sign of the shade
        #ssign = random.randrange(-1, 1, 2)
        ssign = -1.

        # Loop over x range
        for ix in range(self.size):
            # Get the pixel position in the complex plane
            # For the real part
            realPart = float(ix) / self.size * width + xrng[0]

            # Loop over y range
            for iy in range(self.size):
                # Get the pixel position in the complex plane
                # For the imaginary part
                imagPart = float(iy) / self.size * width + yrng[0]

                # Build the complex
                cpxTmp = complex(realPart, imagPart)

                # Initialize iteration counter
                it = 0

                # Loop over iterations
                while(np.abs(cpxTmp) <= self.escrad**2 and it < self.niter):
                    # Quadratic polynomial
                    cpxTmp = cpxTmp**2 + cpxNum
                    # Increment iteration counter
                    it += 1
                
                # Calculate the shade (a cool thing find somewhere on the net)
                shade = 1. - np.sqrt(it/self.niter)

                # Fill the outpout array
                julia[ix][iy] = ssign * shade

        return julia

    def plotJuliaSet(self, julia, fname='juilaset-output', show=False):
        """
        Plot the output Julia set and show it in matplotlib window or
        write it on disk as a png file.

        :param julia: the Julia set
        :param show: if show is `False` th eoutput image will be
            written as a PNG file named `fname`
        :param fname: Name of the output PNG file to write on disk
        """

        # List of beautiful colormap for Julia sets
        cmapList = [
            cm.Blues, cm.Greens, cm.Purples, cm.hot, cm.inferno, 
            cm.binary, cm.rainbow, cm.twilight_shifted, cm.plasma
        ]

        # Randomly chose one colormap
        cmapName = random.choice(cmapList)

        # Plot the image with a gaussian interpolation
        fig = plt.gcf()
        fig.set_size_inches(3., 3.)
        plt.imshow(julia, interpolation='gaussian', cmap=cmapName)
        
        # Disable axis
        plt.axis('off')

        if(show):
            plt.show()
        else:
            # Write on disk
            fig.savefig(fname+".png", dpi=self.dpi, pad_inches=0.05, bbox_inches='tight')

def julia(**kwargs):
    """
    temp
    """
    # Initialize Julia Set instance
    juliaInstance = JuliaSet()

    # If kwargs not empty update the attributes
    if kwargs is not None:
        juliaInstance.param(**kwargs)

    return juliaInstance

if __name__ == "__main__":
    # execute only if run as a script
    genJuliaSet = JuliaSet()
    genJuliaSet.param()
    genJuliaSet.run()

[((3661, 3683), 'random.choice', 'random.choice', (['cpxList'], {}), '(cpxList)\n', (3674, 3683), False, 'import random\n'), ((4146, 4172), 'random.randrange', 'random.randrange', (['(-1)', '(1)', '(2)'], {}), '(-1, 1, 2)\n', (4162, 4172), False, 'import random\n'), ((4374, 4400), 'random.uniform', 'random.uniform', (['(0.98)', '(1.02)'], {}), '(0.98, 1.02)\n', (4388, 4400), False, 'import random\n'), ((4418, 4444), 'random.uniform', 'random.uniform', (['(0.98)', '(1.02)'], {}), '(0.98, 1.02)\n', (4432, 4444), False, 'import random\n'), ((5105, 5130), 'random.uniform', 'random.uniform', (['(-1.0)', '(1.0)'], {}), '(-1.0, 1.0)\n', (5119, 5130), False, 'import random\n'), ((5145, 5170), 'random.uniform', 'random.uniform', (['(-1.0)', '(1.0)'], {}), '(-1.0, 1.0)\n', (5159, 5170), False, 'import random\n'), ((5817, 5866), 'numpy.ones', 'np.ones', (['(self.size, self.size)'], {'dtype': 'np.float32'}), '((self.size, self.size), dtype=np.float32)\n', (5824, 5866), True, 'import numpy as np\n'), ((8018, 8041), 'random.choice', 'random.choice', (['cmapList'], {}), '(cmapList)\n', (8031, 8041), False, 'import random\n'), ((8112, 8121), 'matplotlib.pyplot.gcf', 'plt.gcf', ([], {}), '()\n', (8119, 8121), True, 'import matplotlib.pyplot as plt\n'), ((8166, 8224), 'matplotlib.pyplot.imshow', 'plt.imshow', (['julia'], {'interpolation': '"""gaussian"""', 'cmap': 'cmapName'}), "(julia, interpolation='gaussian', cmap=cmapName)\n", (8176, 8224), True, 'import matplotlib.pyplot as plt\n'), ((8265, 8280), 'matplotlib.pyplot.axis', 'plt.axis', (['"""off"""'], {}), "('off')\n", (8273, 8280), True, 'import matplotlib.pyplot as plt\n'), ((2795, 2817), 'numpy.flip', 'np.flip', (['julia'], {'axis': '(1)'}), '(julia, axis=1)\n', (2802, 2817), True, 'import numpy as np\n'), ((2838, 2882), 'numpy.concatenate', 'np.concatenate', (['(julia, juliamirror)'], {'axis': '(1)'}), '((julia, juliamirror), axis=1)\n', (2852, 2882), True, 'import numpy as np\n'), ((2958, 2980), 'numpy.flip', 'np.flip', (['julia'], {'axis': '(0)'}), '(julia, axis=0)\n', (2965, 2980), True, 'import numpy as np\n'), ((3001, 3045), 'numpy.concatenate', 'np.concatenate', (['(julia, juliamirror)'], {'axis': '(0)'}), '((julia, juliamirror), axis=0)\n', (3015, 3045), True, 'import numpy as np\n'), ((8312, 8322), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (8320, 8322), True, 'import matplotlib.pyplot as plt\n'), ((2657, 2670), 'numpy.abs', 'np.abs', (['julia'], {}), '(julia)\n', (2663, 2670), True, 'import numpy as np\n'), ((7187, 7211), 'numpy.sqrt', 'np.sqrt', (['(it / self.niter)'], {}), '(it / self.niter)\n', (7194, 7211), True, 'import numpy as np\n'), ((6838, 6852), 'numpy.abs', 'np.abs', (['cpxTmp'], {}), '(cpxTmp)\n', (6844, 6852), True, 'import numpy as np\n')]