#!/usr/bin/python

from sys import argv

from numpy import *

import matplotlib.pyplot as plt

from matplotlib.colors import LogNorm

from mpl_toolkits.mplot3d.axes3d import Axes3D

if len(argv)<2:
   print "Please give at least 1 file id \n For example: 14 for Results.EGMF14 \n or 14 15 for EGMF14 and EGMF15"
   exit()

# figure: position of the arriving particles in space -> size circle = energy
#=============================================================================

nbBins = 100

degre= 180/pi

borne=5 # degree

Elim=1e-1 # GeV

fig = plt.figure()

nbplots = len(argv)-1
ind = 1
for fileId in argv[1:]:

   energy,thetaDir,phiDir = loadtxt("Results_EGMF"+fileId+"/Results_momentum",unpack=True,usecols=[0,4,5])
   thetaPos,phiPos = loadtxt("Results_EGMF"+fileId+"/Results_position",unpack=True,usecols=[4,5])
   charge,weight= loadtxt("Results_EGMF"+fileId+"/Results_extra",unpack=True,usecols=[0,2])

   thetaPos = thetaPos*degre
   phiPos = phiPos*degre
   thetaDir = thetaDir*degre
   phiDir = phiDir*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 "Cut at", Elim, "GeV:", fileId, "->", shape(Energy)[0], "particles"
   plotid = 100+nbplots*10+ind

   #Emax = max(Energy)
   #Emin = min(Energy)
   #area = pi * (15 * Energy/Emax)**2 +15
   #colors = Energy/Emax

   #ax = fig.add_subplot(plotid)

   #sax = ax.scatter(theta, phi, s=area, c=colors, alpha=0.3, marker='o')
   #ax.set_xlim((-borne,borne))
   #ax.set_ylim((-borne,borne))
   #ax.set_xlabel("$\\theta$ in degree")
   #ax.set_ylabel("$\\phi$ in degree")
   #ax.set_title("EGMF: $10^{-%.0f"%float(fileId)+"}$Gauss")
 
   #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 = fig.add_subplot(plotid)
   H,xedges,yedges,im = ax.hist2d(theta,phi,nbBins,weights=weight/shape(weight),norm=LogNorm())

   extent = [yedges[0], yedges[-1], xedges[0], xedges[-1]]

   cbar=fig.colorbar(im, ax=ax)

   cbar.ax.set_ylabel("counts")

   ax.set_xlim((-borne,borne))
   ax.set_ylim((-borne,borne))
   ax.set_xlabel("$\\theta$ in degree")
   ax.set_ylabel("$\\phi$ in degree")
   ax.set_title("EGMF: $10^{-%.0f"%float(fileId)+"}$Gauss")

   ind=ind+1

plt.show()