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

# ## Plotly plot of the Riemann surface, real part of square root, $\Re(\sqrt{z})$

# In[5]:


import numpy as np
from numpy import pi, exp
import plotly.graph_objects as go
import plotly
plotly.offline.init_notebook_mode()


# The graphs of complex functions are usually plotted with the hsv colormap, that encodes the argument of the function values. HSV is a bad colormap, since its lightness is not constant. But for comparison with all that have been published in math, related to visualization of the complex functions, we are using first, the hsv colormap, then cmocean.phase, and a cyclic colormap from Colorcet coloction of colormaps. 

# In[24]:


hsv= [[0.0, 'rgb(0, 242, 242)'],
 [0.083, 'rgb(0, 121, 242)'],
 [0.167, 'rgb(0, 0, 242)'],
 [0.25, 'rgb(121, 0, 242)'],
 [0.333, 'rgb(242, 0, 242)'],
 [0.417, 'rgb(242, 0, 121)'],
 [0.5, 'rgb(242, 0, 0)'],
 [0.583, 'rgb(242, 121, 0)'],
 [0.667, 'rgb(242, 242, 0)'],
 [0.75, 'rgb(121, 242, 0)'],
 [0.833, 'rgb(0, 242, 0)'],
 [0.917, 'rgb(0, 242, 121)'],
 [1.0, 'rgb(0, 242, 242)']]


cyclic=[[0.0, '#657be7'], #colorcet cyclic_mygbm_50-90_c46_n256_s25
 [0.05, '#9276d9'],
 [0.1, '#b279c8'],
 [0.15, '#ca7eb6'],
 [0.2, '#dc87a5'],
 [0.25, '#e89495'],
 [0.3, '#eea187'],
 [0.35, '#efb179'],
 [0.4, '#ebc16d'],
 [0.45, '#e1d163'],
 [0.5, '#ccdd63'],
 [0.55, '#b0de74'],
 [0.6, '#94d88a'],
 [0.65, '#79d09e'],
 [0.7, '#5cc6b0'],
 [0.75, '#3dbbc1'],
 [0.8, '#20b0ce'],
 [0.85, '#04a4d9'],
 [0.9, '#1597e2'],
 [0.95, '#3589e9'],
 [1.0, '#657be7']];


# We are plotting the Riemann surface associated to the real part of the complex square root function,
# $w=f(z)$, where $f(z)=\sqrt{z}$. Since Python implements only the principal branch of the function $\sqrt{z}$
# we derive  the graph of $\Re(f(z))$, from the graph of the inverse function, $f^{-1}(w)=w^2$.

# In[25]:


rho = np.linspace(0, 1.0, 100)
phi = np.linspace(-np.pi, np.pi, 300)
rho, phi = np.meshgrid(rho,phi)
w = rho * exp(1j*phi)
zc = w**2
fig = go.Figure(go.Surface(x=np.real(zc), y=np.imag(zc), z=np.real(w), 
                    colorscale=hsv,
                    surfacecolor= np.angle(w),
                    cmin=-pi, cmax=pi,
                    colorbar=dict(thickness=24, len=0.75,
                       tickvals=[-pi, -2*pi/3, -pi/3, 0, pi/3, 2*pi/3, pi], 
                       ticktext=["-\u03c0", "-2\u03c0/3", "-\u03c0/3",
                                 "0", "\u03c0/3", "2\u03c0/3", "\u03c0"],
                       title="arg(f)")       
                    ))
plot_title = 'Riemann surface Re(sqrt(z))),  colored<br> according to the argument values on each branch'
fig.update_layout(title_text=plot_title,
                  title_x=0.5,
                  font=dict(family='Balto'),
                  width=600,
                  height=600)
fig.update_scenes(xaxis_title='Re(z)',
                  yaxis_title='Im(z)', 
                  zaxis_title='Re(f(z))',
                  camera_eye=dict(x=1.5, y=1.5, z=1))


# In[26]:


fig.update_traces(colorscale="phase")


# In[27]:


fig.update_traces(colorscale=cyclic)


# In[14]:


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