import numpy as np


def schwerpunktenergie(emin,emax,gamma):
  if gamma!=-1:
    es=((emin**(gamma+1)+emax**(gamma+1))/2.)**(1/(gamma+1))
  else:
    es=np.sqrt(emin*emax)
  return es

print "Schwerpunktenergie:"
mode=float(raw_input("Select mode (0=individual, 1=EPHIN_lvl3 chs, 2=EPHIN lvl3 ions overview, 3=EPHIN electron overview): "))

if mode==0:
  emin=float(raw_input("Emin: "))
  emax=float(raw_input("Emax: "))
  gamma=float(raw_input("Gamma: "))

  print schwerpunktenergie(emin,emax,gamma)


if mode==1:
  gamma=float(raw_input("Gamma: "))

  emin,emax,emean= 4.3,7.8,6.05
  print "P4/H4: emin: %2.2f emax: %2.2f emean_log: %2.2f e_schwerpunkt:%2.2f" %(emin,emax,emean,schwerpunktenergie(emin,emax,gamma))

  emin,emax,emean=7.3,25,16.15
  print "P8/H8: emin: %2.2f emax: %2.2f emean_log: %2.2f e_schwerpunkt:%2.2f" %(emin,emax,emean,schwerpunktenergie(emin,emax,gamma))

  emin,emax,emean=25,40.9,32.95
  print "P25/H25: emin: %2.2f emax: %2.2f emean_log: %2.2f e_schwerpunkt:%2.2f" %(emin,emax,emean,schwerpunktenergie(emin,emax,gamma))
  
  emin,emax,emean=40.9,53,46.95
  print "P41/H41: emin: %2.2f emax: %2.2f emean_log: %2.2f e_schwerpunkt:%2.2f" %(emin,emax,emean,schwerpunktenergie(emin,emax,gamma))

if mode==2:
  import matplotlib.pyplot as plt
  bins=100 
  gamma_min,gamma_max=-3,1
  emins=[4.3,7.3,25.,40.9]
  emaxs=[7.3,25.,40.9, 53]
  labels=["P4","P8","P25","P41"]
  plt.figure(figsize=(8.27,11.69))
  plt.subplots_adjust(top=0.98,bottom=0.05,hspace=0.2,right=0.95,left=0.08)
  for ene_idx,w in enumerate(range(len(emins))):
    emin=emins[w]
    emax=emaxs[w]
    gammas=np.linspace(gamma_min,gamma_max,bins)
    es=np.zeros(bins)
    eth_min=emin*np.ones(bins)
    eth_max=emax*np.ones(bins)
    for idx,q in enumerate(gammas): 
      if -1.3<q<-0.7:
        eth_min[idx]=np.nan
        eth_max[idx]=np.nan
      if -1.15<q<-0.85:
        es[idx]=np.nan
      else:
        es[idx]=schwerpunktenergie(emin,emax,q)
    plt.plot(gammas,es, "r-", lw=2) 
    plt.plot(gammas,eth_min,"k-") 
    plt.plot(gammas,eth_max,"k-") 
    plt.text(-1,np.sqrt(emin*emax),labels[ene_idx],va="center",ha="center",color="r",fontsize=14)
    plt.text(-1,emin,"%1.1f MeV"%emin,va="center",ha="center",color="k",fontsize=12)
    if ene_idx==len(emins)-1: 
      plt.text(-1,emax,"%1.1f MeV"%emax,va="center",ha="center",color="k",fontsize=12)
    plt.text(gammas[len(gammas)/14],es[0]-0.5,"%1.1f MeV"%es[0],va="center",ha="center",fontsize=14)
    plt.text(gammas[13*len(gammas)/14],es[-1]+0.5,"%1.1f MeV"%es[-1],va="center",ha="center",fontsize=14)

    #plt.plot(gammas,np.sqrt(emin*emax)*np.ones(len(gammas)),"k:") 
  max_de=emaxs[-1]-emins[0]
  plt.ylim(emins[0]-0.03*max_de,emaxs[-1]+0.03*max_de)

  for q in np.arange(5,55,5):
    plt.plot(gammas,q*np.ones(bins),"k:",alpha=0.5,lw=2)
  plt.xlabel("power-law index $\gamma$")
  plt.ylabel("center of gravity energy / MeV")
  plt.show()
  #plt.savefig("ephin_schwerpunktenergie.pdf")


if mode==3:
  import matplotlib.pyplot as plt
  bins=100 
  gamma_min,gamma_max=-3,1
  emins=[0.25,0.67,2.64,4.8]
  emaxs=[0.7,3,6.18,10.4]
  lim_fmts=["-","--","-.",":"]
  labels=["E150","E300","E1300","E3000"]
  plt.figure(figsize=(8.27,11.69))
  plt.subplots_adjust(top=0.98,bottom=0.05,hspace=0.2,right=0.95,left=0.08)
  for ene_idx,w in enumerate(range(len(emins))):
    emin=emins[w]
    emax=emaxs[w]
    gammas=np.linspace(gamma_min,gamma_max,bins)
    es=np.zeros(bins)
    eth_min=emin*np.ones(bins)
    eth_max=emax*np.ones(bins)
    for idx,q in enumerate(gammas): 
      if -1.3<q<-0.7:
        eth_min[idx]=np.nan
        eth_max[idx]=np.nan
      if -1.15<q<-0.85:
        es[idx]=np.nan
      else:
        es[idx]=schwerpunktenergie(emin,emax,q)
    plt.plot(gammas,es, "r-", lw=2) 
    plt.plot(gammas,eth_min,color="k",ls=lim_fmts[ene_idx]) 
    plt.plot(gammas,eth_max,color="k",ls=lim_fmts[ene_idx]) 
    plt.text(-1,np.sqrt(emin*emax),labels[ene_idx],va="center",ha="center",color="r",fontsize=14)
    plt.text(-1,emin,"%1.1f MeV"%emin,va="center",ha="center",color="k",fontsize=12)
    plt.text(-1,emax,"%1.1f MeV"%emax,va="center",ha="center",color="k",fontsize=12)
    plt.text(gammas[len(gammas)/14],es[0]-0.1,"%1.1f MeV"%es[0],va="center",ha="center",fontsize=14)
    plt.text(gammas[13*len(gammas)/14],es[-1]+0.1,"%1.1f MeV"%es[-1],va="center",ha="center",fontsize=14)

    #plt.plot(gammas,np.sqrt(emin*emax)*np.ones(len(gammas)),"k:") 
  max_de=emaxs[-1]-emins[0]
  plt.ylim(emins[0]-0.03*max_de,emaxs[-1]+0.03*max_de)

  for q in np.arange(5,55,5):
    plt.plot(gammas,q*np.ones(bins),"k:",alpha=0.5,lw=2)
  plt.xlabel("power-law index $\gamma$")
  plt.ylabel("center of gravity energy / MeV")
  plt.show()
  #plt.savefig("ephin_schwerpunktenergie.pdf")