import numpy as np
import matplotlib.pyplot as plt
import time
from pylab import *
from datetime import datetime
from datetime import timedelta

from Tkinter import *
import tkMessageBox
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2TkAgg


def about():
  tkMessageBox.showinfo('About...','quick bethe-bloch - bragg peak calculator and plotter (17.12.13)')
  #lab=Label(root,text="Hier nicht!")
  #lab.pack()
#\frac{2m_ec^2\beta^2}{I\cdot(1-\beta^2)}
def showequation():
  win = Toplevel()
  f = plt.figure(figsize=(6,3), dpi=100)
  p1=plt.subplot(111)
  equation=r"$-\frac{dE}{dx}=\frac{4\pi}{m_ec^2}\cdot\frac{nz_P^2}{\beta^2}\cdot\left(\frac{e^2}{4\pi\epsilon_0}\right)^2\cdot(\ln(\frac{2m_ec^2\beta^2}{I\cdot(1-\beta^2)})-\beta^2)$"
  plt.text(0.5,0.7,equation,ha='center',size=20)
  plt.text(0.5,0.5,r"$n=q_T\rho/(m_T u)$",ha='center',size=20)
  plt.text(0.5,0.3,r"$\beta^2=1-(m_Puc^2/(E+E_{rest}))^2$",ha='center',size=20)
  plt.text(0.5,0.1,r"$-\frac{dE}{dx}\propto\frac{q_T}{m_T}\cdot\frac{q_P^2\rho}{\beta^2}\cdot(const.+\ln(\frac{\beta^2}{I\cdot(1-\beta^2)})-\beta^2)$",ha='center',size=20)
  plt.subplots_adjust(left=-0.1,right=1.1,top=1.1,bottom=-0.1)
  canvas=FigureCanvasTkAgg(f,master=win)
  canvas.show()
  canvas.get_tk_widget().pack()
  canvas._tkcanvas.pack(side=TOP, fill=BOTH, expand=YES)
  Button(win, text='OK', command=win.destroy).pack()

def calc_bethe(pm,pc,pe,tm,tc,td,tip,il,dl):
  #(entpm.get(),entpc.get(),entpe.get(),enttm.get(),enttc.get(),enttd.get(),enttip.get(),entil.get(),entdl.get())
  Etot=0
  thisdx=0
  dxleft=dl*10**(-6)
  binning=il*10**(-6)
  Eleft=pe
  dxarray,dearray,deiarray=[],[],[]
  while binning<dxleft and Eleft>0:
    thisdedx=dedx(Eleft,binning,pm,pc,tm,tc,td,tip)
    Eleft-=thisdedx*binning
    Etot+=thisdedx*binning
    dxleft-=binning
    thisdx+=binning
    dearray.append(thisdedx)
    dxarray.append(thisdx-binning/2.)
    try: deiarray.append(thisdedx*binning+deiarray[-1])
    except: deiarray.append(thisdedx*binning)
  if Eleft>0:
    thisdedx=dedx(Eleft,dxleft,pm,pc,tm,tc,td,tip)
    Etot+=thisdedx*dxleft
    thisdx+=binning
    dearray.append(thisdedx)
    dxarray.append(thisdx-binning/2.)
    try: deiarray.append(thisdedx*dxleft+deiarray[-1])
    except: deiarray.append(thisdedx*dxleft)
  dxarr,dedxarr,deiarr=np.array(dxarray)*10**(6),np.array(dearray)*10**(-3),np.array(deiarray)
  labeloutput1.config(text="Total energy deposition: %.2f MeV" %(Etot))
  labeloutput2.config(text=u"Position of max. dE/dx (Bragg-Peak): %i \u03bcm" %dxarr[np.argmax(dedxarr)])
  labeloutput3.config(text=u"Maximum range: %i \u03bcm" %dxarr[np.where(dedxarr>10**(-2))][-1])
  labeloutput4.config(text="Projectile Energy at Bragg-Peak: %.2f MeV" %(float(entpe.get())-deiarr[np.argmax(dedxarr)] ) )
  return dxarr,dedxarr,deiarr


def dedx(E,dx,projectile_mass,projectile_charge,tm,tc,td,tip):   #calculates dE/dx (in MeV) as function of kinetic E (in MeV) and dx (in ym) in silicon
  Z=tc  #14
  rho=td  #2336  #kg/m3
  A=tm  #28.085
  u=1.660538*10**(-27)
  e=1.602*10**(-19)
  n=Z*rho/(A*u)
  me=9.109*10**(-31)
  e0=8.854187817*10**(-12)
  I=tip*e #172*e
  mp=1.672*10**(-27)
  e=1.602*10**(-19)
  c=3*10**8
  ej=E*10**(6)*e +projectile_mass*mp*c**2  #E in joule
  #E=np.logspace(0,2,100)
  #v=np.sqrt(2.*ej*mp)
  v=c*np.sqrt(1-(projectile_mass*mp*c**2/ej)**2)
  b=v/c
  dedx=projectile_charge**2*4*np.pi*n*e**4/(me*b**2*c**2*4**2*np.pi**2*e0**2)  * (np.log( (2*me*c**2*b**2) / (I*(1-b**2)) ) -b**2 )
  dedx= dedx/(e*10**(6))  #back in Mev
  return dedx

def myexit():
  root.quit()
  root.destroy()


def canvasplot(f):
  for child in plotwindow.winfo_children():
    child.destroy()
  canvas=FigureCanvasTkAgg(f,master=plotwindow)
  canvas.show()
  canvas.get_tk_widget().pack()
  if vartoolbar.get():
    toolbar = NavigationToolbar2TkAgg( canvas, plotwindow )
    toolbar.update()
  #canvas._tkcanvas.configure(width = 800, height = 800)
  #canvas._tkcanvas.resizescrollregion()
  canvas._tkcanvas.pack(side=TOP, fill=BOTH, expand=YES)

def replot():
  f = plt.figure(figsize=(8,7), dpi=100)
  a = f.add_subplot(111)
  dxarr,dedxarr,deiarr=calc_bethe(float(entpm.get()),float(entpc.get()),float(entpe.get()),float(enttm.get()),float(enttc.get()),float(enttd.get()),float(enttip.get()),float(entil.get()),float(entdl.get()))
  a.step(dxarr[dedxarr>0],dedxarr[dedxarr>0],'k-') #a.plot(t,s)
  plt.subplots_adjust(left=0.14, right=0.86, top=0.93, bottom=0.1)
  plt.xticks(rotation=30)
  if enttitle.get()=="auto": plt.title(r"%s MeV projectile (m$_P$=%s u, q$_P$=%s e) on target (m$_T$=%s u, q$_T$=%s e)" %(entpe.get(),entpm.get(),entpc.get(),enttm.get(),enttc.get()))
  else: plt.title("%s" %(enttitle.get()))
  plt.grid()
  plt.xlabel(r"x in $\mu$m")
  plt.ylabel(r"dE/dx in keV/$\mu$m")
  if varlog.get():
    plt.xscale("log")
    plt.yscale("log")
  b=a.twinx()
  b.step(dxarr,deiarr,'r-')
  #np.savetxt("integration.dat",np.vstack((dxarr,deiarr)).T)
  b.set_ylabel(r"$\int_0^x \mathrm{d}E/\mathrm{d}x\ \mathrm{d}x /MeV$",color='r')  #('dE / MeV', color='r')
  for tl in b.get_yticklabels():
    tl.set_color('r')
  if varlog.get():
    plt.yscale("log")
  canvasplot(f)

root=Tk()
root.title("Bethe-Bloch Calculator and Plotter")

#plotwindow
plotwindow=Frame(root,width=900,height=900)
plotwindow.pack(expand=YES,side=RIGHT)


# buttonwindow
buttonwindow=Frame(root,width=400,height=900,cursor="trek")
buttonwindow.pack(expand=YES,side=LEFT)
buttonwindowrow=1

#description
labeldescription=Label(buttonwindow,text="Description",font="bold")
labeldescription.grid(row=buttonwindowrow,column=2)
buttonwindowrow+=1
labeldesc=Label(buttonwindow,text="This GUI integrates the Bethe-Bloch equation for quick estimations of the range/energy deposition of given projectiles and targets.\nStandard settings represent a proton hitting a selicon detector.\n \nThe plot can be saved by activating the toolbar and using the save button.\nThe plottitle can be changed - the keyword 'auto' creates a title including the projectile and target mass/charge.\n \nThe integration steplength should be choosen low enough, that unphysical, high energy depositions behind the Bragg peak do not occour (e.g. choose steplength<0.1)",justify=LEFT,wraplength=400 )
#wraplength=30  
labeldesc.grid(row=buttonwindowrow,column=1,columnspan=3)
buttonwindowrow+=1

labeldummy66=Label(buttonwindow,text=" ")
labeldummy66.grid(row=buttonwindowrow,column=2)
buttonwindowrow+=1

# parameter
labeldate=Label(buttonwindow,text="Parameter",font="bold")
labeldate.grid(row=buttonwindowrow,column=2)
buttonwindowrow+=1
labelp=Label(buttonwindow,text="Projectile Mass: ")
labelp.grid(row=buttonwindowrow,column=1)#.pack(side=LEFT)
entpm=Entry(buttonwindow,width=6)
entpm.grid(row=buttonwindowrow,column=2)#.pack(side=RIGHT)
entpm.insert(0, "1")
labelpu=Label(buttonwindow,text="u (e.g. 4He: 4, H: 1)")
labelpu.grid(row=buttonwindowrow,column=3)#.pack(side=LEFT)
buttonwindowrow+=1
labelpc=Label(buttonwindow,text="Projectile Charge: ")
labelpc.grid(row=buttonwindowrow,column=1)#.pack(side=LEFT)
entpc=Entry(buttonwindow,width=6)
entpc.grid(row=buttonwindowrow,column=2)#.pack(side=RIGHT)
entpc.insert(0,"1")
labelpcu=Label(buttonwindow,text="e (e.g. 4He: 2, H: 1)")
labelpcu.grid(row=buttonwindowrow,column=3)#.pack(side=LEFT)
buttonwindowrow+=1
labelp2=Label(buttonwindow,text="Projectile Energy:")
labelp2.grid(row=buttonwindowrow,column=1)#.pack(side=LEFT)
entpe=Entry(buttonwindow,width=6)
entpe.grid(row=buttonwindowrow,column=2)#.pack(side=RIGHT)
entpe.insert(0, "3")
labelpeu=Label(buttonwindow,text="MeV")
labelpeu.grid(row=buttonwindowrow,column=3)#.pack(side=LEFT)
buttonwindowrow+=1
labelt=Label(buttonwindow,text="Target Mass: ")
labelt.grid(row=buttonwindowrow,column=1)#.pack(side=LEFT)
enttm=Entry(buttonwindow,width=6)
enttm.grid(row=buttonwindowrow,column=2)#.pack(side=RIGHT)
enttm.insert(0, "28")
labeltmu=Label(buttonwindow,text="u (e.g. Si: 28)")
labeltmu.grid(row=buttonwindowrow,column=3)#.pack(side=LEFT)
buttonwindowrow+=1
labelt=Label(buttonwindow,text="Target Charge: ")
labelt.grid(row=buttonwindowrow,column=1)#.pack(side=LEFT)
enttc=Entry(buttonwindow,width=6)
enttc.grid(row=buttonwindowrow,column=2)#.pack(side=RIGHT)
enttc.insert(0,"14")
labeltcu=Label(buttonwindow,text="e (e.g. Si: 14)")
labeltcu.grid(row=buttonwindowrow,column=3)#.pack(side=LEFT)
buttonwindowrow+=1
labelt2=Label(buttonwindow,text="Target Density: ")
labelt2.grid(row=buttonwindowrow,column=1)#.pack(side=LEFT)
enttd=Entry(buttonwindow,width=6)
enttd.grid(row=buttonwindowrow,column=2)#.pack(side=RIGHT)
enttd.insert(0, "2336")
labeltdu=Label(buttonwindow,text=u"g/dm\u00B3 (e.g. Si: 2336)")
labeltdu.grid(row=buttonwindowrow,column=3)#.pack(side=LEFT)
buttonwindowrow+=1
labeltip=Label(buttonwindow,text="Mean excitation pot.: ")
labeltip.grid(row=buttonwindowrow,column=1)#.pack(side=LEFT)
enttip=Entry(buttonwindow,width=6)
enttip.grid(row=buttonwindowrow,column=2)#.pack(side=RIGHT)
enttip.insert(0,"172")
labeltdu=Label(buttonwindow,text="eV (e.g. Si: 172)")
labeltdu.grid(row=buttonwindowrow,column=3)#.pack(side=LEFT)
buttonwindowrow+=1
labelil=Label(buttonwindow,text="Steplength: ")
labelil.grid(row=buttonwindowrow,column=1)#.pack(side=LEFT)
entil=Entry(buttonwindow,width=6)
entil.grid(row=buttonwindowrow,column=2)#.pack(side=RIGHT)
entil.insert(0, "0.1"  ) #"0.1")
labeltdu=Label(buttonwindow,text=u"\u03bcm (e.g. 0.1)")
labeltdu.grid(row=buttonwindowrow,column=3)#.pack(side=LEFT)
buttonwindowrow+=1
labeldl=Label(buttonwindow,text="Detector Thickness: ")
labeldl.grid(row=buttonwindowrow,column=1)#.pack(side=LEFT)
entdl=Entry(buttonwindow,width=6)
entdl.grid(row=buttonwindowrow,column=2)#.pack(side=RIGHT)
entdl.insert(0, "150")
labeltdu=Label(buttonwindow,text=u"\u03bcm (e.g. Si: 150)")
labeltdu.grid(row=buttonwindowrow,column=3)#.pack(side=LEFT)
buttonwindowrow+=1
labeldummy2=Label(buttonwindow,text=" ")
labeldummy2.grid(row=buttonwindowrow,column=2)
buttonwindowrow+=1



#plot stuff
labelplot=Label(buttonwindow,text="Plot Options",font="bold")
labelplot.grid(row=buttonwindowrow,column=2)
buttonwindowrow+=1
labeltitle=Label(buttonwindow,text="Plot Title")
labeltitle.grid(row=buttonwindowrow,column=1)#.pack(side=LEFT)
enttitle=Entry(buttonwindow,width=28)
enttitle.grid(row=buttonwindowrow,column=2,columnspan=2)#.pack(side=RIGHT)
enttitle.insert(0, "auto")
buttonwindowrow+=1
varlog=IntVar()
chlog = Checkbutton(buttonwindow, text="Log plot", variable=varlog)
chlog.grid(row=buttonwindowrow,column=1)
vartoolbar=IntVar()
chtool = Checkbutton(buttonwindow, text="Toolbar (bugged)", variable=vartoolbar)
chtool.grid(row=buttonwindowrow,column=2)

buttonwindowrow+=1
labeldummy99=Label(buttonwindow,text=" ")
labeldummy99.grid(row=buttonwindowrow,column=2)
buttonwindowrow+=1
butreplot=Button(buttonwindow,text="(Re-)calculate and plot",command=replot)
butreplot.grid(row=buttonwindowrow,column=1,columnspan=3)  #pack()#side=TOP)
buttonwindowrow+=1
labeldummy3=Label(buttonwindow,text=" ")
labeldummy3.grid(row=buttonwindowrow,column=2)
buttonwindowrow+=1
#output
labelcommands=Label(buttonwindow,text="Results",font="bold")
labelcommands.grid(row=buttonwindowrow,column=2)
buttonwindowrow+=1
labeloutput1=Label(buttonwindow,text="Total energy deposition: ??? MeV") #, width=40,height=10)  #,state=DISABLED)
labeloutput1.grid(row=buttonwindowrow,column=1,columnspan=3)
buttonwindowrow+=1
labeloutput2=Label(buttonwindow,text=u"Position of max. dE/dx (Bragg-Peak): ??? \u03bcm") #, width=40,height=10)  #,state=DISABLED)
labeloutput2.grid(row=buttonwindowrow,column=1,columnspan=3)
buttonwindowrow+=1
labeloutput3=Label(buttonwindow,text=u"Maximum range: ??? \u03bcm") #, width=40,height=10)  #,state=DISABLED)
labeloutput3.grid(row=buttonwindowrow,column=1,columnspan=3)
buttonwindowrow+=1
labeloutput4=Label(buttonwindow,text="Projectile Energy at Bragg-Peak: ??? MeV") #, width=40,height=10)  #,state=DISABLED)
labeloutput4.grid(row=buttonwindowrow,column=1,columnspan=3)
buttonwindowrow+=1
labeldummy4=Label(buttonwindow,text=" ")
labeldummy4.grid(row=buttonwindowrow,column=2)
buttonwindowrow+=1
#commands
labelcommands=Label(buttonwindow,text="Commands",font="bold")
labelcommands.grid(row=buttonwindowrow,column=2)
buttonwindowrow+=1
but=Button(buttonwindow,text="About",command=about)
but.grid(row=buttonwindowrow,column=1)  #.pack()#side=BOTTOM)

butformular=Button(buttonwindow,text="Bethe-Bloch-Equation",command=showequation)
butformular.grid(row=buttonwindowrow,column=2)  #.pack()#side=BOTTOM)

butexit=Button(buttonwindow,text="Exit",command=myexit) 
butexit.grid(row=buttonwindowrow,column=3)  #.pack()#side=BOTTOM)
buttonwindowrow+=1


replot()
root.mainloop()