import numpy as np
import matplotlib
#atplotlib.use("Agg")
from matplotlib import pyplot as plt
import matplotlib.animation as animation
from matplotlib.widgets import Slider
from matplotlib.colors import LogNorm
from matplotlib import gridspec

# ================================================
#      Generate output for movies
# ================================================
  
# name = batch name (no leading underscores)
def generate_output(fullname, res, force=False):
	import os
	name = fullname.split('/')[-1]
	path = '/'.join(fullname.split('/')[:-1])+'/tempy%s' %name
	generate = True
	
	try:
		os.mkdir(path)
	except OSError:
		generate = False
		pass
		
	if force==True:
		generate=True

	if generate==True:
		dyn = open(path+'/dyn.txt','w')
		file = open(fullname,'r')
		next(file)
		
		step = 0
		for line in file:
			data = str(line).split()
			line = next(file)
			
			if line != '\n':
				while True:
					try: nuc = open(path+'/ZY'+str(int(data[1])-2), 'w')
					except:
						data = next(file).split()
						continue
					break

				dyn.write("  ".join([data[1], data[3], data[5]])+"\n")
				
				while line != '\n':
					nuc.write(line)
					line = next(file)
				
				nuc.close()
		
			step += 1
		dyn.close()
		file.close()
		
	return
	
# ================================================
#      User input
# ================================================

fullname = input('Please give me the time file.\n').strip()
name = fullname.split('/')[-1]
path = '/'.join(fullname.split('/')[:-1])+'/tempy%s' %name
res = int(input('What timestep resolution would you like (10 recommended)? (e.g. "1" for every timestep, "5" for every fifth timestep)\n'))
abZ_choose = int(input('Show abZ pattern? ("1" or "0")\n'))

# use `force=True` if you want to re-parse the output
generate_output(fullname, res)#, force=True)

# ================================================

dyn = np.genfromtxt(path+'/dyn.txt')[::res]
finalab = np.genfromtxt(path+'/ZY'+str(res*(len(dyn)-1))).T

ab = {}
for j,Y in enumerate(finalab[2]):
	if finalab[1][j] in ab.keys():
		ab[finalab[1][j]] += Y
	else:
		ab[finalab[1][j]] = Y
		
ab = sorted([[A, Y] for A,Y in ab.items()])
ab = np.array(ab).T

# Solar
sfile = np.genfromtxt('yva_solar_sne08.dat')
norm = 0.0002
sol = []
for i in sfile:
	sol.append([i[0], i[1]*norm/sfile[np.where(sfile[:,0]==195)[0][0],1]])
sol = np.array(sol).T

	
# ================================================
#      Plotting
# ================================================

fig = plt.figure(figsize=(8,5))
fig.subplots_adjust(left=0.1, bottom=0.15, right=0.85)

if abZ_choose==1:
	gs = gridspec.GridSpec(3, 1, height_ratios=[2,1,1])
	ax3 = fig.add_subplot(gs[2])
	ax3.set_xlim(35,100)
	ax3.set_ylim(1E-8,1E-2)
	ax3.set_yscale('log')
	ax3.set_ylabel('Y')
	ax3.set_xlabel('Z')
	zline, = ax3.plot([], [], 'b-', alpha=0.8)
	solzfile = np.genfromtxt('Solar_r')
	solz = [[row[0], (10.0**float(row[1]))*norm/(10.0**float(sfile[58,1]))] for row in solzfile]
	solz = np.array(solz).T
	solarz = ax3.plot(solz[0], solz[1], 'k-', alpha=0.6, lw=1)

else:
	gs = gridspec.GridSpec(2, 1, height_ratios=[2.1,1])
	
ax1 = fig.add_subplot(gs[0])
ax1.set_xlim(20,220)
ax1.set_ylim(10,110)
ax1.set_ylabel('Z', color='black')
ax1.set_xlabel('N', color='black')
ax1.xaxis.set_ticks_position('top')
ax1.xaxis.set_label_position('top') 

plt.plot((10, 70), (20, 20), color='black', lw=2, alpha=0.4)
plt.plot((20, 80), (28, 28), color='black', lw=2, alpha=0.4)
plt.plot((60, 120), (50, 50), color='black', lw=2, alpha=0.4)
plt.plot((120, 190), (82, 82), color='black', lw=2, alpha=0.4)

plt.plot((28,28), (10,35), color='black', lw=2, alpha=0.4)
plt.plot((50, 50), (10, 50), color='black', lw=2, alpha=0.4)
plt.plot((82, 82), (30, 65), color='black', lw=2, alpha=0.4)
plt.plot((126, 126), (55, 90), color='black', lw=2, alpha=0.4)
plt.plot((184, 184), (75, 115), color='black', lw=2, alpha=0.4)

ax2 = fig.add_subplot(gs[1])
ax2.set_xlim(90,260)
ax2.set_ylim(1E-8,1E-2)

ax2.set_yscale('log')
ax2.set_ylabel('Y', color='black')
ax2.set_xlabel('A', color='black')

aline, = ax2.plot([], [], color='orange')
solar = ax2.plot(sol[0], sol[1], color='black', alpha=0.6, lw=1)

# -------------------------------------------------
#            Initialize animation vars
# -------------------------------------------------


time_text = ax1.text(0.5, 0.13, '', color='black', transform=ax1.transAxes)
t9_text = ax1.text(0.5, 0.05, '', color='black', transform=ax1.transAxes)

n = np.arange(0, 250)
z = np.arange(0, 150)
nn, zz = np.meshgrid(n, z)

nuc = np.genfromtxt(path+'/ZY0').T

y = [[np.nan]*250]*150
y = np.array(y, dtype=float)
N =  np.array(nuc[0], dtype=int)
Z =  np.array(nuc[1], dtype=int)

# list of stable isotopes
stable = np.genfromtxt('stable.txt')
for j,row in enumerate(stable):
	y[int(row[0]), int(row[1])] = 1

ax1.pcolormesh(n, z, y, cmap='Greys', vmin=-0.15*0.8)

# ==========================================================

y = np.array([[np.nan]*250]*150)

heatmap = ax1.pcolormesh(nn, zz, y, cmap=plt.cm.get_cmap('afmhot'), norm=LogNorm(vmin=10**-15, vmax=10**-2))
cax = fig.add_axes([0.85, 0.15, 0.02, 0.73])
cb = fig.colorbar(heatmap, cax=cax)
cb.set_label("Abundance", color='black')


# Slider
tax = plt.axes([0.15, 0.02, 0.70, 0.03], facecolor='white')
tslide = Slider(tax, 'time step', 0, len(dyn)-1, valinit=0, valfmt='%1.0f')

# initialization function: plot the background of each frame
# Used for animation
def init():
	heatmap.set_array(np.array([]))
	aline.set_data([], [])
	time_text.set_text('')
	t9_text.set_text('')

	if abZ_choose==1:
		zline.set_data([], [])
		return heatmap, aline, time_text, t9_text, zline

	return heatmap, aline, time_text, t9_text


# animation function.  This is called sequentially
# Used for Animation AND Slider
def animate(i):
	i = int(i)
	
	nuc = np.genfromtxt(path+'/ZY'+str(res*i)).T
	
	y = [[0.0]*250]*150
	y = np.array(y)#, dtype=float)
	
	N =  np.array(nuc[1], dtype=int)-np.array(nuc[0], dtype=int)
	Z =  np.array(nuc[0], dtype=int)
	
	for j,Y in enumerate(nuc[2]):
		nnum = int(nuc[1][j]) - int(nuc[0][j])
		y[int(nuc[0][j]), nnum] = Y
		
	ab = {}
	abZ = {}
	for j,Y in enumerate(nuc[2]):
		if nuc[1][j] in ab.keys():
			ab[nuc[1][j]] += Y
		else:
			ab[nuc[1][j]] = Y
		
		if nuc[0][j] in abZ.keys():
			abZ[nuc[0][j]] += Y
		else:
			abZ[nuc[0][j]] = Y
				
	ab = sorted([[A, Y] for A,Y in ab.items()])
	ab = np.array(ab).T
	abZ = np.array(sorted([[Z,Y] for Z,Y in abZ.items()])).T
		
	
	#abZ = np.genfromtxt(path+'/ZY'+str(res*i)).T
	#heatmap.set_array(y[:-1,:-1].ravel())
	heatmap.set_array(y)
	
	aline.set_data(ab[0], ab[1])
	time_text.set_text('time = {0:0.2e} s'.format(dyn[i][1]))
	t9_text.set_text(r'$T_9$ = {0:0.2e}'.format(dyn[i][2]))
	
	
	fig.canvas.draw_idle()
	
	if abZ_choose==1:
		zline.set_data(abZ[0], abZ[1])
		return heatmap, aline, time_text, t9_text, zline
	return heatmap, aline, time_text, t9_text

	
tslide.on_changed(animate)

# call the animator.  blit=True means only re-draw the parts that have changed.
#anim = animation.FuncAnimation(fig, animate, init_func=init, frames=len(dyn), interval=5)#, blit=True)

# -------------------------------------------------
#                 More plot options
# -------------------------------------------------

ax1.xaxis.grid(True, ls=':')
ax1.yaxis.grid(True, ls=':')

gs.update(wspace=0, hspace=0.02)
plt.show()
plt.close()

#writer = animation.writers['ffmpeg'](fps=5)
#anim.save('rprocess_movie.mp4', writer=writer, dpi=250, savefig_kwargs={'facecolor':'black'})#, fps=20)