from numpy import shape, histogram2d, arange, newaxis, array, nditer, zeros, histogram
from numpy import sqrt, pi, exp, log, log10, sin, cos, where

from matplotlib.pyplot import figure, matshow, savefig, show

from matplotlib.colors import LogNorm

#from scipy.signal import convolve2d

from read import  ReadRadial, ReadAngleVsEnergy

def computeMap(theta,phi,weight,energy,saveId,nbBins=50,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(figsize=(12,9))

   ax = fig.add_subplot(111)

   step = 2*array(borne,dtype=float)/nbBins
   philim=float(borne[0])
   thetalim=float(borne[1])
   bining = array([int(2*philim/step[0])+1,int(2*thetalim/step[1])+1])
   limits = [[-thetalim,thetalim],[-philim,philim]]

   #===============================================================================================#
   #  IMAGE
   #===============================================================================================#

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

   H,xedges,yedges,im = ax.hist2d(theta*cos(phi),theta*sin(phi),bining,weights=weight,
         norm=LogNorm(),range=limits)

   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(saveId+" - isotrop")

   savefig(saveId+"/map_isotrop.png")

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