###################################################################
#
# CPU, a CLASS Plotting Utility 
# v1.3 
# Benjamin Audren, 07.11.2011 
#
# This is a small python program aimed to gain time when comparing two
# spectra, i.e. from CAMB and CLASS, or a non-linear spectrum to a
# linear one.  It is designed to be used in a command line fashion, not
# being restricted to your CLASS directory, though it recognized mainly
# CLASS output format.  Far from perfect, or complete, it could use any
# suggestion for enhancing it, just to avoid losing time on useless
# matters for others.  Be warned that, when comparing with other
# format, the following is assumed: there are no empty line (especially
# at the end of file). Gnuplot comment lines (starting with a # are
# allowed). This issue will cause a non-very descriptive error in CPU,
# any suggestion for testing it is welcome.  Example of use: To
# superimpose two different spectra and see their global shape : python
# CPU output/lcdm_z2_pk.dat output/lncdm_z2_pk.dat To precisely compare
# the non-linear contribution of one-loop w.r.t. to TRG method: python
# CPU -i output/lcdm_1l_z1_pk.dat output/lcdm_1l_z1_pk_nl_density.dat
# output/lcdm_trg_pk_nl_density.dat The documentation available with
# --help should cover any question.
###################################################################

from scipy import *
from scipy import interpolate
import numpy as np
import os,sys,string,io,subprocess
import argparse

# parser for input arguments 

parser = argparse.ArgumentParser(description='CPU, a CLASS Plotting Utility, specify wether you want to superimpose, divide or interpolate different files, and behold!',
  epilog='''A standard usage would be, for instance :
    python CPU output/test_1l_pk.dat output/test_1l_pk_nl_density.dat -i 0
    python CPU output/wmap_cl.dat output/planck_cl.dat''',
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('files',	metavar='Files',
  type=str, nargs='*',
  help='the relative path of the desired file to plot\nfrom the root directory of CLASS')
parser.add_argument('-d, --divide',			
  dest='merging',action='store_const',const='blend_against',default='blend_together',
  help='divide the spectra from different files. The k-values must be strictly identical (default: plot every graph on the same plot)')
parser.add_argument('-i, --interpolate',metavar='index',
  nargs='?',dest='interp',const=True,default=False,
  help='interpolate the spectra on each set of k-values if different. Use to compare two (or more) graphs with different k-values. If nothing more is specified, it will plot both index in gnuplot. Otherwise, just add -i 0 to plot only the first one.')
parser.add_argument('-t',	metavar='pk, cl_lin,etc...',
  help='specify the file type (whether pk or cl (if not specified, cl==cl_lin, other choices are cl_log and cl_ll for log linear), only needed if your file name does not already contain one of these keywords...')
parser.add_argument('-colnum',metavar='2, 3, ...',
  help='specify the column you want to plot. By default, column is 2: you are plotting TT spectrum')
parser.add_argument('-x, --x11',metavar='gnuplot terminal',
  dest='term',action='store_const',const='x11',default='aqua',
  help='if your version of gnuplot does not support the aqua terminal, it will pick instead the x11 one, for crappier displays')
parser.add_argument('-p, --print',	
  dest='printfile',action='store_true',default=False,
  help='print the graph directly in a .png file')
parser.add_argument('-r, --repeat',
  dest='repeat',action='store_true',default=False,
  help='repeat the step for all redshifts with same base name, whatever the desired action')
parser.add_argument('-c, --cleaning',metavar='path',
  nargs='?',dest='cleaning',const=True,default=False,
  help='remove all .plt files in the current directory (if none specified) : keep your folders clean !')

# Remove all .plt, interp and divided files generated in output/ 
# By default, will clean current directory.

def clean(path):
  pattern1,pattern2,pattern3='.plt','_interp.dat','_divided.dat'
  path+="/"
  print 'Cleaning {0} directory from .plt files ...'.format(path)
  i=0 
  for each in os.listdir(path):
    if ( (each.rfind(pattern1)!=-1) or (each.rfind(pattern2)!=-1) or (each.rfind(pattern3)!=-1) ):
      name= "{0}{1}".format(path,each)
      print '  deleting {0}'.format(name)
      if os.remove(name)!= None:
	break
      i+=1
  if i==0:
    print ' ==> Already as clean as possible'
  else:
    print '  ==> Done'

# Errors ###########################
def error_format():
  print "  Hum... have you really provided a .dat file ?"
  print "  We apologise for the inconvenience, but CPU is exiting now"
  exit()

def error_type():
  print "  Spectrum type unrecognized or unsupported yet, sorry!"
  print "  Please try the -t command in addition to your previous call"
  print "  We apologise for the inconvenience"
  exit()

def error_number_of_files():
  print "  You specified a wrong number of files for the operation you demanded, maybe you do not want to divide only one file, for instance ?"
  print "  We apologise for the inconvenience"
  exit()
####################################

# Subfunction called to write on the gnuplot script file the proper file names.

def printstring(names):
  print_file=names[0].rstrip(".dat")
  print_file+=".ps"
  print_string="set terminal po enhanced color\nset output '{0}'\nreplot".format(print_file)
  print '{0} has been generated'.format(print_file)
  return print_string

# Subfunction called to write the most of the gnuplot script file according to options.

def headers_plot_file(spectrum_type,names,plot_line,term):
  tmp=open(names[0],"r")

  # printing option
  if ((args.printfile is True) and (plot_line is not False)):
    print_string=printstring(names)
  else:
    print_string=""
  
  # for interpolation: index choice option
  if plot_line is True:
    if args.interp is not True:
      if args.interp is not False:
	index_string=" index {0} ".format(args.interp)
      else:
	index_string=" "
    else:
      index_string=" "

  # if args colnum not specified, put it to 2,
  if args.colnum is None:
    args.colnum=2


  if spectrum_type=='cl_lin':
    if plot_line is True:
      plot_string="plot "
      for name in names:
	plot_string+="'{0}'".format(name)+"{0}u 1:{1} w l,".format(index_string,args.colnum)
      plot_string=plot_string.rstrip(",")
      plot_string+="\n"
    else:
      plot_string=''
    return "set terminal {0} enhanced\n".format(term)+"set xlabel 'l'\nset ylabel 'l(l+1)C_l / 2{/Symbol p}'\nset xr [2:]\nset format y '%.0t*10^{%T}'\nset key right\nset title 'CLASS output'\n"+plot_string+print_string
  elif spectrum_type=='cl_log':
    for line in tmp:
      if line.rfind('multipoles')!=-1:
	lmax= line.split(None)[-1]
	break
    if plot_line is True:
      plot_string="plot "
      for name in names:
	plot_string+="'{0}'".format(name)+"{0}u 1:{1} w l,".format(index_string,args.colnum)
      plot_string=plot_string.rstrip(",")
      plot_string+="\n"
    else:
      plot_string=''
    return "set terminal {0} enhanced\n".format(term)+"set logscale x\nset xr [2:{0}]\n".format(lmax)+"set xlabel 'l'\nset format y '%.0t*10^{%T}'\n"+"set ylabel 'l(l+1)C_l / 2{/Symbol p}'\nset key left\nset title 'CLASS output'\n"+plot_string+print_string
  elif spectrum_type=='cl_ll':
    for line in tmp:
      if line.rfind('multipoles')!=-1:
	lmax= line.split()[-1]
	break
    if plot_line is True:
      plot_string="plot "
      for name in names:
	plot_string+="'{0}'".format(name)+"{0}u (sqrt($1)):{1} w l,".format(index_string,args.colnum)
      plot_string=plot_string.rstrip(",")
      plot_string=plot_string+"\n"
    else:
      plot_string=''
    return "set terminal {0} enhanced\n".format(term)+"set xlabel 'l'\nset ylabel 'l(l+1)C_l / 2{/Symbol p}'\nset key right\nset title 'CLASS output'\nset xtics ('2' (2),'100' (100), '500' (500), '1000' (1000), '1500' (1500), '2000' (2000),'2500' (2500))"+"\nset format y '%.0t*10^{%T}'\n"+"set xr [2:{0}]\n".format(lmax)+plot_string+print_string
  elif spectrum_type=='pk':
    if plot_line is True:
      plot_string="plot "
      for name in names:
	plot_string+="'{0}'".format(name)+"{0}w l,".format(index_string)
      plot_string=plot_string.rstrip(",")
      plot_string+="\n"
    else:
      plot_string=""
    for line in tmp:
      if line.rfind('redshift')!=-1:
	z= line.split("=")[1]
	z= z.rstrip("\n")
	break
      else:
	z= 'I have no idea'
    return "set terminal {0} enhanced\n".format(term)+"set logscale\nset xlabel 'k (h/Mpc)'\nset ylabel 'P_k (Mpc/h)^3'\nset key right\nset title 'Power spectrum at z={0}'\n".format(z)+plot_string+print_string
  else:
    return None
  tmp.close()

# Print all asked files one next to the other (default operation if files with different k values)
def blend_together(names,spectrum_type,term):
  gnuplot_file=names[0].replace(".dat",".plt")
  if gnuplot_file==names[0]:
    error_format()
  print '  creating {0}'.format(gnuplot_file)
  plotfile = open(gnuplot_file, "w")
  plotfile.write(headers_plot_file(spectrum_type, names,True,term))
  plotfile.close()
  _plot(gnuplot_file)

# Print all asked files with respect to the same k (or anything) values, all y-data being divided by the y data of the first file.
# Only valid if the k values are exactly the same (values and number of values).
def blend_against(names,spectrum_type,term):
  gnuplot_file=names[0].replace(".dat",".plt")
  data_file=names[0].replace(".dat","_divided.dat")
  if gnuplot_file==names[0]:
    error_format()
  print '  creating {0} and {1}'.format(gnuplot_file,data_file)
  string=['' for rows in range(10000)]
  imax=0
  datafile=open(data_file,"w")
  for name in names:
    i=0
    tmp=open(name,"r")
    for line in tmp:
      if ((line.find('#')==-1) and (line.rfind('000000')==-1)):
	string[i]=string[i]+line.rstrip("\n")+"\t"
	if i>imax:
	  imax=i
	i+=1
    tmp.close()
  for i in range(imax):
    datafile.write(string[i]+"\n")
  plotfile = open(gnuplot_file,"w")
  plotfile.write(headers_plot_file(spectrum_type,names,False,term))
  plot_string='unset logscale\nplot '
  x_base=1
  field_base=2
  field=2
  size=len(names)
  for i in range(size):	
    field+=2
    plot_string+="'{0}' u {1}:(${2}/${3}) w l,".format(data_file,x_base,field,field_base)
  plot_string=plot_string.rstrip(',')
  plot_string+="\n"
  plotfile.write(plot_string)
  plotfile.close()
  datafile.close()
  _plot(gnuplot_file)

# If k values are different, an interpolation is done and outputs a data file. For a two files case:
# File 1 contains ( k1 | P1(k1) ), File 2 ( k2 | P2(k2) ). The data file created will contain:
# k1 | P1(k1)        | P2(k1)_interp
#(blank space for gnuplot using index 0, etc) 
# k2 | P1(k2)_interp | P2(k2)
def blend_against_interp(names,spectrum_type,term):
  gnuplot_file=names[0].replace(".dat",".plt")
  data_file=names[0].replace(".dat","_interp.dat")
  if gnuplot_file==names[0]:
    error_format()
  print '  creating {0} and {1}'.format(gnuplot_file,data_file)
  plotfile = open(gnuplot_file,"w")
  l=0
  jmax=[0 for col in range(10)]
  lmax=0
  kmax=10000000
  kmin=0
  spam = np.array([[[0 for col in range(10)] for row in range(10000)] for depth in range(2)],dtype=float)
  for name in names:
    currentfile = open(name,"r")
    print '  reading {0}..'.format(name)
    j=0
    for line in currentfile:
      if (line.find('#')==-1):
	line=line.split()
	spam[0,j,l]=float(line[0])
	spam[1,j,l]=float(line[1])
	j+=1
    jmax[l]=j
    if spam[0,jmax[l]-1,l]<=kmax:
      kmax=spam[0,jmax[l]-1,l]
    if spam[0,0,l]>=kmin:
      kmin=spam[0,0,l]
    l+=1
    currentfile.close()
  lmax=l
  lower_bound=[0 for col in range(10)]
  upper_bound=[0 for col in range(10)]
  #determining the upper and lower bound for each file, to only do interpolation
  for l in range (lmax):
    for j in range (jmax[l]):
      if ((spam[0,j,l]<kmin) and (spam[0,j+1,l]>=kmin)):
	lower_bound[l]=j+1
      if (spam[0,j,l]<=kmax):
	upper_bound[l]=j
  print '  -> done'

  #creating the new data file
  curves=[np.array for row in range(lmax)]
  interpolated=[np.array for row in range(lmax)]
  for l in range (lmax):
    x=spam[0,lower_bound[l]:upper_bound[l],l]
    y=spam[1,lower_bound[l]:upper_bound[l],l]
    curves[l]=interpolate.splrep(x,y)

  datafile = open(data_file,"w")
  for l in range (lmax):
    for ll in range(lmax):
      if ll==l:
	interpolated[ll]=spam[1,lower_bound[l]:upper_bound[l],l]
      else:	  
	x2=spam[0,lower_bound[l]:upper_bound[l],l]
        interpolated[ll]=interpolate.splev(x2,curves[ll],der=0)
    for j in range(upper_bound[l]-lower_bound[l]-2):
      data_string=''
      for ll in range(lmax):
	data_string=data_string+str((interpolated[ll])[j+1])+"\t"
      datafile.write(str(spam[0,j+lower_bound[l]+1,l])+"\t"+data_string+"\n")
    datafile.write("\n\n")
  datafile.close()

  # if index chosen
  if args.interp is not True:
    index_string_interp=" index {0}".format(args.interp)
  else:
    index_string_interp=""
  plotfile = open(gnuplot_file,"w")
  plotfile.write(headers_plot_file(spectrum_type,names,False,term))
  plotfile.write("unset logscale \n")
  if args.t in ['cl_log', 'pk']:
    plotfile.write("set logscale x\n")
  plotfile.write("set xr [{0}:{1}]\n".format(kmin,kmax))
  #plotfile.write("set yr [-0.001:0.005]\n")
  plot_string="plot "
  for l in range(lmax-1):
    plot_string+="'{0}' {1} u 1:(${2}/$2-1) w l,".format(data_file,index_string_interp,l+3)
  plot_string=plot_string.rstrip(",")
  plot_string=plot_string+"\n"
  plotfile.write(plot_string)

  if (args.printfile is True):
    plotfile.write(printstring(names))

  plotfile.close()
  _plot(gnuplot_file)

# Launch session of gnuplot with generated gnuplot script file
def _plot(gnuplot_file):
  os.system("gnuplot -persist '{0}'".format(gnuplot_file))


#######################################################
################## MAIN PART ##########################
#######################################################

print '~~~ CPU, a CLASS Plotting Utility ~~~'
args = parser.parse_args()

# check if the user want to clean its directory first 
if args.cleaning is not False:
  if args.cleaning is not True:
    clean(args.cleaning)
  else:
    clean(os.getcwd())
  exit()

# if there are no argument in the input, print usage
if len(args.files)==0:
  parser.print_usage()
  exit()

# if the first file name contains cl or pk, infer the type of desired spectrum
if ((args.files[0].rfind('cl')!=-1) and (args.t is None)):
  spectrum_type='cl_lin'
  args.t = 'cl_lin'
elif args.files[0].rfind('pk')!=-1:
  spectrum_type='pk'
  args.t = 'pk'
elif args.t is not None:
  spectrum_type=args.t
else:
  error_type()

# repeater
temp_path=args.files[0].split("/")
if len(temp_path)> 1:
  path=temp_path[-2]
else:
  path=os.getcwd()
list_file=['' for i in range(len(args.files)*10)]
if (temp_path[-1].rfind("z")!=-1 and args.repeat is True):
  root_name=temp_path[-1].split("z")[-2]
  for any in range (0,len(args.files)):
    extension_name=args.files[any].split("/")[-1].split("_")[-1]
    i=0
    for each in os.listdir(path):
      if (("z" in each) and (each.split("z")[-2]==root_name) and (each.split("_")[-1]==extension_name)):
	if len(temp_path)>1:
	  list_file[any+len(args.files)*i]=temp_path[-2]+"/"+each
	else:
	  list_file[any+len(args.files)*i]=each
	i+=1
  if args.repeat is True:
    repeat_len=i-1
  else:
    repeat_len=0
else:
  for any in range(0,len(args.files)):
    list_file[any]=args.files[any]
  repeat_len=0


# actual computation
for i in range(0,repeat_len+1):
  local_files=['' for k in range(len(args.files))]
  for j in range(0,len(args.files)):
    local_files[j]=list_file[j+len(args.files)*i]
  if args.interp is False:
    if args.merging=='blend_together':
      blend_together(local_files,spectrum_type,args.term)
    else:
      if len(local_files)<2:
	error_number_of_files()
      else:
	blend_against(local_files,spectrum_type,args.term)
  else:
    print '**interpolating (please wait)'
    blend_against_interp(local_files,spectrum_type,args.term)
  
exit()