from numpy import shape, pi, histogram2d, arange, newaxis, exp, log, array, sin, cos
from scipy.signal import convolve2d
from matplotlib.pyplot import figure, show, matshow
from matplotlib.colors import LogNorm
from Read import ReadEnergy, ReadWeight, ReadMomentumAngle, ReadPositionAngle, ReadNphot_source
from Read import ReadD_source, Readz_source, ReadEGMF
from Constants import degre

thetalim = 180
philim = 180
step = 0.25 #deg
nbBins = array([int(2*philim/step)+1,int(2*thetalim/step)+1])
philim=float(philim)
thetalim=float(thetalim)
limits = [[-thetalim,thetalim],[-philim,philim]]

def thetaphi(fileId,Elim=0.1):
   '''
      Compute arrival angles (theta,phi) selected with energy upper than Elim 
      normalize to 1 source photon
      Input: directory name
      Input (optional): lower energy limit in GeV (default = 100MeV)
      Output: energy (GeV), theta, phi (degre), weight
   '''
   energy = ReadEnergy(fileId)
   weight = ReadWeight(fileId)
   thetaDir,phiDir = ReadMomentumAngle(fileId)*degre
   thetaPos,phiPos = ReadPositionAngle(fileId)*degre
   Nmax=ReadNphot_source(fileId) # Nb photons emitted by the source

   theta = thetaDir - thetaPos
   phi = phiDir -phiPos

   cond=(energy>Elim) & (abs(theta)<thetalim) & (abs(phi)<philim)
   energy = energy[cond]
   weight=weight[cond]/Nmax
   theta=theta[cond]
   phi=phi[cond]

   print "file", fileId, "->", shape(energy)[0], "events"
   return energy, theta, phi, weight

def isotrop():
   return array([[1]])

def Gaussian2D(center=[0,0],fwhm=0.01):
   x = arange(-thetalim,thetalim+step,step)
   y = arange(-philim,philim+step,step)
   y = y[:,newaxis]
   return exp(-4*log(2) * ((x-center[0])**2 + (y+center[1])**2) / fwhm**2)

def drawMap(files,Elim=0.1,source=isotrop(),borne=[180,180]):
   '''
      Plot 2D histogram of arrival direction of the photons with energy upper than Elim
      Input: list of directories
      Input (optional): lower energy limit (in GeV (default = 100MeV)
      Input (optional): nature of the source (default isotrop) 
      Input (optional): limits on theta, phi 
      Output: 2D histogram of theta cos(phi), theta sin(phi)
   '''
   print "Cut at", Elim, "GeV" 
   fig = figure()
   nbplots = len(files)

   ind=1
   for fileId in files:
      energy, theta, phi, weight = thetaphi(fileId,Elim)

      H,xedges,yedges = histogram2d(theta*cos(phi),theta*sin(phi),nbBins,weights=weight,range=limits)
      conv=convolve2d(H,source,boundary="wrap",mode="same")

      ax = fig.add_subplot(100+nbplots*10+ind)
      im = ax.matshow(conv,norm=LogNorm(),aspect='auto',
                        extent=[xedges[0],xedges[-1],yedges[0],yedges[-1]])
      #H,xedges,yedges,im = ax.hist2d(theta,phi,nbBins,weights=weight,norm=LogNorm())
      cbar=fig.colorbar(im, ax=ax)
      cbar.ax.set_ylabel("counts normalize to 1 photon emitted")
      ax.set_xlim((-borne[0],borne[0]))
      ax.set_ylim((-borne[1],borne[1]))
      ax.set_xlabel("$\\theta$ $\\cos(\\phi)$ [deg]")
      ax.set_ylabel("$\\theta$ $\\sin(\\phi)$ [deg]")
      ax.grid(b=True,which='major')
      zsource=Readz_source(fileId)
      Dsource=ReadD_source(fileId)
      EGMF=ReadEGMF(fileId)
      ax.set_title("source at %.0f"%Dsource+"Mpc (z=%.4f"%zsource+"): EGMF="+str(EGMF)+" Gauss"   )
    
      ind=ind+1

   show()

#def drawMap(files):
#   print "Cut at", Elim, "GeV" 
#   fig = figure()
#   nbplots = len(files)
#
#   ind=1
#   for fileId in files:
#      energy, theta, phi, weight = thetaphi(fileId)
#
#      ax = fig.add_subplot(100+nbplots*10+ind)
#
#      Emax = max(energy)
#      Emin = min(energy)
#      colors = energy/Emax
#      area = pi * (15 * colors)**2 +15
#      sax = ax.scatter(theta, phi, c=colors, s=area, alpha=0.3, marker='o')
#
#      if Emin < 1e-3:
#         liminf = "%.0f"%(Emin*1e6)+"keV"
#      elif Emin < 1:
#         liminf = "%.0f"%(Emin*1e3)+"MeV"
#      else: 
#         liminf = "%.0f"%(Emin)+"GeV"
#      limsup = "%.0f"%(Emax*1e-3)+"TeV"
#      cbar = fig.colorbar(sax, ticks=[0.01,1])
#      cbar.ax.set_yticklabels([liminf,limsup])
#
#      ax.set_xlim((-borne,borne))
#      ax.set_ylim((-borne,borne))
#      ax.set_xlabel("$\\theta$ [deg]")
#      ax.set_ylabel("$\\phi$ [deg]")
#      zsource=Readz_source(fileId)
#      Dsource=ReadD_source(fileId)
#      EGMF=ReadEGMF(fileId)
#      ax.set_title("source at %.0f"%Dsource+"Mpc (z=%.4f"%zsource+"): EGMF="+str(EGMF)+" Gauss"   )
# 
#      ind=ind+1
#
#   show()