#!/usr/bin/python

from sys import argv
import numpy as np
import matplotlib.pyplot as plt
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 = 50
degre= 180/np.pi
borne=5 # degree
Elim=1e-3 # GeV

fig = plt.figure()

nbplots = len(argv)-1
ind = 1
for fileId in argv[1:]:
   energy,thetaDir,phiDir = np.loadtxt("Results_EGMF"+fileId+"/Results_momentum",unpack=True,usecols=[0,4,5])
   thetaPos,phiPos = np.loadtxt("Results_EGMF"+fileId+"/Results_position",unpack=True,usecols=[4,5])

   plotid = 100+nbplots*10+ind
   ax = fig.add_subplot(plotid)

   Energy = energy[energy>Elim]
   print "Cut at", Elim, "GeV:", np.shape(Energy)
   thetaPos = thetaPos[energy>Elim]*degre
   phiPos = phiPos[energy>Elim]*degre
   thetaDir = thetaDir[energy>Elim]*degre
   phiDir = phiDir[energy>Elim]*degre
   theta = thetaDir - thetaPos
   phi = phiDir -phiPos

   ## random position of the detector
   #a=np.random.rand(2,1)
   #w=a[0,0]*2.-1.
   #u=np.sqrt(1-w*w)*np.cos(2*np.pi*a[1,0])
   #v=np.sqrt(1.-w*w)*np.sin(2.*np.pi*a[1,0])
   ## 2D -> use only u and v
   #theta_detector = u*180 #degre
   #phi_detector = v*90 #degre

   #FOV = 0.1 #degre

   #energy = energy[energy>Elim 
   #                  and np.absolute(theta-theta_detector)<FOV 
   #                  and np.absolute(phi-phi_detector)<FOV]
   #print "Cut at", Elim, "GeV:", np.shape(Energy)
   #theta = theta_detector - theta[energy>Elim 
   #               and np.absolute(theta-theta_detector)<FOV
   #               and np.absolute(phi-phi_detector)<FOV]
   #phi = phi_detector - phi[energy>Elim 
   #            and np.absolute(theta-theta_detector)<FOV 
   #            and np.absolute(phi-phi_detector)<FOV]
   #Energy=np.append((Energy,energy),axis=1)
   #Theta=np.append((Theta,theta),axis=1)
   #Phi=np.append((Phi,phi),axis=1)

   Emax = max(Energy)
   Emin = min(Energy)
   area = np.pi * (15 * Energy/Emax)**2 +15
   colors = Energy/Emax
   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])

   ind=ind+1

plt.show()