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, savefig
from matplotlib.colors import LogNorm
from Read import  resultDirectory
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 isotrop():
   return array([[1]]), "isotropic"

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), "jet_"+str(fwhm)

def printMap(theta,phi,weight,fileId,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)
   '''
   fig = figure()

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

   ax = fig.add_subplot(111)
   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")
   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')
   ax.set_title(fileId+" - "+source[1])

   savefig(resultDirectory+fileId+"/map_"+source[1]+".png")
   #show()