# -*- coding: utf-8 -*-
"""
Created on Sun Dec 16 20:35:02 2018

@author: Buet Blanche
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

plt.close("all")

   
epsilon = 10**(-6)
maxIt = 50


N = 1000
image = np.zeros((N, N, 3)) # on réinitialise à une image noire

# initialiser la discretisation de [-1,1] * [-1,1]

px = np.linspace(-1,1,N)
py = np.linspace(1,-1,N) 

# Newton

def f(z):
    return z**3 - 1

def fprime(z):
    return 3*z**2

roots = np.array([complex(-0.5, 0.5*np.sqrt(3)), complex(-0.5, -0.5*np.sqrt(3)), 1.])

# Newton method
def newtonModified(fun, funprime, z0, epsilon, maxIt):
    It = 0
    b = -1
    while (b == -1) & (It < maxIt):
        z1 = z0 - fun(z0)/funprime(z0)
        z0 = z1
        It += 1
        for r in range(roots.size):
            if np.abs(z1 - roots[r]) < epsilon:
                b = r
    return (z1, It, b)

# bassins d'attraction des racines + nb d'itérations
image2 = np.ones((N, N, 3)) # on crée une image blanche

for ix in range(N):
    for iy in range(N):
        z0 = complex(px[ix], py[iy])
        sol, It, b = newtonModified(f, fprime, z0, epsilon, maxIt)
        # on colorie en fonction du nombre d'itérations image2        
        image2[iy, ix, :] = It/maxIt*np.ones(3)
        if b== 0:
            image[iy, ix, 0] = 1 # en rouge
            image[iy, ix, 1] = 1 
        elif b == 1:
            image[iy, ix, 1] = 1 # en vert
            image[iy, ix, 2] = 1
        elif b == 2:
            image[iy, ix, 2] = 1 # en bleu
            image[iy, ix, 0] = 1
        # et si b == -1 on laisse en noir
        
plt.imshow(image)
plt.savefig('newton1.pdf')

plt.figure(6)
plt.imshow(image2)
plt.savefig('newton2.pdf')