#! /usr/bin/env python

import ppgplot, Numeric, math
import sys, string, fileinput
# import Scientific

#from Scientific.Functions.LeastSquares import leastSquaresFit
# Abbreviate the fit command
#fit=Scientific.Functions.LeastSquares.leastSquaresFit

# define linear function for use in fits
def linear(parameters, x):
    a=parameters[0]
    b=parameters[1]
    return a*x+b

# Function to work out avg, sx and sigmax from a list of numbers
def stat(numbers):
	sum=0
	# Calculate avg, sx and sigmax
	if len(numbers)>1:
	    	for i in numbers:
	    		    sum=sum+i
	    		    # Store total sum for later
	    		    total=sum
	    	avg=sum/len(numbers)
	    	sum=0
	    	for i in numbers:
	    		    sum=sum+(i-avg)**2
	    	sx=math.sqrt(sum/(len(numbers)-1))
	    	sigmax=math.sqrt(sum/len(numbers))
	else:
	    	avg=numbers[0]
	    	sx=0
	    	sigmax=0
	    	sum=numbers[0]
	# Calculate mean
	numbers.sort()
	if (len(numbers) % 2)==1:
	    	mean=numbers[(len(numbers)-1)/2]
	else:
	    	mean=(numbers[(len(numbers)/2)-1]+numbers[(len(numbers)/2)])/2
	# Return results
	return avg, sx, sigmax, total, mean

# print help
def print_help():
	print """
	\n Plotit.py by Enno Middelberg, October 2003
        \n Your choice:\n
        h  - print this message
        ---- Flagging ----
	a  - delete above
	b  - delete below
	l  - delete to left
	r  - delete to right
	q  - delete rectangle
	k  - keep rectangle (delete outside)
	o  - delete nearest point
        ---- Analyze data ----
	p  - print nearest point
        t  - statistics of points in rectangle
	i  - plot histogram of x data
	j  - plot histogram of y data
	f  - fit linear function to data
        ---- Change plot ----
	s  - re-scale
	z  - zoom in
	<> - change plot symbol
	x  - exit"""

# return coordinates of rectangle drawn by user
# arguments: anchor of rectangle (first corner's coords)
def get_rectangle(x1, y1):
	action=ppgplot.pgband(2, 0, x1, y1)
	x2=action[0]
	y2=action[1]
	if x2<x1:
		i=x1
		x1=x2
		x2=i
	if y2<y1:
		i=y1
		y1=y2
		y2=i
        return (x1, x2, y1, y2)

# return stats of points inside a rectangle
def rectangle_stats(xarray, yarray, x1, x2, y1, y2):
	xcopy=[]
	ycopy=[]
        for i in range(len(xarray)):
		if x1<=xarray[i]<=x2 and y1<=yarray[i]<=y2:
			xcopy.append(xarray[i])
			ycopy.append(yarray[i])
        n=len(xcopy)
	xstat=stat(xcopy)
        ystat=stat(ycopy)
        return(xstat, ystat, n)
        

# Function to plot and edit data
def plotit(xarray, yarray, xlabel, ylabel, title):
        # print help
        print_help()
        print "\n Processing %i points" % len(xarray)
	action=(0,0, "s")
	ppgplot.pgopen("/xwin")
	ppgplot.pgask(0)
	pttype=2
 	while not (action[2]=="x" or action[2]=="X"):
		# get array statistics
		xlist=xarray.tolist()
		ylist=yarray.tolist()
		xstat=stat(xlist)
		ystat=stat(ylist)
		# compute error rectangle corners
		rect_xmin=xstat[0]-xstat[1]
		rect_xmax=xstat[0]+xstat[1]
		rect_ymin=ystat[0]-ystat[1]
		rect_ymax=ystat[0]+ystat[1]
		if action[2]=="s":
			# find plot boundaries - xmin and xmax
			xmin=xarray[0]
			xmax=xarray[0]
			for i in xarray:
				if i<xmin:
					xmin=i
				if i>xmax:
					xmax=i
			# find plot boundaries - ymin and ymax
			ymin=yarray[0]
			ymax=yarray[0]
			for i in yarray:
				if i<ymin:
					ymin=i
				if i>ymax:
					ymax=i
			# Store data for later use
			dataxmin=xmin
			dataxmax=xmax
			dataymin=ymin
			dataymax=ymax
			# add and subtract 5% as extra margin
			xmin=xmin-0.05*(xmax-xmin)
			xmax=xmax+0.05*(xmax-xmin)
			# add and subtract 5% as extra margin
			ymin=ymin-0.05*(ymax-ymin)
			ymax=ymax+0.05*(ymax-ymin)
			print "\n New plot dimensions: xmin=%.1f, xmax=%.1f, ymin=%.1f, ymax=%.1f" % (xmin, xmax, ymin, ymax)

		# initialize plot window
		ppgplot.pgenv(xmin, xmax, ymin, ymax, 0, 1)
		ppgplot.pglab(xlabel, ylabel, title)

		# plot data
		ppgplot.pgpt(xarray, yarray, pttype)

		# plot hollow 1 sigma rectangle with cross in it
		ppgplot.pgsfs(2)
		ppgplot.pgrect(rect_xmin, rect_xmax, rect_ymin, rect_ymax)

		# start flagging
		action=ppgplot.pgband(7, 0, 0, 0)

		# report nearest data point
		if action[2]=="p" or action[2]=="P":
			x1=action[0]
			y1=action[1]
			min=math.sqrt((xmax-xmin)**2+(ymax-ymin)**2)
			for i in range(len(xarray)):
				dist=math.sqrt((xarray[i]-x1)**2+(yarray[i]-y1)**2)
				if dist<min:
					min=dist
					index=i
			print "\n Nearest data point: x=%f  y=%f" % (xarray[index], yarray[index])
		# delete above horizontal line
		if action[2]=="a":
			ylim=action[1]
			print "\n Flag above y=%f" % ylim
			xcopy=[]
			ycopy=[]
			for i in range(len(yarray)):
				if yarray[i]<=ylim:
					xcopy.append(xarray[i])
					ycopy.append(yarray[i])
			xarray=Numeric.array(xcopy)
			yarray=Numeric.array(ycopy)
		# delete below horizontal line
		if action[2]=="b":
			ylim=action[1]
			print "\n Flag below y=%f" % ylim
			xcopy=[]
			ycopy=[]
			for i in range(len(yarray)):
				if yarray[i]>=ylim:
					xcopy.append(xarray[i])
					ycopy.append(yarray[i])
			xarray=Numeric.array(xcopy)
			yarray=Numeric.array(ycopy)
		# delete right of vertical line
     		if action[2]=="r":
			xlim=action[0]
			print "\n Flag right of x=%f" % xlim
			xcopy=[]
			ycopy=[]
			for i in range(len(xarray)):
				if xarray[i]<=xlim:
					xcopy.append(xarray[i])
					ycopy.append(yarray[i])
			xarray=Numeric.array(xcopy)
			yarray=Numeric.array(ycopy)
		# delete left of vertical line
		if action[2]=="l":
			xlim=action[0]
			print "\n Flag left of x=%f" % xlim
			xcopy=[]
			ycopy=[]
			for i in range(len(xarray)):
				if xarray[i]>=xlim:
					xcopy.append(xarray[i])
					ycopy.append(yarray[i])
			xarray=Numeric.array(xcopy)
			yarray=Numeric.array(ycopy)
		# delete inside or outside of rectangle
		if action[2]=="q" or action[2]=="k":
                        key=action[2]
                        (x1, x2, y1, y2)=get_rectangle(action[0], action[1])
			print "\n Limits: x=%f and x=%f" % (x1, x2)
			print "         y=%f and y=%f" % (y1, y2)
			xcopy=[]
			ycopy=[]
			if key=="q":
				for i in range(len(xarray)):
					if not (x1<=xarray[i]<=x2 and y1<=yarray[i]<=y2):
						xcopy.append(xarray[i])
						ycopy.append(yarray[i])
			if key=="k":
				for i in range(len(xarray)):
					if x1<=xarray[i]<=x2 and y1<=yarray[i]<=y2:
						xcopy.append(xarray[i])
						ycopy.append(yarray[i])
			xarray=Numeric.array(xcopy)
			yarray=Numeric.array(ycopy)
			# clear "action" or last click it will be evaluated in next loop
			action=(0,0,"")
		# zoom in
		if action[2]=="z":
			(xmin, xmax, ymin, ymax)=get_rectangle(action[0], action[1])
			print "\n New plot dimensions: xmin=%.1f, xmax=%.1f, ymin=%.1f, ymax=%.1f" % (xmin, xmax, ymin, ymax)
			# clear "action" or last click it will be evaluated in next loop
			action=(0,0,"")
		# decrease plot symbol number
		if action[2]=="<":
			pttype=pttype-1
		# increase plot symbol number
		if action[2]==">":
			pttype=pttype+1
		# delete nearest point (click left=="A")
		if action[2]=="o" or action[2]=="A":
			x1=action[0]
			y1=action[1]
			min=math.sqrt((xmax-xmin)**2+(ymax-ymin)**2)
			for i in range(len(xarray)):
				dist=math.sqrt((xarray[i]-x1)**2+(yarray[i]-y1)**2)
				if dist<min:
					min=dist
					index=i
			xcopy=[]
			ycopy=[]
			for i in range(len(xarray)):
				if i!=index:
					xcopy.append(xarray[i])
					ycopy.append(yarray[i])
			xarray=Numeric.array(xcopy)
			yarray=Numeric.array(ycopy)
                # stats inside rectangle
                if action[2]=="t":
                        (x1, x2, y1, y2)=get_rectangle(action[0], action[1])
                        (xstat, ystat, n)=rectangle_stats(xarray, yarray, x1, x2, y1, y2)
                        print "\n Statistics of %i points in rectangle xmin=%.1f, xmax=%.1f, ymin=%.1f, ymax=%.1f" % (n, x1, x2, y1, y2)
                        print "\n x: avg(x)=%f, mean=%f, sx=%f, sum=%f" % (xstat[0], xstat[4], xstat[1], xstat[3])
                        print " y: avg(x)=%f, mean=%f, sx=%f, sum=%f" % (ystat[0], ystat[4], ystat[1], ystat[3])
                # stats inside viewport
                if action[2]=="v":
                        (xstat, ystat, n)=rectangle_stats(xarray, yarray, xmin, xmax, ymin, ymax)
                        print "\n Statistics of %i points in viewport xmin=%.1f, xmax=%.1f, ymin=%.1f, ymax=%.1f" % (n, xmin, xmax, ymin, ymax)
                        print "\n x: avg(x)=%f, mean=%f, sx=%f, sum=%f" % (xstat[0], xstat[4], xstat[1], xstat[3])
                        print " y: avg(x)=%f, mean=%f, sx=%f, sum=%f" % (ystat[0], ystat[4], ystat[1], ystat[3])
                # histogram of x
		if action[2]=="i":
			inp=raw_input("\n Number of bins: ")
			nbins=int(inp)
			ppgplot.pgeras()
			length=len(xarray)
			ppgplot.pghist(length, xarray, dataxmin, dataxmax, nbins, 0)
			ppgplot.pgband(2,0,0,0)
                # histogram of y
		if action[2]=="j":
			inp=raw_input("\n Number of bins: ")
			nbins=int(inp)
			ppgplot.pgeras()
			length=len(yarray)
			ppgplot.pghist(length, yarray, dataymin, dataymax, nbins, 0)
			ppgplot.pgband(2,0,0,0)
                # linear fit to data
		if action[2]=="f":
			data=[]
			x=xarray.tolist()
			y=yarray.tolist()
			for i in range(len(x)):
				data.append([x[i], y[i]])
			a=fit(linear, (1, 1), data)
			print a
			print "\n Linear fit is: f(x)=%f * x + %f" % (a[0][0], a[0][1])
			ylist=[]
			for x in xarray:
				ylist.append(linear((a[0][0], a[0][1]), x))
			yfunc=Numeric.array(ylist)
			ppgplot.pgline(xarray, yfunc)
			ppgplot.pgband(2,0,0,0)
                # print help
                if action[2]=="h":
                        print_help()
	return xarray, yarray

		
# read data from file
data1=[]
data2=[]
for x in fileinput.input(sys.argv[1]):
	list=string.split(x)
	data1.append(float(list[0]))
	data2.append(float(list[1]))
x=Numeric.array(data1)
y=Numeric.array(data2)

# plot and edit data, merge x and y arrays into data
data=Numeric.array(plotit(x,y,"x", "y", sys.argv[1]))

# swap axis for output
data_out=Numeric.swapaxes(data, 0, 1)

# print data and final stats
for i in range(len(data_out[:,0])):
	print "%f   %f" % (data_out[i][0], data_out[i][1])
print "\n %i data points left. Statistics:" % len(data_out[:,0])

# Slice through array, convert to list and compute final stats
xstat=stat(data_out[:,0].tolist())
print "\n x: avg(x)=%f, mean=%f, sx=%f, sum=%f" % (xstat[0], xstat[4], xstat[1], xstat[3])
ystat=stat(data_out[:,1].tolist())
print " y: avg(x)=%f, mean=%f, sx=%f, sum=%f" % (ystat[0], ystat[4], ystat[1], ystat[3])

print