%matplotlib inline
%config InlineBackend.figure_format='svg'
import numpy as np
import pylab as pl

1+1j

(1+1j)**2

2**2j

abs(1+1j)

def zeta0(s):
    x = 0
    for i in range(1,1000000):
        x += 1 / (i ** s)
    return x

zeta0(0.5+14.134j)

#import scipy.special
#scipy.special.zeta(0.5+14.134j, 0)

import scipy.misc as scm
scm.comb(1000000, 10)

scm.comb(10,2,exact=True)

from scipy.misc import comb
LOWER_THRESHOLD=1.0e-6
UPPER_BOUND=1.0e+4
INFINITY=300

def zeta(s):
   outer_sum = 0.0
   prev = np.inf
   if s == 1:
      return np.inf
   for m in range(1,INFINITY):
      inner_sum = 0.0
      for j in range(1, m+1):
         k = comb(m-1, j-1, exact=True) * j**(-s)
         inner_sum += k if j%2 else -k
      inner_sum = inner_sum * 2**(-m)
      outer_sum += inner_sum
      if abs(prev - inner_sum) < LOWER_THRESHOLD:
         break
      if abs(outer_sum) > UPPER_BOUND:
         break
      prev = inner_sum
   return outer_sum / (1-(2**(1-s)))

zeta(0.5+14.134j)

x = np.linspace(0,50,1000)
%time y = np.vectorize(zeta)(0.5 + x*1j)

#x = np.linspace(0,30,1000)
#%time y = np.vectorize(zeta)(0.8 + x*1j)
pl.grid()
pl.plot(x, np.abs(y))
pl.figure()
pl.grid()
pl.plot(np.real(y), np.imag(y))

re = np.linspace(-5, 5, 51)
im = np.linspace(-40, 40, 101)
IM,RE = np.meshgrid(im, re)
%time Z = np.vectorize(zeta)(RE + IM*1j)

pl.imshow(np.abs(Z), vmin=0, vmax=5, aspect=0.4)

#%matplotlib 
%matplotlib inline
from mpl_toolkits.mplot3d import Axes3D
fig = pl.figure(figsize=(10,6))
ax = Axes3D(fig)
ZZ = np.abs(Z)
ZZ[ZZ > 10] = 10
ax.plot_wireframe(RE, IM, ZZ)
#ax.plot_wireframe(RE, IM, np.log(np.abs(Z)))
#ax.plot_surface(RE, IM, np.log(np.abs(Z)), rstride=1, cstride=1, cmap='hot')