from numpy import array, cos, sin 
from matplotlib.pyplot import figure, matshow, savefig, show
from matplotlib.colors import LogNorm
from matplotlib import rcParams
from src.constants import degre

label_size = 20
rcParams['xtick.labelsize'] = label_size 
rcParams['ytick.labelsize'] = label_size 

def image(theta,phi,weight,saveId,step=5,borne=[180,180]):
   '''
      Plot 2D histogram of arrival direction of the photons with energy upper than Elim
      Input: theta and phi in rad (important)
      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=(14,12))
   ax = fig.add_subplot(111) 

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

   H,xedges,yedges,im = ax.hist2d(theta*cos(phi)/degre,theta*sin(phi)/degre,
           bining,weights=weight,norm=LogNorm(),range=limits)
   cbar=fig.colorbar(im, ax=ax)
   cbar.ax.set_ylabel("Surface brightness $dN/(dSd\\Omega_d)$",fontsize=label_size)
                                             
   ax.grid(b=True,which='major')
   ax.set_xlim((-borne[0],borne[0]))
   ax.set_ylim((-borne[1],borne[1]))
   ax.set_xlabel("$\\theta_d$ $\\cos(\\phi_d)$ [deg]",fontsize=label_size)
   ax.set_ylabel("$\\theta_d$ $\\sin(\\phi_d)$ [deg]",fontsize=label_size)

   savefig(saveId+"/image.png")