#!/scisoft/i386/bin/python
#
#interp.py
#
#Given an input file consisting of
#wavelength, n and k, 
#interpolate the file out to the specified wavelength
#this does a cubic spline fit of the original data
#a linear fit (or remove loops now in doing special n&k fitting
#for the extrapolated region
#it also saves a pdf file of the resulting data for comparison

import os
import numpy,pylab,matplotlib
from scipy import interpolate
from scipy import polyfit, polyval
from optparse import OptionParser

__version__ = "1.0 (2007 June 12)"


def readData(fileName):
	try:
		f=open(fileName,'r')
		
	except IOError:
		print "Problem reading file:",fileName
		sys.exit(1)
		
	lines=f.readlines()
	chunk=lines[0].split()
	numCols=len(chunk)
	numRows=len(lines)
	
	dataArr=numpy.zeros((numRows,numCols),dtype=numpy.float32)
	
	for i in range(0,numRows,1):
		chunk=lines[i].split()
		for j in range(0,numCols,1):
			dataArr[i,j]=float(chunk[j])
	
	return dataArr


#these are just some default values for anyone calling
#the function from inside python instead of the csh prompt
def interpNK(inFile,outFile,endWave=4000,limit=True,fitpoints=3,startval=500):

	""" Useage(from shell prompt): interp.py -i file -w end wavelength 
	
		The code first uses a cubic spline to interpolate
		between the given points, then it extrapolates out
		to the specified wavelength
		
		startval is the wavelength to start the n&k limiting at (the paper starts
			at 500, but some of the dust prop files have actual data out farther
			than that
			
		fitpoints is the number of points at the end of the actual data to use
		for the linear extracpolation fit
		
		set LIMIT= FALSE, if you dont want the n&k limiting done		
	
		
	
	"""

	if(not(limit)):
		print "Not using the n&k 500 um limits\n"
	
	#read in the input file
	data=readData(inFile)
	wave=data[:,0]
	n=data[:,1]
	k=data[:,2]
	
	if( endWave < wave[-1]):
		print "your file extends past your specified end wavelength"
		print endWave, wave[-1]
		raise ValueError
	
	#now interpolate the points
	
	fitn=interpolate.splrep(wave,n,s=0)
	fitk=interpolate.splrep(wave,k,s=0)
	
	#extrapolate the new y's
	
#	use this for full extrapolation
#	newxe=numpy.arange(endWave) + wave[0] #give us a new range from the first wave to end wave
	newxe=wave

#	use this for linear fit
	moreWaves=(endWave - wave[-1])
	newx=numpy.arange(moreWaves) + wave[-1] + 1.

	print "min newx: ",min(newx)
	print "max newx: ",max(newx)


	""" This was an attempt at spline extrapolation, but it only
		works inside the original data range
	"""	
	newyn=interpolate.splev(newxe,fitn,der=0)
	newyk=interpolate.splev(newxe,fitk,der=0)
	
	
	#simple linear extrapolation using the last few points in the data 
	npts=fitpoints
	fitpointsx=wave[-(npts-1):]
	fitpointsyn=n[-(npts-1):]
	fitpointsyk=k[-(npts-1):]
	
	#get a linear fit to both n and k
	(ar,br)=polyfit(fitpointsx,fitpointsyn,1)
	fityn=polyval([ar,br],newx)

	(ar,br)=polyfit(fitpointsx,fitpointsyk,1)
	fityk=polyval([ar,br],newx)
	
	
	#now save out the file with new n and k values
	try:
		o=open(outFile,'w')
	except:
		print "Error opening output file: ",outFile
		sys.exit(0)

	results=numpy.zeros((len(newx)+len(newxe),3))

	#for spline fit section		
	for i in range(0,len(newxe),1):
		if(limit):
			#for limiting n and k vals , see Chiang, E. I. et al. 2001, ApJ, 547, 1077
			if(newxe[i] > startval): 
				loc=numpy.where(newxe >= startval)
				nval=newyn[loc[0][0]]
				kval=newyk[loc[0][0]]/(newxe[i]/startval)
			
				newline=str(newxe[i])+"\t"+str(nval)+"\t"+str(kval)+"\n"
				results[i,:]=[newxe[i],nval,kval]
			else:
				newline=str(newxe[i])+"\t"+str(newyn[i])+"\t"+str(newyk[i])+"\n"
				results[i,:]=[newxe[i],newyn[i],newyk[i]]
		else:
			newline=str(newxe[i])+"\t"+str(newyn[i])+"\t"+str(newyk[i])+"\n"
			results[i,:]=[newxe[i],newyn[i],newyk[i]]
			
		o.write(newline)
	
	start=i+1
	#for linear fit section
	
	for i in range(start,len(newx)+start,1):
	
		if(limit):
			#for limiting n and k vals in dorchner files, see Chiang, E. I. et al. 2001, ApJ, 547, 1077
			if( newx[i-start] >startval): 
				loc=numpy.where(newx >= startval)
				nval=fityn[loc[0][0]]
				kval=fityk[loc[0][0]]/(newx[i-start]/startval)

				newline=str(newx[i-start])+"\t"+str(nval)+"\t"+str(kval)+"\n"
				results[i,:]=[newx[i-start],nval,kval]
			else:
				newline=str(newx[i-start])+"\t"+str(fityn[i-start])+"\t"+str(fityk[i-start])+"\n"
				results[i,:]=[newx[i-start],fityn[i-start],fityk[i-start]]
		else:
			newline=str(newx[i-start])+"\t"+str(fityn[i-start])+"\t"+str(fityk[i-start])+"\n"
			results[i,:]=[newx[i-start],fityn[i-start],fityk[i-start]]

			
		o.write(newline)
		
	o.close()
	
	
	print "Saved extrapolated & fitted data to: ",outFile
	
	#now plot the results to see how we did

#	params={'font.size': 9,
#			'text.fontsize': 9}
#	pylab.rcParams.update(params)
	
	pylab.figure(1)
	pylab.axis('equal')
	pylab.figtext(0.2,0.95,outFile,fontsize=12,fontweight='bold')

	font={ 'fontsize': 11,
		   'fontname' : 'Times'}
	
	pylab.subplot(221)
	pylab.plot(wave,n)
	pylab.title('Original',fontsize=10)
	pylab.xlabel('wave',font)
	pylab.ylabel('n',font)
	#pylab.title('Original Data')
	
	pylab.subplot(222)
	pylab.plot(wave,n,'b.',results[:,0],results[:,1],'r-')
	pylab.title('Extrapolated',fontsize=10)
	pylab.xlabel('wave',font)
	pylab.ylabel('n',font)
	#pylab.title('Extrapolated n -cubic spline')
	
	pylab.subplot(223)
	pylab.title('Original',fontsize=10)
	pylab.plot(wave,k)
	pylab.xlabel('wave',font)
	pylab.ylabel('k',font)
	#pylab.title('Original Data')
	
	pylab.subplot(224)
	pylab.plot(wave,k,'b.',results[:,0],results[:,2],'r-')
	pylab.title('Extrapolated',fontsize=10)
	pylab.xlabel('wave',font)
	pylab.ylabel('k',font)
	#pylab.title('Extrapolated k - cubic spline')
	outplot=outFile+".pdf"
	pylab.savefig(outplot)
	pylab.show()
	
	print "Plot saved to: ",outplot
	return results

if __name__ == "__main__":

	#parse the command line
	usage = "usage: %prog [options]"
	parser = OptionParser(usage=usage)


	parser.add_option("-i","--infile",dest="inFile",default="INDEF",type="string",
		help="Input filename")
		
	parser.add_option("-o","--outfile",dest="outFile",default="INDEF",type="string",
		help="Output filename")
		
	parser.add_option("-w","--wave",dest="endWave",default=0,type="float",
		help="Ending wavelength")
		
	parser.add_option("-l","--nolimit",dest="nolimit",default=True,action="store_false",
		help="Dont limit the n&k values as  in the Chiang paper")
	
	parser.add_option("-f","--fitpoints",dest="fitpoints",default=3,type="int",
		help="Number of points at the end for the fit")
		
	parser.add_option("-s","--startval",dest="startval",default=500,type="float",
		help="Wavelength to start n&k limits at")	
		
	(options, args)  = parser.parse_args()


	#check the inputs
	
	if(not(os.access(options.inFile,os.F_OK))):
		print "Unable to access file: ",options.inFile
		raise ValueError
	
	if(os.access(options.outFile,os.F_OK)):
		print "Output file already exists: ",options.outFile
		raise ValueError
		
	if(options.endWave == 0):
		print "You did not specify an output wavelength"
		raise ValueError
		

	interpNK(options.inFile,options.outFile,options.endWave,options.nolimit,options.fitpoints,options.startval)
