import os
import time
import numpy
import sys
from os.path import basename
#from pylab import *
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, RadioButtons, Button

# This python program will plot the temperature or (if available) the 
# H-filed dependence of each data point in a BT-1 file.
# It is appropriate ONLY for those times when sitting on
# a particular Bragg peak and scanning T or H.
#
# The chosen detector (via the slider) dependence against T or H (if a BH1 file is 
# available) is plotted. Data can be saved as det#.dat with the save button.
#
# Warning - If there is random lines on the plot- you do not have the right type of 
# measurement to make this plotting program useful
#
# Craig M Brown 10/1/2015
# NIST Center for Neutron Research


f = sys.argv[1]
file = open(f)
detectors = []
columns = []
monitors = []

# read data
for i in range(0,18):
	# skip headers we dont care about those for this
	line = file.readline().strip()
	i +=1
# Get column names from
#     $ M4=     11.000  T= 250.770 C=   0.81 N=           0
column_names = [s.split()[-1] for s in line.split('=')[:-1]]
column_number = dict((c,i) for i,c in enumerate(column_names))
while '=' in line:
	fields = [s.split()[0] for s in line.split('=')[1:]]
	#print fields
	columns.append([float(s) for s in fields])
	#print columns
	line = file.readline()
	if line.startswith('$'):
		monitors.append([int(s) for s in line[1:].split()])
		line = file.readline()
	parts = []
	while line.startswith(' '):
		parts.append(line)
		line = file.readline()
	joined_line = "".join(s.strip() for s in parts)
	if joined_line.endswith(','): joined_line = joined_line[:-1]
	detectors.append([int(s) for s in joined_line.split(',')])
counts = numpy.array(detectors)
monitors = numpy.array(monitors)
columns = numpy.array(columns)
detectors = numpy.array(detectors)


Data_points =  len(columns)
num_dets =  len(detectors[0])

# extract a column for detectors angles, temperatures depending on the run type:
if column_names[0] == 'M4':
	angles = columns[:,0]
elif column_names[0] == 'T':
	temps = columns[:,0]
if column_names[1] == 'T':
	temps = columns[:,1]

#Is there a field file?
basen =  os.path.splitext(f)[0]
Hfile = basen+".bH1"
try:
	fileH = open(Hfile)
	Hfile = 1
	Hfield =[]
	print "we have a H file"
except:
	Hfile = 0
	print "we dont have a H file"
		
if Hfile == 1:
	for i in range(0,3):
		# skip headers we dont care about those for this
		line = fileH.readline().strip()
	while line != '' and ord(line[0]) != 0:
		try:
			line = fileH.readline().strip()
			Hfield.append(line.split()[1])
		except:
			pass
			
# assume what to plot.
# always a detector ct, and T or H.

det=0
d0 = 1
#fig, ax = plt.subplots()
fig = plt.figure()
ax = fig.add_subplot(111)
plt.subplots_adjust(left=0.15, bottom=0.25)
lx, = plt.plot(temps, detectors[:,d0-1], linewidth=2, color='r')
plt.ylabel('Counts')
plt.xlabel('Temperature')
plt.grid(True)
#ax.autoscale_view()
axes = plt.gca()
#plt.title('Detector Number :{0}'.format(det))

axcolor = 'lightgoldenrodyellow'

if Hfile:
	rax = plt.axes([0.67, 0.05, 0.15, 0.1], axisbg=axcolor)
	radio = RadioButtons(rax, ('Temp', 'H Field'))

bx = plt.axes([0.84, 0.05, 0.1, 0.1], axisbg=axcolor)
button  = Button (bx, 'save',color=axcolor)

axhte = plt.axes([0.1, 0.1, 0.55, 0.03], axisbg=axcolor)
detno = Slider(axhte, '', 1, 32, valinit=d0)
detno.valtext.set_visible(False)

def update(val):
	global det
	det = detno.val
	det=int(det)
	dd=detectors[:,det-1]
	lx.set_ydata(dd)
	axes.set_ylim(min(dd), max(dd))
	plt.title('')
	plt.title('Detector Number :{0}'.format(det))
	fig.canvas.draw_idle()

def swap_X(label):
		global temps
		global lx
		global Hfield
		global plt
		if label == 'Temp':
			lx.set_xdata(temps)
			axes.set_xlim(min(temps), max(temps))
			#plt.xlabel('Temperature')
			axes.set_xlabel('Temperature')
			fig.canvas.draw_idle()
		else:
			lx.set_xdata(Hfield)
			axes.set_xlim(float(min(Hfield)),float(max(Hfield)))
			#plt.xlabel('Field')
			axes.set_xlabel('H Field')
			fig.canvas.draw_idle()

def saveme(event):
	global det
	print "Saving the current data to "+str(det)+".dat"
	x = lx.get_data()[0]
	y = lx.get_data()[1]
	fileout = str(det)+".dat"
	file = open(fileout, "w")
	for index in range(len(x)):
		file.write(str(x[index]) + " " + str(y[index]) + "\n")
	file.close()
	
detno.on_changed(update)
if Hfile:
	radio.on_clicked(swap_X)
button.on_clicked(saveme)

plt.show()