import math, cmath
from cmath import exp, sin, cos
from math import pi, tanh
import numpy as np
from numpy import vectorize

import matplotlib as mp
import matplotlib.pyplot as plot
from colorsys import hls_to_rgb

# line width of the frame grid
lw_grid=2

# color of the axes
axes_color = [0.28,0.28,0.26]

# lightness gradient
Lgrad = 0.5

# number of iterations
#  100: preview
# 1000: publication
N = 100

# viewport
X=[-2,6]
Y=[-4,4]


# function to be graphed
def f1(z):
  return z*z-(4+2j)*z+5


style = {
  "axes.linewidth": lw_grid
}
mp.rcParams.update(style)

def colorize(z):
  n,m = z.shape
  c = np.zeros((n,m,3))
  c[np.isinf(z)] = (1.0, 1.0, 1.0)
  c[np.isnan(z)] = (0.5, 0.5, 0.5)
  idx = ~(np.isinf(z) + np.isnan(z))

  # hue
  H = 1.0-np.angle(z[idx])/(2*pi)

  # lightness
  L = (1.0 - 1.0/(1.0+abs(z[idx])**Lgrad))

  c[idx] = [hls_to_rgb(h,l,1.0) for h,l in zip(H,L)]
  return c

fig = plot.figure()
ax = fig.add_subplot(1,1,1)
ax.spines['bottom'].set_color(axes_color)
ax.spines['top'].set_color(axes_color) 
ax.spines['right'].set_color(axes_color)
ax.spines['left'].set_color(axes_color)

ax.xaxis.set_tick_params(width=lw_grid,length=-6,pad=10,color=axes_color)
ax.yaxis.set_tick_params(width=lw_grid,length=-6,pad=10,color=axes_color)
ax.set_aspect('equal')
ax.axis(X+Y)
ax.set_xticks(range(X[0],X[1]+1,1))
ax.set_yticks(range(Y[0],Y[1]+1,1))

axis_x = np.linspace(X[0],X[1],N)
axis_y = np.linspace(Y[0],Y[1],N)
x,y = np.meshgrid(axis_x,axis_y)
z = x - y*1j

w = np.zeros((N,N),dtype='complex')
w = vectorize(f1)(z).conj()

# interpolation='bicubic'
ax.imshow(colorize(w), interpolation='none',
  extent=(X[0],X[1],Y[0],Y[1]))

# contour lines
cl = ax.contour(x, -y, abs(w), [2,5,10,20],
  colors="#606060", alpha=0.4, linewidths=3.2)
ax.clabel(cl, inline=1, fontsize=16, fmt="%g", colors="#000000",
  manual=[[4.0,-2.6],[1.1,-2.5],[0.9,-1.8],[0.9,-1.2]])

plot.savefig("plot.jpg",bbox_inches='tight')