# from https://mathstodon.xyz/@tpfto/115282738536916157

import mpmath as mpm
import numpy as np
import plotly.graph_objs as go

# prepare a polar mesh
rh = np.linspace(0.2, 0.8, 60)
th = np.linspace(-np.pi, np.pi, 90)
rG, tG = np.meshgrid(rh, th)
zz = rG * np.exp(1j * tG)

# Jacobi epsilon function, https://dlmf.nist.gov/22.16.E30
def ellipjep(zz, m):
    km = mpm.ellipk(m)
    em = mpm.ellipe(m)
    qm = mpm.qfrom(m=m)
    fac = mpm.mp.pi / (2 * km)
    uu = fac * zz
    zn = fac * mpm.jtheta(4, uu, qm, 1) / mpm.jtheta(4, uu, qm)
    res = zz * (em / km) + zn
    return res

# vectorize the elliptic functions
jaccn = np.vectorize((lambda u, m: complex(mpm.ellipfun("cn", u, m))))
jacdn = np.vectorize((lambda u, m: complex(mpm.ellipfun("dn", u, m))))
jacsn = np.vectorize((lambda u, m: complex(mpm.ellipfun("sn", u, m))))
jacep = np.vectorize((lambda u, m: complex(ellipjep(u, m))))

# period
cc = float(2 * mpm.ellipk(0.5))

# precompute elliptic functions over mesh
st = jacsn(cc * zz, 0.5)
ct = jaccn(cc * zz, 0.5)
dt = jacdn(cc * zz, 0.5)
et = jacep(cc * zz, 0.5)

# Chen-Gackstatter surface, https://doi.org/10.1007/BF01456948
x = np.real(
    (np.pi + cc * cc / 2) * zz - cc * et
    + (ct * dt * ((2 * np.pi / cc) / st - cc * st)) / (st**2)
)
y = np.imag(
    (np.pi - cc * cc / 2) * zz + cc * et
    + (ct * dt * ((2 * np.pi / cc) / st + cc * st)) / (st**2)
)
z = np.sqrt(6 * np.pi) * np.real((1 / st) ** 2) - np.sqrt(1.5 * np.pi)

# 'birdie' from https://tpfto.wordpress.com/2019/04/17/faking-a-birds-colors-and-miscellanea/
birdie = [
    "#3E26A8", "#4536D4", "#484CEF", "#4562FC", "#347AFC", "#2C90F0", "#21A3E3", "#0CB3D3",
    "#12BEB9", "#34C79C", "#5CCC76", "#94C94A", "#C8C129", "#F1BA37", "#FEC933", "#F5E327",
    "#F9FB14"
]

# and now, the plot...
surface = go.Surface(x=x, y=y, z=z, colorscale=birdie, showscale=False)
data = [surface]

pset = dict(
    gridcolor="#EEE8D5", zerolinecolor="#EEE8D5",
    showbackground=True, backgroundcolor="#839496"
)
layout = go.Layout(
    title="Chen-Gackstatter surface", scene=dict(xaxis=pset, yaxis=pset, zaxis=pset)
)

fig = go.Figure(data=data, layout=layout)
fig.show()