#!/bin/python

from sys import argv

from numpy import append, savetxt, shape, array, newaxis, zeros, arange

from Modules.Read import ReadEnergy, ReadTime, ReadExtraFile, ReadProfile

from Modules.Read import ReadMomentumAngle, ReadPositionAngle, resultDirectory

from Modules.Spectrum import spectrum

from Modules.Map import printMap, isotrop

from Modules.Angle import radial, angle_vs_energy
from Modules.Timing import timing

from Modules.Constants import degre

def InProgress(fileId,step,i,Nmax):
   print "Work on ", fileId, "(step "+ step +": ", (i*100)/Nmax, "% done"

if shape(argv)[0] < 2: 
   print "not enough arguments (at least 1)"
   exit()

#================================================================#
#  Parameters
#================================================================#

PowerSpectrum=[1,1.5,2,2.5] 

powerlaw_index = 2
Elim = 1e-1 # GeV
thetalim = 20 # degre

for fileId in argv[1:]:
   # read files
   time = ReadTime(fileId)
   energy = ReadEnergy(fileId)
   weightini, generation, theta_arrival, Esource = ReadExtraFile(fileId,[2,3,4,5])
   nbPhotonsEmitted=ReadProfile(fileId,[3]) 

   #=============================================================================#
   #  IMAGING 
   #=============================================================================#
   thetaDir,phiDir = ReadMomentumAngle(fileId)*degre
   thetaPos,phiPos = ReadPositionAngle(fileId)*degre
   theta = thetaDir - thetaPos
   phi = phiDir -phiPos

   # apply source spectrum
   weight_source = (Esource/min(Esource))**(1-powerlaw_index)
   weight = weightini* weight_source

   # apply selection
   cond=(energy>Elim) & (abs(theta)<thetalim) & (abs(phi)<thetalim)
   weight=weight[cond]
   theta=theta[cond]
   phi=phi[cond]
   printMap(theta,phi,weight,fileId,source=isotrop(),borne=[thetalim,thetalim])

   Gen_contrib = []
   Gen_cont = zeros((int(max(generation))+1))

   Source = []
   Spectrum = []
   Radial = []
   Angle_Energy = []
   Timing = []

   for powerlaw_index in PowerSpectrum:

      # apply source spectrum

      weight_source = (Esource/min(Esource))**(1-powerlaw_index)

      weight = weightini* weight_source

      #=============================================================================#

      # GENERATION CONTRIBUTION
      #=============================================================================#
      NbTotEvents = sum(weight)
      if Gen_contrib==[]:
         Gen_contrib = arange(0,max(generation)+1,1)[:,newaxis]

      for gen in list(set(generation)):
         Gen_cont[int(gen)] = sum(weight[generation==gen])/NbTotEvents *100
      Gen_contrib = append(Gen_contrib,Gen_cont[:,newaxis],axis=1)

      #=============================================================================#

      #  SPECTRUM (SOURCE AND MEASURED)
      #=============================================================================#
      nbBins = 100
      # draw source spectrum
      Es=array(list(set(Esource)))

      Ws= (Es/min(Es))**(1-powerlaw_index)

      Es,Fs = spectrum(Es,Ws,nbBins=nbBins)

      Es=Es[:,newaxis]
      Fs=Fs[:,newaxis]
      if Source==[]:
         Source = Es
         Source = append(Source,Fs,axis=1)
      else: 
         Source = append(Source,Fs,axis=1)

      # primary gamma-rays contribution and full spectrum
      ener,flux,flux_0 = spectrum(energy,weight,generation,nbBins)

      ener=ener[:,newaxis]
      flux=flux[:,newaxis]
      flux_0=flux_0[:,newaxis]
      if Spectrum==[]:
         Spectrum = ener
         Spectrum = append(Spectrum,flux,axis=1)
         Spectrum = append(Spectrum,flux_0,axis=1)
      else:
         Spectrum = append(Spectrum,flux,axis=1)
         Spectrum = append(Spectrum,flux_0,axis=1)

      #=============================================================================#
      #  RADIAL DISTRIBUTION AND ANGLE VERSUS ENERGY
      #=============================================================================#
      nbBins = 100
      ener,angle = angle_vs_energy(theta_arrival,energy,weight,nbPhotonsEmitted)
      ener=ener[:,newaxis]
      angle=angle[:,newaxis]
      if Angle_Energy==[]:
         Angle_Energy = ener
         Angle_Energy = append(Angle_Energy,angle,axis=1)
      else:
         Angle_Energy = append(Angle_Energy,angle,axis=1)

      theta,dndtheta2 = radial(theta_arrival,weight,nbPhotonsEmitted)
      theta=theta[:,newaxis]
      dndtheta2=dndtheta2[:,newaxis]
      if Radial==[]:
         Radial = theta
         Radial = append(Radial,dndtheta2,axis=1)
      else:
         Radial = append(Radial,dndtheta2,axis=1)

      #=============================================================================#
      #  TIME DISTRIBUTION AND TIME DELAY VERSUS ANGLE
      #=============================================================================#
      nbBins = 100
      delta_t,dNdt = timing(time,weight,nbBins)
      delta_t=delta_t[:,newaxis]
      dNdt=dNdt[:,newaxis]
      if Timing==[]:
         Timing = delta_t
         Timing = append(Timing,dNdt,axis=1)
      else:
         Timing = append(Timing,dNdt,axis=1)

      InProgress(fileId,"1",PowerSpectrum.index(powerlaw_index)+1,shape(PowerSpectrum)[0])

   savetxt(resultDirectory+fileId+"/Generation.txt",Gen_contrib)
   savetxt(resultDirectory+fileId+"/Spectrum.txt",Spectrum)
   savetxt(resultDirectory+fileId+"/Source_spectrum.txt",Source)
   savetxt(resultDirectory+fileId+"/Angle_vs_Energy.txt",Angle_Energy)
   savetxt(resultDirectory+fileId+"/Radial_distribution.txt",Radial)
   savetxt(resultDirectory+fileId+"/Timing.txt",Timing)