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

# # Images of an accretion disk around a Kerr black hole
# 
# This Jupyter/SageMath notebook is relative to the lectures
# [Geometry and physics of black holes](https://relativite.obspm.fr/blackholes/).

# In[1]:


version()


# In[2]:


get_ipython().run_line_magic('display', 'latex')


# In[3]:


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


# ## Functions $\ell_{\rm c}(r_0)$ and $q_{\rm c}(r_0)$ for critical null geodesics
# 
# We use $m=1$ and denote $r_0$ simply by $r$.

# In[4]:


a, r = var('a r') 


# In[5]:


lsph(a, r) = (r^2*(3 - r) - a^2*(r + 1))/(a*(r -1))
lsph


# In[6]:


qsph(a, r) = r^3 / (a^2*(r - 1)^2) * (4*a^2 - r*(r - 3)^2)
qsph


# The radii $r_+$ and $r_-$ of the two horizons:

# In[7]:


rp(a) = 1 + sqrt(1 - a^2)
rm(a) = 1 - sqrt(1 - a^2)


# ## Critical radii $r_{\rm ph}^{**}$, $r_{\rm ph}^*$, $r_{\rm ph}^+$, $r_{\rm ph}^-$, $r_{\rm ph}^{\rm ms}$ and $r_{\rm ph}^{\rm pol}$ 

# In[8]:


rph_ss(a) = 1/2 + cos(2/3*asin(a) + 2*pi/3)
rph_ss


# In[9]:


rph_s(a) = 4*cos(acos(-a)/3 + 4*pi/3)^2
rph_s


# In[10]:


rph_p(a) = 4*cos(acos(-a)/3)^2
rph_p


# In[11]:


rph_m(a) = 4*cos(acos(a)/3)^2
rph_m


# We add the radius of the marginally stable orbit:

# In[12]:


rph_ms(a) = 1 - (1 - a^2)^(1/3)
rph_ms


# as well as the radius of outer and inner polar orbits:

# In[13]:


rph_pol(a) = 1 + 2*sqrt(1 - a^2/3)*cos(1/3*arccos((1 - a^2)/(1 - a^2/3)^(3/2)))
rph_pol


# In[14]:


rph_pol_in(a) = 1 + 2*sqrt(1 - a^2/3)*cos(1/3*arccos((1 - a^2)/(1 - a^2/3)^(3/2)) + 2*pi/3)
rph_pol_in


# In[15]:


a0 = 0.95
# a0 = 1
show(LatexExpr(r'r_{\rm ph}^{**} = '), n(rph_ss(a0)))
show(LatexExpr(r'r_{\rm ph}^{\rm ms} = '), n(rph_ms(a0)))
show(LatexExpr(r'r_{\rm ph}^{*} = '), n(rph_s(a0)))
show(LatexExpr(r'r_- = '), n(rm(a0)))
show(LatexExpr(r'r_+ = '), n(rp(a0)))
show(LatexExpr(r'r_{\rm ph}^+ = '), n(rph_p(a0)))
show(LatexExpr(r'r_{\rm ph}^- = '), n(rph_m(a0)))
show(LatexExpr(r'r_{\rm ph}^{\rm pol} = '), n(rph_pol(a0)))
show(LatexExpr(r'r_{\rm ph}^{\rm pol,in} = '), n(rph_pol_in(a0)))


# ### Prograde timelike ISCO radius compared to various radii of spherical photon orbits

# In[16]:


def r_isco(a, retrograde=False):
    eps = -1 if not retrograde else 1
    a2 = a^2
    z1 = 1 + (1 - a2)^(1/3) * ((1 + a)^(1/3) + (1 - a)^(1/3))
    z2 = sqrt(3*a2 + z1^2)
    return 3 + z2 + eps*sqrt((3 - z1)*(3 + z1 + 2*z2))


# In[17]:


r_isco(0.5), r_isco(0.95), r_isco(0.99), r_isco(1.)


# In[18]:


g = plot(r_isco(a), (a, 0, 1), color='red', thickness=2, 
         legend_label=r'$r_{\rm ISCO}^+$', axes_labels=[r'$a/m$', r'$r/m$'], 
         frame=True, axes=False, gridlines=True)
g += plot(rph_p(a), (a, 0, 1), color='green', thickness=2,
          legend_label=r'$r_{\rm ph}^+$')
g += plot(rph_m(a), (a, 0, 1), color='green', thickness=2,
          linestyle='--', legend_label=r'$r_{\rm ph}^-$')
g += plot(rph_pol(a), (a, 0, 1), color='green', thickness=2,
          linestyle=':', legend_label=r'$r_{\rm ph}^{\rm pol}$')
g.set_legend_options(handlelength=2.5)
g.save('gik_rISCO_rph.pdf')
g


# In[19]:


fr(a) = r_isco(a) - rph_m(a)
find_root(fr, 0.6, 0.7)


# In[20]:


fr(a) = r_isco(a) - rph_pol(a)
find_root(fr, 0.8, 1)


# In[21]:


def alpha(a, th_obs, r0):
    if a == 1:
        ell = - r0^2 + 2*r0 + 1
    else:
        ell = lsph(a, r0)
    return - ell / sin(th_obs)

def Theta(a, th_obs, r0):
    if a == 1:
        ell = - r0^2 + 2*r0 + 1
        q = r0^3 * (4 - r0)
    else:
        ell = lsph(a, r0)
        q = qsph(a, r0)
    return q + cos(th_obs)^2 * (a^2 - ell^2/sin(th_obs)^2)

def beta(a, th_obs, r0, eps_theta=1):
    return eps_theta * sqrt(Theta(a, th_obs, r0,))


# In[22]:


def r0_bounds(a, th_obs, outer=True):
    r"""
    Return `(r0_min, r0_max)`
    """
    if outer:
        r1 = n(rph_p(a))
        r2 = n(rph_m(a))
        r3 = rph_pol(a)
    else:
        r1 = n(rph_ss(a))
        r2 = 0
        r3 = rph_pol_in(a)
    #
    # Computation of rmin:
    try:
        if a == 1:
            th_crit = n(asin(sqrt(3) - 1))
            if n(th_obs) < th_crit or n(th_obs) > n(pi) - th_crit:
                rmin = find_root(lambda r: Theta(a, th_obs, r), r1, r3)
            else:
                rmin = 1
        else:
            rmin = find_root(lambda r: Theta(a, th_obs, r), r1, r3)
    except TypeError:  # special case th_obs = pi/2
        rmin = r1    
    #
    # Computation of rmax:
    try:
        rmax = find_root(lambda r: Theta(a, th_obs, r), r3, r2)
    except TypeError:  # special case th_obs = pi/2
        rmax = r2    
    #
    return (rmin, rmax)


# In[23]:


def shadow_plot(a, th_obs, orientation=0, outer=True, color=None, number_colors=5, 
                range_col=None, thickness=2, linestyle='-', plot_points=200,
                legend='automatic', legend_loc=(1.02, 0.36), fill=True, 
                fillcolor='grey', draw_NHEKline=True, 
                draw_spin=False, spin_arrow_length=7, spin_arrow_options=None,
                frame=True, axes=True, axes_labels='automatic', gridlines=True):
    if a==0:
        # Case a = 0:
        rs = 3*sqrt(3)
        if color is None:
            color = 'black'
        if legend == 'automatic':
            legend = None
        g = parametric_plot((rs*cos(x), rs*sin(x)), (x, 0, 2*pi), 
                            color=color, thickness=thickness,
                            linestyle=linestyle, legend_label=legend, 
                            fill=fill, fillcolor=fillcolor, frame=frame, 
                            axes=axes, gridlines=gridlines)
    else:
        # Case a != 0
        rmin, rmax = r0_bounds(a, th_obs, outer=outer)
        if rmin > 0:
            rmin = 1.00000001*rmin
            rmax = 0.99999999*rmax
        else:
            rmin = 0.9999999*rmin
            rmax = 1.0000001*rmax
        print("rmin : ", rmin, "  rmax : ", rmax)
        co = cos(orientation)
        so = sin(orientation)
        fa = lambda r: co*alpha(a, th_obs, r) - so*beta(a, th_obs, r)
        fb = lambda r: so*alpha(a, th_obs, r) + co*beta(a, th_obs, r)
        fam = lambda r: co*alpha(a, th_obs, r) - so*beta(a, th_obs, r, eps_theta=-1)
        fbm = lambda r: so*alpha(a, th_obs, r) + co*beta(a, th_obs, r, eps_theta=-1)
        if range_col is None:
            range_col = r0_bounds(a, pi/2, outer=outer)
        rmin_col, rmax_col = range_col 
        print("rmin_col : ", rmin_col, "  rmax_col : ", rmax_col)
        dr = (rmax_col - rmin_col) / number_colors
        rm = rmin_col + int((rmin - rmin_col)/dr)*dr
        r1s = rmin
        r_ranges = []
        while rm + dr < rmax:
            col = hue((rm - rmin_col)/(rmax_col - rmin_col + 0.1))
            r2s = rm + dr
            r_ranges.append((r1s, r2s, col))
            rm += dr
            r1s = r2s
        if color is None:
            col = hue((rm - rmin_col)/(rmax_col - rmin_col + 0.1))
        else:
            col = color
        r_ranges.append((r1s, rmax, col))
        g = Graphics()
        legend_label = None  # a priori
        if a == 1 and draw_NHEKline:
            th_crit = asin(sqrt(3) - 1)
            if th_obs > th_crit and th_obs < pi - th_crit:
                # NHEK line
                alpha0 = -2/sin(th_obs)
                beta0 = sqrt(3 - cos(th_obs)**2 *(6 + cos(th_obs)**2))/sin(th_obs)
                alpha1 = co*alpha0 - so*beta0
                beta1 = so*alpha0 + co*beta0
                alpha2 = co*alpha0 + so*beta0
                beta2 = so*alpha0 - co*beta0
                if legend == 'automatic':
                    legend_label = r"$r_0 = m$"
                if color is None:
                    colNHEK = 'maroon'
                else:
                    colNHEK = color
                g += line([(alpha1, beta1), (alpha2, beta2)], color=colNHEK, 
                          thickness=thickness, linestyle=linestyle,
                          legend_label=legend_label)
                if fill:
                    g += polygon2d([(fa(rmax), fb(rmax)), (alpha1, beta1), (alpha2, beta2)], 
                                   color=fillcolor, alpha=0.5)
        for rg in r_ranges:
            r1s, r2s = rg[0], rg[1]
            col = rg[2]
            if legend:
                if legend == 'automatic':
                    if draw_NHEKline and abs(r1s - 1) < 1e-5:
                        legend_label = r"${:.2f}\, m < r_0 \leq {:.2f}\, m$".format(
                                       float(r1s), float(r2s))
                    else:
                        legend_label = r"${:.2f}\, m \leq r_0 \leq {:.2f}\, m$".format(
                                       float(r1s), float(r2s))
                else:
                    legend_label = legend
            g += parametric_plot((fa, fb), (r1s, r2s), plot_points=plot_points, color=col, 
                                 thickness=thickness, linestyle=linestyle,
                                 legend_label=legend_label, 
                                 frame=frame, axes=axes, gridlines=gridlines)
            g += parametric_plot((fam, fbm), (r1s, r2s), plot_points=plot_points, color=col, 
                                 thickness=thickness, linestyle=linestyle)
        if fill:
            g += parametric_plot((fa, fb), (rmin, rmax), fill=True, fillcolor=fillcolor, 
                                 thickness=0)
            g += parametric_plot((fam, fbm), (rmin, rmax), fill=True, fillcolor=fillcolor, 
                                 thickness=0)
        if draw_spin:
            if not spin_arrow_options:
                spin_arrow_options = {}
            if 'color' not in spin_arrow_options:
                spin_arrow_options['color'] = color
            g += arrow2d((0,0), (-so*spin_arrow_length, co*spin_arrow_length), 
                         **spin_arrow_options)
        # end of case a != 0
    g.set_aspect_ratio(1)
    if axes_labels:
        if axes_labels == 'automatic':
            g.axes_labels([r"$(r_{\mathscr{O}}/m)\; \alpha$", 
                           r"$(r_{\mathscr{O}}/m)\; \beta$"])
        else:
            g.axes_labels(axes_labels)
    if legend:
        g.set_legend_options(handlelength=2, loc=legend_loc)
    return g


# In[24]:


shadow_plot(1, pi/2, fill=False, color='red', number_colors=1, 
            thickness=1.5, linestyle=':', legend=False)


# In[25]:


shadow_plot(1, pi/2, orientation=pi/6, fill=True, color='red', 
            number_colors=1, thickness=1.5, 
            draw_spin=True, spin_arrow_options={'color': 'blue', 'width': 3},
            legend=False)


# Half-width of the image in units of $m/r_{\mathscr{O}}$:

# In[26]:


gyoto_field_of_view = 74  # Gyoto field of view for M87* (in microarcseconds)
scale_M87 = 3.66467403690525  # m/r for M87* (in microarcseconds)
fsize = gyoto_field_of_view / 2 / scale_M87
#extent = (-fsize, fsize, -fsize, fsize)
print("fsize =", fsize, "m/r_obs")


# In[27]:


def empty_plot(size, frame=True, axes=False, axes_labels='automatic', 
               gridlines=False):
    g = Graphics()
    g._extra_kwds['frame'] = frame
    g._extra_kwds['axes'] = axes
    g._extra_kwds['gridlines'] = gridlines
    if axes_labels:
        if axes_labels == 'automatic':
            g.axes_labels([r"$(r_{\mathscr{O}}/m)\; \alpha$", 
                           r"$(r_{\mathscr{O}}/m)\; \beta$"])
        else:
            g.axes_labels(axes_labels)
    g.set_aspect_ratio(1)
    g.set_axes_range(-size, size, -size, size)
    return g


# Default resolution of SageMath images:

# In[28]:


dpi = Graphics.SHOW_OPTIONS['dpi']
dpi


# Default resolution of Matplotlib images:

# In[29]:


from matplotlib import rcParams
rcParams['figure.dpi']


# In[30]:


gyoto_images = {'a95_th00': (0.95, 0),
                'a95_th30': (0.95, pi/6),
                'a95_th60': (0.95, pi/3),
                'a95_th90': (0.95, pi/2)}


# In[31]:


#gyoto_images = {'a1_th90_rin193': (1., pi/2),
#                'a1_th90_rin193-5000': (1., pi/2, 5000, 1800, 2500),
#                'a99_th90_rin193': (0.99, pi/2, 2500, 900, 1250)}


# In[32]:


get_ipython().run_line_magic('display', 'plain')
get_ipython().run_line_magic('matplotlib', 'notebook')
frame = True
figsize_factor = 1  # default: 1
axes_labels = 'automatic'
#axes_labels = False
test_size = False
raw = False
for gimage, param in gyoto_images.items():
    img = mpimg.imread(gimage + '.png')
    print(gimage, " resolution:", img.shape)
    # size of the image in inches:
    figsize = (float(img.shape[0])/dpi*figsize_factor, 
               float(img.shape[1])/dpi*figsize_factor)
    print("figsize: ", figsize)
    a0, th_obs = param[:2]
    print("a/m =", a0, "  theta_obs =", th_obs)
    if len(param) > 2:
        resol = param[2]
        imin = param[3] - 1
        jmin = param[4] - 1
    else:
        resol = img.shape[0]
        imin = 0
        jmin = 0
    imax = imin + img.shape[1]
    jmax = jmin + img.shape[0]
    xmin = (imin - resol/2)/resol * 2*fsize
    xmax = (imax - resol/2)/resol * 2*fsize
    ymin = (jmin - resol/2)/resol * 2*fsize
    ymax = (jmax - resol/2)/resol * 2*fsize
    extent = (xmin, xmax, ymin, ymax)
    print("extent: ", extent)
    # Critical curve as a SageMath graphics object from shadow_plot
    if th_obs == 0:
        th_obs = 0.001
    if not raw:
        gcrit = shadow_plot(a0, th_obs, fill=False, color='red', number_colors=1, 
                            thickness=1.5, linestyle=':', frame=frame, axes=False, 
                            axes_labels=axes_labels, gridlines=False, legend=False)
    else:
        gcrit = empty_plot(fsize, axes_labels=axes_labels, frame=frame)
    if test_size:
        axes_labels_bck = gcrit.axes_labels()
        gcrit += point((-fsize, 0), size=60, color='red', zorder=100)
        gcrit += point((fsize, 0), size=60, color='red', zorder=100)
        gcrit += point((0, -fsize), size=60, color='red', zorder=100)
        gcrit += point((0, fsize), size=60, color='red', zorder=100)
        gcrit.axes_labels(axes_labels_bck)
    gcrit.set_axes_range(xmin, xmax, ymin, ymax)
    # Matplotlib figure corresponding to gcrit:
    options = gcrit.SHOW_OPTIONS.copy()
    options.update(gcrit._extra_kwds)
    options['figsize'] = figsize
    options['axes_pad'] = 0
    options.pop('dpi')           # strip meaningless options for matplotlib
    options.pop('transparent')   #
    options.pop('fig_tight')     #
    fcrit = gcrit.matplotlib(**options)
    # Adding the Gyoto image onto it:
    ax = fcrit.axes[0]
    ax.imshow(img, extent=extent)
    # Save result to png and pdf
    ext = '_raw' if raw else '_crit'
    fcrit.savefig(gimage + ext + '.png', dpi=dpi, pad_inches=0, 
                  bbox_inches='tight', transparent=True)
    fcrit.savefig(gimage + ext + '.pdf', pad_inches=0, 
                  bbox_inches='tight')
    # Only for display in the current notebook:
    fig = plt.figure(figsize=figsize, dpi=dpi, frameon=False)
    imgc = mpimg.imread(gimage + '_crit.png')
    if axes_labels:
        imgcp = plt.imshow(imgc)
    else:
        imgcp = plt.imshow(imgc, extent=extent)
    print(" ")


# In[ ]: