#!/bin/python

from sys import argv
from numpy import append,savetxt, shape, array, newaxis, empty
from Modules.Read import ReadEnergy, ReadTime, ReadExtraFile, ReadProfile
from Modules.Spectrum import spectrum
from Modules.Angle import radial, angle_vs_energy
from Modules.Timing import timing

def error(error_type):
   print error_type
   exit()

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

folder = "Results/"

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

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

   for powerlaw_index in [1,1.5,2,2.5]: 
      # apply source spectrum
      weight = weightini*(Esource)**(1-powerlaw_index)

      #=============================================================================#
      #  SPECTRUM (SOURCE AND MEASURED)
      #=============================================================================#
      nbBins = 100
      # draw source spectrum
      Es=array(list(set(Esource)))
      Ws= (Es)**(1-powerlaw_index)
      Es,Fs = spectrum(Es,Ws,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
      cond = (generation==0)
      ener,flux_0 = spectrum(energy[cond],weight[cond],nbBins)
      # full spectrum
      ener,flux = spectrum(energy,weight,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)

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