from numpy import shape, pi
from matplotlib.pyplot import figure, show
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

Elim=1e-1 # GeV
borne=5 # degree
nbBins = 100

def thetaphi(fileId):
   energy = ReadEnergy(fileId)
   weight = ReadWeight(fileId)
   thetaDir,phiDir = ReadMomentumAngle(fileId)*degre
   thetaPos,phiPos = ReadPositionAngle(fileId)*degre

   theta = thetaDir - thetaPos
   phi = phiDir -phiPos

   cond=(energy>Elim) & (abs(theta)<borne) & (abs(phi)<borne)
   energy = energy[cond]
   weight=weight[cond]
   theta=theta[cond]
   phi=phi[cond]

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

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()

def drawMapHisto(files):
   print "Cut at", Elim, "GeV" 
   fig = figure()
   nbplots = len(files)

   ind=1
   for fileId in files:
      energy, theta, phi, weight = thetaphi(fileId)
      Nmax=ReadNphot_source(fileId) # Nb photons emitted by the source

      ax = fig.add_subplot(100+nbplots*10+ind)
      H,xedges,yedges,im = ax.hist2d(theta,phi,nbBins,range=[[-borne,borne],[-borne,borne]],weights=weight/Nmax,norm=LogNorm())
      extent = [yedges[0], yedges[-1], xedges[0], xedges[-1]]
      cbar=fig.colorbar(im, ax=ax)
      cbar.ax.set_ylabel("counts normalize to 1 photon emitted")
      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()