#!/usr/bin/python

import sys
import os.path
import re
import time
import numpy as np

import FormFactor

###########################################################################################
###########################################################################################
###########################################################################################
###########################################################################################

###############
############### Miscellaneous

def file_exists(arg, name):
	if os.path.exists(arg):
		print(name,":\t",arg)
	else:
		raise IOError(arg,"The file does not exist")

###############

def check_int(strg, name):
	try:
		return int(strg)
	except ValueError:
		raise ValueError(name," must be an int number")

###############

def check_binary(arg, name):
	if arg ==0 or arg == 1:
		pass
	else:
		raise Exception(name," must be 0 or 1")

###############

def check_float_positive(strg, name):
	try:
		float(strg)
	except ValueError:
		raise ValueError(name," must be a float number")
	if float(strg) >= 0:
		return float(strg)
	else:
		raise Exception(name," must be a positive number")

###########################################################################################
###########################################################################################
###########################################################################################
###########################################################################################

class Reading:
	'Read the Options & Parameters file'
	EMPTY=[]

	def __init__ ( self , ARGV ):

		if len(ARGV)>3:
			sys.exit("--- Too many arguments")
		elif len(ARGV)==1:
			sys.exit("--- Specify the parameters file")

		file_exists(ARGV[1], "Parameters File")

		with open(ARGV[1], 'r') as f:
			LINES = f.read().splitlines()
		
		TEMP=[]
		OPT=[]
		PARA=[]
		for i in range(len(LINES)):
			if not LINES[i]=='':
				TEMP.append(LINES[i])
			else:	
				TEMP.append("#")
				Reading.EMPTY.append(i)

		line = 0
		while line < Reading.EMPTY[0]:
			OPT.append(TEMP[line].split())
			line += 1
		line = Reading.EMPTY[1]+1
		while line <= len(TEMP)-1:
			PARA.append(TEMP[line].split())
			line += 1

		self.OPT=OPT
		self.PARA=PARA

###########################################################################################
###########################################################################################

	def Options ( self ):
		'Read Options'	

		err_str =  "\n--- Something missing or wrong.\n--- You need:\n"
		err_str	+= "DATA_FILE:		<filename>\n"
		err_str	+= "TITLE:			string\n"
		err_str	+= "SAVE_FOLDER:		string\n"
		err_str	+= "q_MIN(nm^-1):		float\n"
		err_str	+= "q_MAX(nm^-1):		float\n"
		err_str	+= "MODEL:			string\n"
		err_str	+= "TEMPERATURE(none=0):	float or 0\n"
		err_str	+= "BINNING:		4\n"
		err_str	+= "ERROR(0..1):		float\n"
		err_str	+= "PLOT(0-1):		int 0/1"
	
		options={}

		for i in range(len(self.OPT)):
			if re.match( "^DATA_FILE:$", self.OPT[i][0] ):
				options['datafile'] = str(self.OPT[i][1])
				file_exists(options['datafile'], "Data File")
			
			elif re.match( "^TITLE:$", self.OPT[i][0] ):
				options['title'] = str(self.OPT[i][1])
				print("Title:\t\t", options['title'])

			elif re.match( "^SAVE_FOLDER:$", self.OPT[i][0] ):
				options['folder'] = str(self.OPT[i][1])
				if os.path.exists(options['folder'])==0:
					os.makedirs(options['folder'])
					print("--- Creating new folder: ", options['folder'])
					print("Saving Folder:\t", options['folder'])
				elif os.path.exists(options['folder'])==1:
					print("Saving Folder:\t", options['folder'])

			elif re.match( "^q_MIN\(nm\^-1\):$", self.OPT[i][0] ):
				options['qmin'] = check_float_positive(self.OPT[i][1], "qmin")
				print("q range:\t(%s .." %options['qmin'],	)

			elif re.match( "^q_MAX\(nm\^-1\):$", self.OPT[i][0] ):
				options['qmax'] = check_float_positive(self.OPT[i][1], "qmax")
				print("%s) nm^-1\n" %options['qmax'],	)

			elif re.match( "^MODEL:$", self.OPT[i][0] ):
				options['function'] = str(self.OPT[i][1])
				print("Function:\t", options['function'])

			elif re.match( "^TEMPERATURE\(none=0\):$", self.OPT[i][0] ):
				options['temp'] = check_float_positive(self.OPT[i][1], "Temperarure")
				if options['temp'] != 0:
					print("Temperarure:\t", options['temp'])
				else:
					print("Temperarure:\tNo Simulated Annealing")

			elif re.match( "^BINNING:$", self.OPT[i][0] ):
				options['bing'] = check_int(self.OPT[i][1], "Binning")
				print("Binning:\t", options['bing'])

			elif re.match( "^ERROR\(0\.\.1\):$", self.OPT[i][0] ):
				options['err_mul'] = check_float_positive(self.OPT[i][1], "Error Mult.")
				print("Error Mult.:\t", options['err_mul'])

			elif re.match( "^PLOT\(0-1\):$", self.OPT[i][0] ):
				options['plot'] = check_int(self.OPT[i][1], "Plot")
				check_binary(options['plot'], "Plot")
				print("Plot:\t\t", options['plot'])

			else:
				raise Exception(err_str)

		if len(options)!=10:
			raise Exception(err_str)

		return options


###########################################################################################
###########################################################################################

	def Parameters ( self ):
		'Read Parameters'	

		############### Define parameters arrays
		NAME=[]
		PAR=[]
		FIX=[]
		PRIOR=[]		
		L_LIM=[]
		H_LIM=[]
		for i in range(len(self.PARA)):
			if not re.match( "^#" , self.PARA[i][0] ):
				NAME.append(self.PARA[i][0])
				PAR.append(float(self.PARA[i][1]))
				FIX.append(self.PARA[i][2])
				PRIOR.append(self.PARA[i][3])				
				L_LIM.append(float(self.PARA[i][4]))
				H_LIM.append(float(self.PARA[i][5]))
	
		############### Check if there are free parameters
		n=0
		for i in range(len(FIX)):
			if FIX[i]=="f":
				n += 1
		if n==len(PAR):
			sys.exit("--- All the parameters are fixed; Choose the \"Plot\" option")

		############### Check & Print parameters features
		for i in range(len(NAME)):
			if FIX[i]=="f":
				f="Fixed"
				print(i,"\t%4s\t%.3e\t%s" %(NAME[i],PAR[i],f))
			else:
				if PAR[i]==0: 
					sys.exit("--- 0 value as free parameter is not allowed")
				if FIX[i]=="-" and PRIOR[i]=="-":
					f="Free"
					print(i,"\t%4s\t%.3e\t%s\t\thard boundaries (%.3e, %.3e)" %(NAME[i],PAR[i],f,L_LIM[i],H_LIM[i]))
				elif FIX[i]=="-" and PRIOR[i]!="-":
					f="Prior"
					print(i,"\t%4s\t%.3e\t%s %.5s\thard boundaries (%.3e, %.3e)" %(NAME[i],PAR[i],f,PRIOR[i],L_LIM[i],H_LIM[i]))					
				else:
					sys.exit("--- Free/Fixed -> -/f")
			
			if H_LIM[i]-L_LIM[i]<=0:
				sys.exit("--- Check low & high parameter boundary")
			if PAR[i]<L_LIM[i]:
				sys.exit("--- Check low & high parameter boundary")
			if PAR[i]>H_LIM[i]:
				sys.exit("--- Check low & high parameter boundary")
	
		return ( NAME , PAR , FIX, PRIOR, L_LIM , H_LIM )

###########################################################################################

	def Parameters_refresh ( self ):
		'Read Parameters'	

		############### Define parameters arrays
		NAME=[]
		PAR=[]
		FIX=[]
		PRIOR=[]		
		L_LIM=[]
		H_LIM=[]
		for i in range(len(self.PARA)):
			if not re.match( "^#" , self.PARA[i][0] ):
				NAME.append(self.PARA[i][0])
				PAR.append(float(self.PARA[i][1]))
				FIX.append(self.PARA[i][2])
				PRIOR.append(self.PARA[i][3])				
				L_LIM.append(float(self.PARA[i][4]))
				H_LIM.append(float(self.PARA[i][5]))
	
		return ( NAME , PAR , FIX, PRIOR, L_LIM , H_LIM )
		
###########################################################################################
###########################################################################################
###########################################################################################
###########################################################################################

class ReadData:
	"""Read raw data from file. Format:
        1st column: q value
        2nd column: I value
        3rd column: error
    """


	def __init__ ( self , datafile ):

		DATA = []

		with open(datafile, 'r') as d:
			DLINES = d.read().splitlines()
		for row in DLINES:
			DATA.append(row.split())

		DATA.pop(0)

		for i in range(len(DATA)):
			for j in range(len(DATA[i])):
				DATA[i][j] = float(DATA[i][j])
	
		self.DATA = DATA

###########################################################################################
###########################################################################################
	
	def SelectColumns ( self , qmin , qmax , bing , err_mul ):
		'Select q / I(q) / err(q) data	-	Change I(q) value from mm^-1 to nm^-1'

		Q = []
		IDATA = []
		ERR = []

		for i in range(len(self.DATA)):
			if ( self.DATA[i][0] >= qmin and self.DATA[i][0] <= qmax ) :
				if i%bing == 0 :
					Q.append(self.DATA[i][0])
					if self.DATA[i][1] > 0 :
						IDATA.append(1e-6*self.DATA[i][1]) 
						if err_mul == 0 :
							ERR.append(0.01*1e-6*self.DATA[i][1])
						elif self.DATA[i][2]*err_mul > 0.01*self.DATA[i][1] :
							ERR.append(1e-6*self.DATA[i][2]*err_mul)
						else:
							ERR.append(0.01*1e-6*self.DATA[i][1]) 
					else:
						Q.pop()	

		return ( np.array(Q), np.array(IDATA), np.array(ERR) )

###########################################################################################
###########################################################################################
###########################################################################################
###########################################################################################

class PlotData:
	'Plot Results'

	def __init__ ( self , title , folder , qmin , qmax ):

		self.fitFile="./"+folder+"/"+title+".fit"
		self.gnuFile="./"+folder+"/"+title+".gp"
		self.intensityFile="./"+folder+"/"+title+"_I.pdf"
		self.structureFile="./"+folder+"/"+title+"_Sq.pdf"
		self.errorFile="./"+folder+"/"+title+"_err.pdf"
		self.qmin=str(qmin)
		self.qmax=str(qmax)
		self.gnuStr0= str( "set term pdfcairo enhanced solid color font \"Helvetica,14\" \
						lw 2 rounded size 5 inch, 3.5 inch\n \
						set encoding iso_8859_1\n \
						set grid\n \
						set key box\n \
						set bars small\n \
						set xtics add ("+str(qmin)+","+str(qmax)+")\n \
						set ytics format \"10^\{%T\}\"\n \
						set mytics 10\n \
						set title \""+title+"\"\n" )

###########################################################################################
###########################################################################################

	def Form ( self, Q, IDATA, ERR, I, PAR, pltOptions ):
		'Plot the best fit of a Form Factor'

		############### Check Cluster Structure Factor
		if pltOptions['debyeCtr'] != 0:

			############### Save constant value
			if conv == 0:
				C = str(1e6*PAR[-1])
			else:
				C = str(1e6*PAR[-2])
			
			############### Prepare Form Factor
			PAR_FORM=list(PAR)
			for i in range(pltOptions['parDiff']):
				PAR_FORM.pop(-1)
			PAR_FORM.append(0.0)
			if conv == 1:
				PAR_FORM.append(PAR[-1])
			FORM = [ pltOptions['Form'](q, PAR_FORM) for q in Q ]
			
			############### Prepare Cluster Structure Factor
			if conv == 0:
				Io =	PAR[-4]
				Rg =	PAR[-3]
				D =		PAR[-2]
			else:
				Io =	PAR[-5]
				Rg =	PAR[-4]
				D =		PAR[-3]	
			DEBYE = [ pltOptions['debyeCtr'](q,Io,Rg,D) for q in Q ]
			
		############### Write raw data file
		fit = open ( self.fitFile, "w" )
		if pltOptions['debyeCtr'] == 0:
			for i in range(len(Q)):
				fit.write( "%f\t%f\t%f\t%f\n" %( Q[i], 1e6*IDATA[i], 1e6*ERR[i], 1e6*I[i] ) )
		else:
			for i in range(len(Q)):
				fit.write( "%f\t%f\t%f\t%f\t%f\t%f\n" %( Q[i], 1e6*IDATA[i], 1e6*ERR[i], 1e6*I[i], 1e6*FORM[i], DEBYE[i] ) )
		fit.close()

		############### Write GNUPLOT script
		gnu = open ( self.gnuFile, "w" )
		gnu.write( self.gnuStr0 )

		gnu.write( "set logscale\n" )
		gnu.write( "set xlabel \"q (nm^{-1})\"\n" )
		gnu.write( "set ylabel \"I (mm^{-1})\"\n" )
		gnu.write( "set out '"+self.intensityFile+"'\n" )
		gnu.write( "p ["+self.qmin+":"+self.qmax+"][] '"+self.fitFile+"' \
						u 1:2 w p pt 6 ps 0.5 lt -1 t '', '' u 1:2:3 w errorbars pt 7 ps 0.3 lt 1 t 'data', \
						'' u 1:4 w l lt 3 lw 2.5 t 'Fit'\n" )
		
		gnu.write( "set ylabel \"| I^{data}-I^{fit} |/I^{data} (mm^{-1})\"\n" )
		gnu.write( "set out '"+self.errorFile+"'\n" )
		gnu.write( "p ["+self.qmin+":"+self.qmax+"][] '"+self.fitFile+"' \
					u 1:($3/$2) w l lt rgb \"web-green\" lw 2.5 t 'err.', \
					'' u 1:(abs($2-$4)/$2) w p pt 7 ps 0.6 lt 3 t '', \
					'' u 1:(abs($2-$4)/$2) w p pt 6 ps 0.6 lt -1 t ''\n" )

		if pltOptions['debyeCtr'] != 0:
			gnu.write( "unset logscale y\n" )
			gnu.write( "set xtics auto\n" )
			gnu.write( "set xtics add ("+self.qmax+")\n" )
			gnu.write( "set ylabel \"S(q)\"\n" )
			gnu.write( "set out '"+self.structureFile+"'\n" )
			gnu.write( "p ["+self.qmin+":"+self.qmax+"][] '"+self.fitFile+"' \
						u 1:(($2-"+C+")/$5) w p pt 6 ps 0.5 lt -1 t '', \
						'' u 1:(($2-"+C+")/$5) w p pt 7 ps 0.3 lt 1 t 'data', \
						'' u 1:(1+$6) w l lt rgb \"purple\" lw 2.5 t '1+D(q)'\n" )
						

		gnu.close()
	
		############### Run GNUPLOT script
		os.system("gnuplot "+self.gnuFile)		
		############### Remove GNUPLOT script
		os.remove(self.gnuFile)
		############### Print saved files
		print(self.fitFile)
		print(self.intensityFile)
		print(self.errorFile)

		return

###########################################################################################
###########################################################################################

	def Form_Porod ( self , Q , IDATA , ERR , I , PAR ):
		'Plot the best fit of a Form Factor'

		############### Porod function values
		if conv==0:
			Pval =	PAR[-3]
			Pexp =	-PAR[-2]
		else:
			Pval =	PAR[-4]
			Pexp =	-PAR[-3]
	
		############### Write raw data file
		fit = open ( self.fitFile, "w" )
		for i in range(len(Q)):
			fit.write( "%f\t%f\t%f\t%f\t%f\n" %( Q[i], 1e6*IDATA[i], 1e6*ERR[i], 1e6*I[i],  1e6*Pval*Q[i]**Pexp ) )
		fit.close()

		############### Write GNUPLOT script
		gnu = open ( self.gnuFile, "w" )
		gnu.write( self.gnuStr0 )

		gnu.write( "set logscale\n" )
		gnu.write( "set xlabel \"q (nm^{-1})\"\n" )
		gnu.write( "set ylabel \"I (mm^{-1})\"\n" )
		gnu.write( "set out '"+self.intensityFile+"'\n" )
		gnu.write( "p ["+self.qmin+":"+self.qmax+"][] '"+self.fitFile+"' \
					u 1:2 w p pt 7 ps 0.6 lt -1 t '', '' u 1:2:3 w errorbars pt 7 ps 0.5 lt 1 t 'data', \
					'' u 1:4 w l lt 3 lw 2.5 t 'Fit', \
					'' u 1:5 w l lt rgb \"web-green\" lw 2 t 'Porod', \
					'' u 1:($4-$5) w l lt 0 lw 3 t ''\n" )

		gnu.write( "set ylabel \"| I^{data}-I^{fit} |/I^{data} (mm^{-1})\"\n" )
		gnu.write( "set out '"+self.errorFile+"'\n" )
		gnu.write( "p ["+self.qmin+":"+self.qmax+"][] '"+self.fitFile+"' \
					u 1:($3/$2) w l lt -1 lw 2.5 t 'err.', \
					'' u 1:(abs($2-$4)/$2) w p pt 6 ps 0.6 lt 1 t ''\n" )

		gnu.close()
	
		############### Run GNUPLOT script
		os.system("gnuplot "+self.gnuFile)		
		############### Remove GNUPLOT script
		os.remove(self.gnuFile)
		############### Print saved files
		print(self.fitFile)
		print(self.intensityFile)
		print(self.errorFile)

		return

###########################################################################################
###########################################################################################

	def Structure (self, Q, IDATA, ERR, I, PAR, pltOptions):
		'Plot the best fit of a Form Factor'

		############### Save constant value
		if conv == 0:
			C = str(1e6*PAR[-1])
		else:
			C = str(1e6*PAR[-2])

		############### Prepare Form Factor
		PAR_FORM=list(PAR)
		for i in range(pltOptions['parDiff']):
			PAR_FORM.pop(-1)
		PAR_FORM.append(0.0)
		if conv == 1:
			PAR_FORM.append(PAR[-1])
		FORM = [ pltOptions['Form'](q, PAR_FORM) for q in Q ]

		############### Check Cluster Structure Factor
		if pltOptions['debyeCtr'] != 0:
			############### Prepare Cluster Structure Factor
			if conv == 0:
				Io =	PAR[-4]
				Rg =	PAR[-3]
				D =		PAR[-2]
			else:
				Io =	PAR[-5]
				Rg =	PAR[-4]
				D =		PAR[-3]	
			DEBYE = [ pltOptions['debyeCtr'](q,Io,Rg,D) for q in Q ]
	
		############### Write raw data file
		fit = open ( self.fitFile, "w" )
		if pltOptions['debyeCtr'] == 0:
			for i in range(len(Q)):
				fit.write( "%f\t%f\t%f\t%f\t%f\n" %( Q[i], 1e6*IDATA[i], 1e6*ERR[i], 1e6*I[i], 1e6*FORM[i] ) )
		else:
			for i in range(len(Q)):
				fit.write( "%f\t%f\t%f\t%f\t%f\t%f\n" %( Q[i], 1e6*IDATA[i], 1e6*ERR[i], 1e6*I[i], 1e6*FORM[i], DEBYE[i] ) )
		fit.close()

		############### Write GNUPLOT script
		gnu = open ( self.gnuFile, "w" )
		gnu.write( self.gnuStr0 )

		gnu.write( "set logscale\n" )
		gnu.write( "set xlabel \"q (nm^{-1})\"\n" )
		gnu.write( "set ylabel \"I (mm^{-1})\"\n" )
		gnu.write( "set out '"+self.intensityFile+"'\n" )
		gnu.write( "p ["+self.qmin+":"+self.qmax+"][] '"+self.fitFile+"' \
					u 1:2 w p pt 6 ps 0.5 lt -1 t '', '' u 1:2:3 w errorbars pt 7 ps 0.3 lt 1 t 'data', \
					'' u 1:5 w l lt rgb \"web-green\" lw 2.5 t 'FormFactor', \
					'' u 1:4 w l lt 3 lw 2.5 t 'Fit'\n" )
		
		gnu.write( "set ylabel \"| I^{data}-I^{fit} |/I^{data} (mm^{-1})\"\n" )
		gnu.write( "set out '"+self.errorFile+"'\n" )
		gnu.write( "p ["+self.qmin+":"+self.qmax+"][] '"+self.fitFile+"' \
					u 1:($3/$2) w l lt rgb \"web-green\" lw 2.5 t 'err.', \
					'' u 1:(abs($2-$4)/$2) w p pt 7 ps 0.6 lt 3 t '', \
					'' u 1:(abs($2-$4)/$2) w p pt 6 ps 0.6 lt -1 t ''\n" )

		gnu.write( "unset logscale\n" )
		gnu.write( "set xtics auto\n" )
		gnu.write( "set xtics add ("+self.qmax+")\n" )
		gnu.write( "set ylabel \"S(q)\"\n" )
		gnu.write( "set out '"+self.structureFile+"'\n" )
		if pltOptions['debyeCtr'] == 0:
			gnu.write( "p ["+self.qmin+":"+self.qmax+"][] '"+self.fitFile+"' \
						u 1:(($2-"+C+")/$5) w p pt 6 ps 0.5 lt -1 t '', \
						'' u 1:(($2-"+C+")/$5) w p pt 7 ps 0.3 lt 1 t 'data', \
						'' u 1:(($4-"+C+")/$5) w l lt 3 lw 2.5 t 'S(q)'\n" )
		else:
			gnu.write( "p ["+self.qmin+":"+self.qmax+"][] '"+self.fitFile+"' \
						u 1:(($2-"+C+")/$5) w p pt 6 ps 0.5 lt -1 t '', \
						'' u 1:(($2-"+C+")/$5) w p pt 7 ps 0.3 lt 1 t 'data', \
						'' u 1:(1+$6) w l lt rgb \"purple\" lw 2.5 t '1+D(q)', \
						'' u 1:(($4-"+C+")/$5-$6) w l lt rgb \"dark-green\" lw 2.5 t 'S(q)', \
						'' u 1:(($4-"+C+")/$5) w l lt 3 lw 2.5 t 'S(q)+D(q)'\n" )

		gnu.close()
	
		############### Run GNUPLOT script
		os.system("gnuplot "+self.gnuFile)		
		############### Remove GNUPLOT script
		os.remove(self.gnuFile)
		############### Print saved files
		print(self.fitFile)
		print(self.intensityFile)
		print(self.errorFile)
		print(self.structureFile)

		return





###########################################################################################
###########################################################################################
###########################################################################################
###########################################################################################