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

# ## Plotly interactive visualization of complex valued functions ##

# In this Jupyter Notebook we generate via [Plotly](https://plot.ly), an interactive plot of a complex valued function, $f:D\subset\mathbb{C}\to\mathbb{C}$. A complex function is visualized using a version of the 
# [domain coloring method](http://nbviewer.jupyter.org/github/empet/Math/blob/master/DomainColoring.ipynb).
# 
# Compared with other types of domain coloring, the Plotly interactive plot  is much more informative. It displays for each point $z$ in a rectangular region of the complex plane, 
# not only the hsv color associated to $\arg(f(z)$ (argument of  $f(z)$), but also  the values  $\arg(f(z)$  and $\log(|f(z)|)$ (the  log modulus). 

# First we define a Plotly hsv (hue, saturation, value) colorscale, adapted to the range of the numpy functions, `np.angle`, respectively, `np.arctan2`. 
# We plot a Heatmap over a rectangular region in the complex plane, colored via this  colorscale, according to
# the values of  the  $\arg(f(z))$. Over the Heatmap are plotted a few contour lines of the log modulus $\log(|f(z)|)$.
# 

# In[ ]:


import numpy as np
import numpy.ma as ma
from numpy import pi
import matplotlib.pyplot as plt
import matplotlib.colors


# In[ ]:


def hsv_colorscale(S=1, V=1): 
    if S < 0 or S > 1 or V < 0 or V > 1:
        raise ValueError('Parameters S (saturation), V (value, brightness) must be in [0,1]')
   
    argument = np.array([-pi, -5*pi/6, -2*pi/3, -3*pi/6, -pi/3,  -pi/6, 0, 
                       pi/6, pi/3, 3*pi/6, 2*pi/3, 5*pi/6, pi])
    
    H = argument/(2*np.pi)+1
    H = np.mod(H,1)
    Sat = S*np.ones_like(H)
    Val = V*np.ones_like(H)
    
    HSV = np.dstack((H, Sat, Val))
    RGB = matplotlib.colors.hsv_to_rgb(HSV) 
   
    colormap = (255* np.squeeze(RGB)).astype(int) 
    #Define and return the Plotly hsv colorscale adapted to polar coordinates for complex valued functions
    step_size = 1 / (colormap.shape[0]-1)
    return [[round(k*step_size, 4), f'rgb{tuple(c)}']  for k, c in enumerate(colormap)]   


# In[ ]:


pl_hsv=hsv_colorscale()
pl_hsv


# The following two functions compute data needed for visualization:

# In[ ]:


def evaluate_function(func,  re=(-1,1), im=(-1,1), N=100):
    # func is the complex function to be ploted
    #re, im are the interval ends on the real and imaginary axes, defining the rectangular region in the complex plane
    #N gives the number of points in an interval of length 1
    l = re[1]-re[0]
    h = im[1]-im[0]
    resL = int(N*l) #horizontal resolution
    resH = int(N*h) #vertical resolution
    X = np.linspace(re[0], re[1], resL)
    Y = np.linspace(im[0], im[1], resH)
    x, y = np.meshgrid(X,Y)
    z = x+1j*y
    w = func(z)
    argument = np.angle(w)
    modulus = np.absolute(w)
    log_modulus = ma.log(modulus)
    return X,Y, argument, log_modulus


# In[ ]:


def get_levels(fmodul, nr=10):
    #define the levels for contour plot of the modulus |f(z)|
    #fmodul is the log modulus of f(z) computed on meshgrid
    #nr= the number of  contour lines
    mv = np.nanmin(fmodul)
    Mv = np.nanmax(fmodul)
    size = (Mv-mv)/float(nr)
    return [mv+k*size for k in range(nr+1)]


# In[ ]:


import plotly.graph_objects as go


# We extract the contour lines from an attribute of the `matplotlib.contour.QuadContourSet` object returned by the `matplotlib.pyplot.contour`.
# 
# The function defined in the next cell retrieves the points on the contour lines segments, and defines the corresponding Plotly traces:

# In[ ]:


def plotly_contour_lines(contour_data):
    #contour_data is a matplotlib.contour.QuadContourSet object returned by plt.contour
    contours=contour_data.allsegs  # 
    if len(contours)==0:
        raise ValueError('Something wrong hapend in computing contour lines')
    #contours is a list of lists; if contour is a list in contours, its elements are arrays. 
    #Each array defines a segment of contour line at some level
    xl = []# list of x coordinates of points on a contour line
    yl = []#         y
    #lists of coordinates for contour lines consisting in one point:
    xp = []# 
    yp = []

    for k,contour in enumerate(contours):
        L = len(contour)
        if L!= 0: # sometimes the list of points at the level np.max(modulus) is empty
            for ar in contour:
                if ar.shape[0] == 1:
                    xp += [ar[0,0], None]
                    yp += [ar[0,1], None]
                else: 
                    xl += ar[:,0].tolist()
                    yl += ar[:,1].tolist()
                    xl.append(None) 
                    yl.append(None)

    lines = go.Scatter(x=xl, 
                       y=yl, 
                       mode='lines',  
                       name='modulus', 
                       line=dict(width=1,
                            color='#a5bab7', 
                            shape='spline', 
                            smoothing=1),
                       hoverinfo='skip'
                 )   
    if len(xp) == 0:
        return lines
    else: 
        points = go.Scatter(x=xp, y=yp, mode='markers',
                            marker=dict(size=4, color='#a5bab7'), 
                            hoverinfo='none')
        return lines, points


# Set the layout of the plot:

# In[ ]:


def  set_plot_layout(title, width=500, height=500):
    
    return go.Layout(title_text=title,
                     title_x=0.5,
                     width=width,
                     height=height,
                     showlegend=False,
                     xaxis_title='Re(z)',
                     yaxis_title='Im(z)')
                    
 


# Define a function that associates to each point $z$ in a meshgrid, the strings containing the values to be displayed when hovering the mouse over that point: 

# In[ ]:


def text_to_display(x, y, argum, modulus):
    m, n = argum.shape
    return [['z='+'{:.2f}'.format(x[j])+'+' + '{:.2f}'.format(y[i])+' j'+'<br>arg(f(z))='+'{:.2f}'.format(argum[i][j])+\
       '<br>log(modulus(f(z)))='+'{:.2f}'.format(modulus[i][j])  for j in range(n)] for i in range(m)]


# Finally, the function `plotly_plot` calls all above defined functions in order to generate the phase plot 
# of a given complex function:

# In[ ]:


def plotly_plot(f, re=(-1,1), im=(-1,1), N=50, nr=10, title='', width=500, height=500, **kwargs):
    x, y, argument, log_modulus=evaluate_function(f, re=re, im=im, N=N)
    
    levels = get_levels(log_modulus, nr=nr)
    plt.figure(figsize=(0.05,0.05))
    plt.axis('off')
    cp = plt.contour(x,y, log_modulus, levels=levels)
    cl = plotly_contour_lines(cp)
    text = text_to_display(x,y, argument, log_modulus.data)
    tickvals = [-np.pi, -2*np.pi/3, -np.pi/3, 0, np.pi/3, 2*np.pi/3, np.pi]
   
    #define the above values as strings with pi-unicode
    ticktext=['-\u03c0', '-2\u03c0/3', '-\u03c0/3', '0', '\u03c0/3', '2\u03c0/3', '\u03c0']
   
    dom = go.Heatmap(x=x, y=y, z=argument,  colorscale=pl_hsv, 
                    text=text, hoverinfo='text',
                    colorbar=dict(thickness=20, tickvals=tickvals, 
                    ticktext=ticktext,  
                    title='arg(f(z))'))
    if len(cl) == 2  and isinstance(cl[0], go.Scatter):
        data = [dom, cl[0], cl[1]]
    else: 
        data = [dom, cl]
    
    layout = set_plot_layout(title=title)    
    fig=go.Figure(data=data, layout =layout)
    return fig


# As an example we take the function $f(z)=\sin(z)/(1-cos(z^3))$: 

# In[ ]:


fig = plotly_plot(lambda z: np.sin(z)/(1-np.cos(z**3)), re=(-2,2), im=(-2,2), 
                nr=22, title='$f(z)=\\sin(z)/(1-\\cos(z^3))$');
fig.update_layout(xaxis_range=[-2,2], yaxis_range=[-2,2])
fig.show()


# In[1]:


from IPython.display import IFrame    
url = "https://chart-studio.plotly.com/~empet/13957"
IFrame(url, width=700, height=700)


# In[2]:


from IPython.core.display import HTML
def  css_styling():
    styles = open("./custom.css", "r").read()
    return HTML(styles)
css_styling()


# In[ ]: