from numpy import sqrt, arccos, shape, log10, histogram
from matplotlib.pyplot import figure, show, ylim, setp
from matplotlib import gridspec

from Read import ReadMomentum, ReadPosition, ReadEnergy, ReadWeight, ReadNphot_source

from Analytic import Analytic_theta

from Map import thetaphi

from Constants import degre

Eliminf = 1e0 #GeV
Elimsup = 1e5 #GeV

def compute_theta(fileName):

   '''
      Compute arrival angle
      Input: directory name
      Output: energy, weight, arrival angle
   '''

   position = ReadPosition(fileName)
   momentum = ReadMomentum(fileName)
   energy = ReadEnergy(fileName)
   weight = ReadWeight(fileName)

   hyp=sqrt(position[0]**2+position[1]**2+position[2]**2)
   position /= hyp
   hyp=sqrt(momentum[0]**2+momentum[1]**2+momentum[2]**2)
   momentum /= hyp

   costheta=position[0]*momentum[0]+position[1]*momentum[1]+position[2]*momentum[2]
   cond=(costheta<=1)&(costheta>=-1)&(energy>Eliminf)&(energy<Elimsup)
   theta = arccos(costheta[cond])*degre
   weight= weight[cond]
   energy= energy[cond]
   print "file", fileName, "->", shape(energy)[0], "events" 

   return energy, weight, theta

def angle_vs_energy(fileName):

   '''
      Compute arrival angle versus energy
      Input: directory name
      Output: energy, arrival angle
   '''

   energy,weight,theta = compute_theta(fileName)
   nbBins = 40

   energy =log10(energy)

   angle,ener =histogram(energy,nbBins,weights=theta*weight)

   nb,ener =histogram(energy,nbBins,weights=weight)

   ener = 10**ener
   enercenter=(ener[1:nbBins+1]+ener[0:nbBins])/2

   maxflux = ReadNphot_source(fileName)
   angle = angle/(nb*maxflux)

   return enercenter, angle

def drawAngle(files):

   '''
      Plot arrival angle versus energy + analytic expression
      Input: list of directories
      Output: graph of arrival angle versus energy
   '''

   fig = figure()
   gs = gridspec.GridSpec(2, 1, height_ratios=[4,1]) 
   fig.subplots_adjust(hspace=0.001)
   ax1 = fig.add_subplot(gs[0])
   ax2 = fig.add_subplot(gs[1],sharex=ax1)

   for fileId in files:
      energy,angle = angle_vs_energy(fileId)
      p=ax1.plot(energy,angle,'.',label=fileId)

      yfit = Analytic_theta(energy,fileId)
      ax1.plot(energy,yfit,'--',color=p[0].get_color(),linewidth=2,label="analytic")

      error=angle/yfit-1
      ax2.plot(energy,error,'+',color=p[0].get_color())

   ax1.legend(loc="best")
   ax1.set_xscale('log')
   ax1.set_yscale('log')
   ax1.set_ylim([0,180])
   ax1.grid(b=True,which='major')
   ax1.set_ylabel("$\\theta$ [degre]")
   ax2.set_xscale('log')
   ax2.grid(b=True,which='major')
   ax2.set_xlabel("energy [GeV]")

   ax2.set_ylabel("error")

   step=0.2*(max(error)-min(error))
   ylim(min(error)-step, max(error)+step)
   xticklabels = ax1.get_xticklabels()
   setp(xticklabels, visible=False)

   show()

##############################################################################################
def radial(fileId,thetamax=90):
   '''
      Compute flux versus arrival angle
      Input: directory name
      Input (optionnal): maximal angle desired (by default full sky)
      Output: arrival angle, flux (dn/dtheta)
   '''
   nbBins = 50
   energy,weight,theta = compute_theta(fileId)

   theta=log10(theta)
   dn,angle = histogram(theta,nbBins,range=[log10(1e-6),log10(thetamax)],weights=weight)
   angle=10**angle
   thetacenter = (angle[1:nbBins+1]+angle[0:nbBins])/2
   dtheta = angle[1:nbBins+1]-angle[0:nbBins]
   dndtheta = dn/dtheta
   maxflux = ReadNphot_source(fileId)
   dndtheta /= maxflux

   print "integral: ", sum(dtheta*dndtheta)

   return thetacenter, dndtheta

def drawRadial(files,thetamax=90):
   '''
      Plot flux versus arrival angle
      Input: list of directories
      Input (optionnal): maximal angle desired (by default full sky)
      Output: graph of flux versus angle
   '''

   fig = figure()
   ax = fig.add_subplot(111)

   for fileId in files:

      theta,dndtheta2=radial(fileId,thetamax)

      ax.plot(theta,dndtheta2,drawstyle='steps-mid',label=fileId)

   ax.legend(loc="best")
   ax.set_xscale('log')

   ax.set_yscale('log')
   ax.grid(b=True,which='major')

   ax.set_xlabel("$\\theta$ [deg]")
   ax.set_ylabel("$dN/d\\theta$ [deg$^{-1}$]")

   
   show()