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 Read import  resultDirectory, ReadRadial, ReadAngleVsEnergy
from Analytic import Analytic_theta

from Constants import degre

thetalim = 180

philim = 180
step = 0.25 #deg

bining = array([int(2*philim/step)+1,int(2*thetalim/step)+1])

philim=float(philim)
thetalim=float(thetalim)
limits = [[-thetalim,thetalim],[-philim,philim]]

def computeMap(theta,phi,weight,energy,dir_source,saveId,jet_opening_angle=180,Elim=1e-3,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)
   '''

   # apply selection
   cond = (dir_source*degre <= jet_opening_angle) & (energy>Elim) #& (abs(theta)<thetalim) & (abs(phi)<thetalim)
   theta    = theta[cond]
   phi      = phi[cond]
   weight   = weight[cond]
   energy   = energy[cond]

   fig = figure()

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

   ax = fig.add_subplot(111)

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

   if jet_opening_angle == 180:
      ax.set_title(saveId+" - isotrop")
      savefig(resultDirectory+saveId+"/map_isotrop.png")
   else:
      ax.set_title(saveId+" - jet opening angle: %d degre"%(int(jet_opening_angle)))
      savefig(resultDirectory+saveId+"/map_"+str(int(jet_opening_angle))+".png")

   if borne != [180,180]: 
      H,xedges,yedges,im = ax.hist2d(theta*cos(phi),theta*sin(phi),bining,weights=weight,
            norm=LogNorm(),range=limits)

   #===============================================================================================#
   #  RADIAL DISTRIBUTION AND ENERGY VERSUS ANGLE
   #===============================================================================================#

   theta2 = zeros(shape(H)[0]*shape(H)[1])
   N_evts = zeros(shape(H)[0]*shape(H)[1])

   theta_cos_phi = xedges[::-1]
   theta_sin_phi = yedges[::-1]

   it = nditer(H,flags=["multi_index"])

   k = 0
   while not it.finished:
      i = int(it.multi_index[1])
      j = int(it.multi_index[0])
      theta2[k]= sqrt(theta_cos_phi[i]**2+theta_sin_phi[j]**2)
      N_evts[k]= it[0]

      k += 1
      it.iternext()

   theta2=log10(theta2)
   thetamax=max(borne)
   dn,angle = histogram(theta2,nbBins,range=[log10(1e-2),log10(thetamax)],weights=N_evts)
   angle=10**angle
   thetacenter = (angle[1:nbBins+1]+angle[0:nbBins])/2
   dtheta = angle[1:nbBins+1]-angle[0:nbBins]
   dndtheta = dn/dtheta*thetacenter

   energy =log10(energy)
   angle,ener =histogram(energy,nbBins,weights=weight*theta)
   nb,ener =histogram(energy,nbBins,weights=weight)
   ener = 10**ener
   enercenter=(ener[1:nbBins+1]+ener[0:nbBins])/2
   angle /= nb 

   return thetacenter, dndtheta , enercenter, angle

def drawRadial(files):
   '''
      Plot flux versus arrival angle
      Input: list of directories
      Output: graph of flux versus angle
   '''
   fig = figure()
   ax = fig.add_subplot(111)

   for fileId in files:
      theta,dndtheta2 = ReadRadial(fileId,[0,1])
      ax.plot(theta,dndtheta2,drawstyle='steps-mid',label=fileId)

   ax.legend(loc="best")
   ax.set_xscale('log')
   ax.set_yscale('log')
   ax.grid(b=True,which='major')
   ax.set_xlabel("$\\theta$ [deg]")
   ax.set_ylabel("$\\theta*dN/d\\theta$")
   
   show()

def drawAngle_vs_energy(files):
   '''
      Plot arrival angle versus energy + analytic expression
      Input: list of directories
      Output: graph of arrival angle versus energy
   '''
   fig = figure()
   ax1 = fig.add_subplot(111)

   for fileId in files:
      ener,angle = ReadAngleVsEnergy(fileId,[0,1])
      p=ax1.plot(ener,angle,drawstyle='steps-mid',label=fileId)

      yfit = Analytic_theta(ener,fileId)
      ax1.plot(ener,yfit,'--',color="m",linewidth=2,label="analytic")

   ax1.legend(loc="best")
   ax1.set_xscale('log')
   ax1.set_yscale('log')
   ax1.set_ylim([0,180])
   ax1.grid(b=True,which='major')
   ax1.set_ylabel("$\\theta$ [rad]")
   ax1.set_xlabel("energy [GeV]")

   show()