import matplotlib.pyplot as plt
from matplotlib import _pylab_helpers
# import cmath
import numpy as np
import mpmath
mpmath.mp.dps = 12

def zeta_relist_im_abs(re_list,im):
 z_re = np.zeros(len(re_list))
 z_im = np.zeros(len(re_list))
 z_abs = np.zeros(len(re_list))
 for i in range(len(re_list)):
 s = re_list[i] + im*1j
 tmp = mpmath.zeta(s)*mpmath.gamma(0.5*s)*mpmath.power(mpmath.pi,-0.5*s)
 z_re[i] = tmp.real
 z_im[i] = tmp.imag
 z_abs[i] = abs(tmp)
 # z_arg[i] = cmath.phase(tmp)
 return z_re,z_im,z_abs

def zeta_re_imlist_abs(re,im_list):
 z_re = np.zeros(len(im_list))
 z_im = np.zeros(len(im_list))
 z_abs = np.zeros(len(im_list))
 for i in range(len(im_list)):
 s = re + im_list[i]*1j
 tmp = mpmath.zeta(s)*mpmath.gamma(0.5*s)*mpmath.power(mpmath.pi,-0.5*s)
 z_re[i] = tmp.real
 z_im[i] = tmp.imag
 z_abs[i] = abs(tmp)
 # z_arg[i] = cmath.phase(tmp)
 return z_re,z_im,z_abs


def main():
 RE_N=11
 IM_N=51
 # RE_AMP=0.1
 IM_AMP=1
 
 fig,axes = plt.subplots(1, 4, figsize=(13,5), gridspec_kw= {'width_ratios': [1, 3, 2, 2]})
 axes[0].set_xticks([0.5,0.75,1])
 axes[0].set_xlim((0.5,1))
 axes[0].set_yticks([0,14,14.5,20])
 axes[0].set_ylim((0,20))
 axes[0].grid()
 axes[1].grid()
 axes[2].grid()
 axes[3].grid()

 # ハンドル取得のため一度plot
 lines00, = axes[0].plot([0,1], [0,0],"r") # s Re part sweep
 lines01, = axes[0].plot([0,1], [0,0],"b") # s Im part sweep s_re
 lines03, = axes[0].plot([0,1], [0,0],"g") # s Im part sweep s_re max
 lines10, = axes[1].plot([0,1], [0,0],"r") # zeta_xi(s) Re part sweep
 lines11, = axes[1].plot([0,1], [0,0],"b") # zeta_xi(s) Im part sweep
 lines13, = axes[1].plot([0,1], [0,0],"g") # zeta_xi(s) Im part sweep s_re max
 lines21, = axes[2].plot([0,1], [0,0],"b") # abs(zeta(s)) Im part sweep
 lines23, = axes[2].plot([0,1], [0,0],"g") # abs(zeta(s)) Im part sweep s_re max
 # lines24, = axes[2].plot([0,1], [0,0],"y") # diff of abs(zeta(s))
 lines3, = axes[3].plot([0,1], [0,0],"r") # abs(zeta(s)) Re part sweep 

 while True:
 manager = _pylab_helpers.Gcf.get_active()
 if manager is None:
 print('closed')
 break

 points = plt.ginput(n=1, timeout=86400, mouse_add=1, mouse_pop=2, mouse_stop=3)
 # a = plt.ginput(n=-1, mouse_add=1, mouse_pop=3, mouse_stop=2)
 # n=-1でインプットが終わるまで座標を取得
 # mouse_addで座標を取得（左クリック）
 # mouse_popでUndo（ミドルクリック）
 # mouse_stopでインプットを終了する（右クリック）

 if len(points)==1:
 re = 0.5
 re_max,im = points[0]
 print(re_max, im)
 
 re_list = np.zeros(RE_N)
 im_list_dummy = np.zeros(RE_N)

 re_list_dummy = np.zeros(IM_N)
 re_list_dummy3 = np.zeros(IM_N)
 im_list = np.zeros(IM_N)

 for i in range(RE_N):
 re_list[i] = re + (re_max-re)*i/(RE_N-1)
 im_list_dummy[i] = im

 for i in range(IM_N):
 re_list_dummy[i] = re
 re_list_dummy3[i] = re_list[-1]
 im_list[i] = im + IM_AMP*(i-(IM_N-1)/2)/((IM_N-1)/2)

 axes[0].set_ylim(None)
 lines00.set_data(re_list , im_list_dummy)
 lines01.set_data(re_list_dummy , im_list )
 lines03.set_data(re_list_dummy3, im_list )
 
 # Re sweep
 z_re1,z_im1,z_absr = zeta_relist_im_abs(re_list, im)

 # Im sweep
 z_re2,z_im2,z_abs2 = zeta_re_imlist_abs(re, im_list)
 z_re4,z_im4,z_abs4 = zeta_re_imlist_abs(re_list[-1], im_list) # [-1] means last element of the list

 lines10.set_data(z_re1, z_im1)
 lines11.set_data(z_re2, z_im2)
 lines13.set_data(z_re4, z_im4)

 # z_abs_diff = z_abs4-z_abs2
 lines21.set_data(im_list, z_abs2)
 lines23.set_data(im_list, z_abs4)
 #lines24.set_data(im_list, z_abs_diff)
 
 #tmp_abs_max1 = max(abs(z_abs_diff.min()),abs(z_abs_diff.max()))
 tmp_abs_max2 = abs(z_abs2.max())
 tmp_abs_max1 = abs(z_abs4.max())
 tmp_abs_max = max(tmp_abs_max1,tmp_abs_max2)
 #tmp_min = min(0,z_abs_diff.min())
 axes[2].set_xlim(im_list[0],im_list[-1])
 #axes[2].set_ylim(tmp_min,tmp_abs_max)
 axes[2].set_ylim(0,tmp_abs_max)

 lines3.set_data(re_list, z_absr)
 axes[3].set_xlim(re_list[0],re_list[-1])
 tmp_abs_max = abs(z_absr.max())
 tmp_abs_min = abs(z_absr.min())
 axes[3].set_ylim(tmp_abs_min,tmp_abs_max)

 tmp_abs_max1 = max(abs(z_re1.min()),abs(z_re1.max()),abs(z_re2.min()),abs(z_re2.max()),abs(z_re4.min()),abs(z_re4.max()))
 tmp_abs_max2 = max(abs(z_im1.min()),abs(z_im1.max()),abs(z_im2.min()),abs(z_im2.max()),abs(z_im4.min()),abs(z_im4.max()))
 
 axes[1].set_xlim((-tmp_abs_max1, tmp_abs_max1))
 axes[1].set_ylim((-tmp_abs_max2, tmp_abs_max2))
 #tmp_abs_max = max(tmp_abs_max1,tmp_abs_max2)
 #axes[1].set_xlim((-tmp_abs_max, tmp_abs_max))
 #axes[1].set_ylim((-tmp_abs_max, tmp_abs_max))
 else:
 print('quit by right click')
 break

 # plt.savefig('fig_test.png')
 # plt.show()
 # 次の描画(≒操作受付)までの待ち時間(秒)
 plt.pause(0.05)


if __name__ == '__main__':
 main()