import pexpect as px
import sys,time,os
import matplotlib.pyplot as plt

if len(sys.argv) != 8:
    print("Usage: python3 stopit_pex.py <z> <a> <thickness_um_layer1> <thickness_um_layer2> <energy_start_MeV>  <energy_stop_MeV> <energy_step>")
    print("Modify the \#absorber section in script to change stopping material")
    quit()

stoppeeZ=int(sys.argv[1])
stoppeeA=int(sys.argv[2])
layer1_thickness = float(sys.argv[3]) #um
layer2_thickness = float(sys.argv[4]) #um
min_stoppee_energy=float(sys.argv[5])
max_stoppee_energy=float(sys.argv[6])
step_stoppee_energy=float(sys.argv[7])

#plotting prep
'''
plt.ion()
fig=plt.figure()
ax=fig.add_subplot(111)
ax.autoscale_view(True,True,True)
line,=ax.plot([],[],'b-')
ax.set_xlim(0,0.5*max_stoppee_energy)
ax.set_ylim(0,0.05*max_stoppee_energy)
'''
stopit_child = px.spawn("./a.out")
stopit_child.logfile = None

def run_stopit(child,Z,A,Energy_MeV,thickness_um):
	child.expect("7 stop\r")

	#stoppee
	child.send("1\r")
	child.expect("Enter Z and A of stopee.\r")
	child.send(str(Z)+" "+str(A)+"\r")
	child.expect("Enter energy in MeV.\r")
	child.send(str(Energy_MeV)+"\r")

	#absorber - standard medium, Si
	child.expect("7 stop\r")
	child.send("2\r")
	child.expect("\r")
	child.send("1\r")
	child.expect("He, and H2\)\r")
	child.send("2\r")
	child.expect("H2-gas\r")
	child.send("2\r") #first 2 for 'define absorber', then 1 for '1 layer', then 2 for 'standardized material', then '2' for Si.
	child.expect("\r")
	child.send(str(thickness_um)+"\r")

	#run!
	child.expect("7 stop\r")
	child.send("4\r")
	child.expect("stopit.log\r")
	#get all the output that shows up before the string 'stopit.log' above.
	#the split() splits it up at all whitespaces. we can pick out the substrings of interest from this bunch now
	results = child.before.decode('utf-8').split()

	return [results[12],results[48],results[53]]

layer1_edep = []
layer2_edep = []
stoppee_energy = min_stoppee_energy #in MeV
while(stoppee_energy < max_stoppee_energy):
	[in_en,e_lost ,e_left ] = run_stopit(stopit_child,stoppeeZ,stoppeeA,stoppee_energy,layer1_thickness)
	[in_en,e_lost2,e_left2] = run_stopit(stopit_child,stoppeeZ,stoppeeA,e_left        ,layer2_thickness)

	print(stoppee_energy,e_lost,e_lost2)
	sys.stdout.flush()
	layer1_edep.append(float(e_lost))
	layer2_edep.append(float(e_lost2))
	'''
	line.set_data(layer2_edep,layer1_edep) #dE in y-axis, E in x-axis when possible
	ax.relim()
	plt.draw()
	plt.pause(0.004)
	'''
	stoppee_energy += step_stoppee_energy

stopit_child.expect("7 stop\r")
stopit_child.send("7\r")
stopit_child.expect(px.EOF)
stopit_child.close()

plt.show(block=True)