#!/usr/bin/env python -i

import csv, sys, os, array, warnings, subprocess
import matplotlib.pyplot as plt
import numpy as np



def cloudy_abund(fname):
	'''
	This reads data from a file produced as part of a cloudy PL loop -
	'''

	IP=[]
	abundance=[]
	inp=open(fname)
	for line in inp.readlines():
		data=line.split()
		eval_dat=[]
		for i in range(len(data)):
			eval_dat.append(float(data[i]))
		IP.append (eval_dat[0])
		abundance.append(eval_dat[1:])		
		

	return IP,np.transpose(np.array(abundance))

def python_abund(fname='test.dat'):
	'''

	'''
	
	t_or_ip=[]
	inp=open(fname,'r')
	lines=inp.readlines()
	data=lines[0].split()
	nion=len(data)-1
	ntemp=len(lines)
	abund=np.zeros((ntemp,nion))
	for i in range(ntemp):
		data=lines[i].split()
		t_or_ip.append(float(data[0]))
		for j in range(nion):
			abund[i][j]=float(data[j+1])
			
	abund2=np.transpose(abund)

	return t_or_ip,abund2

'''
                    UNLV

Synopsis:  

This routine generates comparison plots from data files generated by
either cloudy_pl_loop or python_pl_loop. 


Command line usage (if any):

	usage: cloudy_pl_loop *run1* *run2*.
		the expected file names are 

		cl_hydrogen_*run1*.dat
or		py_hydrogen_*run1*.dat.

The cl or py prefix tells the code which type of data we have.

Description:  

Primary routines:

Notes:
									   
History:

081214 nsh Commented

'''



if __name__ == "__main__":		# allows one to run from command line without running automatically with write_docs.py


	run_name1=sys.argv[1]
	run_name2=sys.argv[2]



	'''
	from matplotlib import rc
	rc('font',**{'family':'serif','serif':['Times']})
	rc('font',size=10)
	#rc('text', usetex=True)
	'''
	xlim=[-7,1]
	ylim=[0.0001,2.000]


	try:
		open('py_hydrogen_'+run_name1+'.dat','r')
		type1='p'
	except:
		try:
			open('cl_hydrogen_'+run_name1+'.dat','r')
			type1='c'
		except:
			print "File names not constructed as expected"
			raise SystemExit

	if type1=='c':
		IP1,h1=cloudy_abund('cl_hydrogen_'+run_name1+'.dat')
		IP1,he1=cloudy_abund('cl_helium_'+run_name1+'.dat')
		IP1,c1=cloudy_abund('cl_carbon_'+run_name1+'.dat')
		IP1,n1=cloudy_abund('cl_nitrogen_'+run_name1+'.dat')
		IP1,o1=cloudy_abund('cl_oxygen_'+run_name1+'.dat')
		IP1,fe1=cloudy_abund('cl_iron_'+run_name1+'.dat')
	elif type1=='p':
		IP1,h1=python_abund('py_hydrogen_'+run_name1+'.dat')
		IP1,he1=python_abund('py_helium_'+run_name1+'.dat')
		IP1,c1=python_abund('py_carbon_'+run_name1+'.dat')
		IP1,n1=python_abund('py_nitrogen_'+run_name1+'.dat')
		IP1,o1=python_abund('py_oxygen_'+run_name1+'.dat')
		IP1,fe1=python_abund('py_iron_'+run_name1+'.dat')
		IP1=np.log10(IP1)


	try:
		open('py_hydrogen_'+run_name2+'.dat','r')
		type2='p'
	except:
		try:
			open('cl_hydrogen_'+run_name2+'.dat','r')
			type2='c'
		except:
			print "File names not constructed as expected"
			raise SystemExit




	if type2=='c':
		IP2,h2=cloudy_abund('cl_hydrogen_'+run_name2+'.dat')
		IP2,he2=cloudy_abund('cl_helium_'+run_name2+'.dat')
		IP2,c2=cloudy_abund('cl_carbon_'+run_name2+'.dat')
		IP2,n2=cloudy_abund('cl_nitrogen_'+run_name2+'.dat')
		IP2,o2=cloudy_abund('cl_oxygen_'+run_name2+'.dat')
		IP2,fe2=cloudy_abund('cl_iron_'+run_name2+'.dat')
	elif type2=='p':
		IP2,h2=python_abund('py_hydrogen_'+run_name2+'.dat')
		IP2,he2=python_abund('py_helium_'+run_name2+'.dat')
		IP2,c2=python_abund('py_carbon_'+run_name2+'.dat')
		IP2,n2=python_abund('py_nitrogen_'+run_name2+'.dat')
		IP2,o2=python_abund('py_oxygen_'+run_name2+'.dat')
		IP2,fe2=python_abund('py_iron_'+run_name2+'.dat')
		IP2=np.log10(IP2)






	xlabel="Ionization parameter (cloudy definition)"
	pdetails="1cm thick shell, 1e11cm radius, Nh=1e7"
	size=(9,5)




	fig=plt.figure(figsize=size,dpi=300)
	ax1=fig.add_subplot(111)
	ax1.set_color_cycle(['k','r'])
	ax1.set_title("Hydrogen ionization state",size=10)
	for i in range(2):
		ax1.semilogy(IP1,h1[i],label=run_name1+'_H'+str(i+1))
	for i in range(2):
		ax1.semilogy(IP2,h2[i],'x',ms=2,label=run_name2+'_H'+str(i+1))
	ax1.set_xlabel(xlabel)
	ax1.set_ylabel("Relative abundance")
	ax1.set_ylim(ylim)
	ax1.set_xlim(xlim)
	ax1.legend(loc='lower right',prop={'size':6})
	plt.savefig('hydrogen_plot_'+run_name1+'_'+run_name2+'.png')



	fig=plt.figure(figsize=size,dpi=300)
	ax2=fig.add_subplot(111)
	ax2.set_color_cycle(['k','r','b'])
	ax2.set_title("Helium ionization state",size=10)
	for i in range(3):
		ax2.semilogy(IP1,he1[i],label=run_name1+'_He'+str(i+1))
	for i in range(3):
		ax2.semilogy(IP2,he2[i],'x',ms=2,label=run_name2+'_He'+str(i+1))
	ax2.set_xlabel(xlabel)
	ax2.set_ylabel("Relative abundance")
	ax2.set_ylim(ylim)
	ax2.set_xlim(xlim)
	ax2.legend(loc='lower left',prop={'size':6})
	plt.savefig('helium_plot_'+run_name1+'_'+run_name2+'.png')


	fig=plt.figure(figsize=size,dpi=300)
	ax3=fig.add_subplot(111)
	ax3.set_color_cycle(['k','r','g','b','c','y','m'])
	ax3.set_title("Carbon ionization state",size=10)
	for i in range(7):
		ax3.semilogy(IP1,c1[i],label=run_name1+'_C'+str(i+1))
	for i in range(7):
		ax3.semilogy(IP2,c2[i],'x',ms=2,label=run_name2+'_C'+str(i+1))
	ax3.set_xlabel(xlabel)
	ax3.set_ylabel("Relative abundance")
	ax3.set_ylim(ylim)
	ax3.set_xlim(xlim)
	ax3.legend(loc='lower left',prop={'size':6})
	plt.savefig('carbon_plot_'+run_name1+'_'+run_name2+'.png')


	fig=plt.figure(figsize=size,dpi=300)
	ax4=fig.add_subplot(111)
	ax4.set_color_cycle(['k','r','g','b','c','y','m','k'])
	ax4.set_title("Nitrogen ionization state",size=10)
	for i in range(8):
		ax4.semilogy(IP1,n1[i],label=run_name1+'_N'+str(i+1))
	for i in range(8):
		ax4.semilogy(IP2,n2[i],'x',ms=2,label=run_name2+'_N'+str(i+1))
	ax4.set_xlabel(xlabel)
	ax4.set_ylabel("Relative abundance")
	ax4.set_ylim(ylim)
	ax4.set_xlim(xlim)
	ax4.legend(loc='lower left',prop={'size':6})
	plt.savefig('nitrogen_plot_'+run_name1+'_'+run_name2+'.png')


	fig=plt.figure(figsize=size,dpi=300)
	ax=fig.add_subplot(111)
	ax.set_color_cycle(['k','r','g','b','c','y','m','k','r'])
	ax.set_title("Oxygen ionization state",size=10)
	for i in range(9):
		ax.semilogy(IP1,o1[i],label=run_name1+'_O'+str(i+1))
	for i in range(9):
		ax.semilogy(IP2,o2[i],'x',ms=2,label=run_name2+'_O'+str(i+1))
	ax.set_xlabel(r"$\rm{\log(U)}$")
	ax.set_ylabel("Relative abundance")
	ax.set_ylim(ylim)
	ax.set_xlim(xlim)
	ax.legend(loc='lower left',prop={'size':6})
	plt.savefig('oxygen_plot_'+run_name1+'_'+run_name2+'.png')



	xlim=[-7,3]
	ylim=[0.0001,2.000]
	fig=plt.figure(figsize=size,dpi=300)
	ax=fig.add_subplot(111)
	ax.set_color_cycle(['k','r','g','b','c','y','m','k','r'])
	ax.set_title("Iron ionization state PL modelled as PL in thin shell python")
	for i in range(9):
		ax.semilogy(IP1,fe1[i],label=run_name1+'_Fe'+str(i+1))
	for i in range(9):
		ax.semilogy(IP2,fe2[i],'x',label=run_name2+'_Fe'+str(i+1))
	ax.set_xlabel(xlabel)
	ax.set_ylabel("Relative abundance")
	ax.set_ylim(ylim)
	ax.set_xlim(xlim)
	ax.legend(loc='lower right',prop={'size':8})
	plt.savefig('iron_low_plot_'+run_name1+'_'+run_name2+'.png')

	xlim=[-1,3]
	ylim=[0.0001,2.000]
	fig=plt.figure(figsize=size,dpi=300)
	ax=fig.add_subplot(111)
	ax.set_color_cycle(['k','r','g','b','c','y','m','k','r'])
	ax.set_title("Iron ionization state PL modelled as PL in thin shell python")
	for i in range(9):
		ax.semilogy(IP1,fe1[i+9],label=run_name1+'_Fe'+str(i+1))
	for i in range(9):
		ax.semilogy(IP2,fe2[i+9],'x',label=run_name2+'_Fe'+str(i+1))
	ax.set_xlabel(xlabel)
	ax.set_ylabel("Relative abundance")
	ax.set_ylim(ylim)
	ax.set_xlim(xlim)
	ax.legend(loc='lower right',prop={'size':8})
	plt.savefig('iron_mid_plot_'+run_name1+'_'+run_name2+'.png')


	xlim=[0,7]
	ylim=[0.0001,2.000]
	fig=plt.figure(figsize=size,dpi=300)
	ax=fig.add_subplot(111)
	ax.set_color_cycle(['k','r','g','b','c','y','m','k','r'])
	ax.set_title("Iron ionization state PL modelled as PL in thin shell python")
	for i in range(9):
		ax.semilogy(IP1,fe1[i+18],label=run_name1+'_Fe'+str(i+1))
	for i in range(9):
		ax.semilogy(IP2,fe2[i+18],'x',label=run_name2+'_Fe'+str(i+1))
	ax.set_xlabel(xlabel)
	ax.set_ylabel("Relative abundance")
	ax.set_ylim(ylim)
	ax.set_xlim(xlim)
	ax.legend(loc='lower right',prop={'size':8})
	plt.savefig('iron_high_plot_'+run_name1+'_'+run_name2+'.png')



