#!/bin/python

import os, shutil

from xml.dom               import minidom
from numpy                 import append, savetxt, shape, array, newaxis, zeros, arange, select, where
from modules.read          import ReadResults, ReadProfile, resultdir
from modules.spectrum      import spectrum
from modules.maps          import computeMap

from modules.arrival_angle import arrivalAngle

from modules.timing        import timing

from modules.observables   import observables_vs_delay, observables_vs_energy

from modules.constants     import degre, day, yr


def PSF_Taylor2011(E): # E in GeV -> PSF in degre
   return 1.7 * E**(-0.74) * (1+(E/15)**2)**0.37

xmlfile = minidom.parse("simulations.xml")
simulations = xmlfile.getElementsByTagName("simu")

for simu in simulations:
   fileId = simu.getAttribute("simulation_dir")
   output_dir = resultdir+simu.getAttribute("id")+"/"
   if not os.path.exists(output_dir):
      os.makedirs(output_dir)

   Emin = float(simu.getAttribute("Emin")) 
   Emax = float(simu.getAttribute("Emax"))

   print "#=============================================================================#"

   print " Reading data from ", fileId 

   print " Output directory: ", output_dir 

   # copy profile
   shutil.copy(fileId+"/profile.dat",output_dir+"profile.dat")

   # read files

   n_phot = ReadProfile("../"+fileId,[4]) 

   results=ReadResults(fileId)
   # select photons only

   cond = (results[0]%2==0) & (results[2]>=Emin) & (results[2]<=Emax)
   results = results[:,cond]

   generation = results[0]

   weight = results[1]/n_phot

   energy = results[2]
   time = results[3]
   arrival_pos1 = results[4]*degre
   arrival_pos2 = results[6]*degre
   theta = (results[4]-results[6])*degre
   phi = (results[5]-results[7])*degre
   theta_arrival = results[8]*degre
   Esource =results[9]

   print " [", Emin,"GeV,",Emax,"GeV] band ->",shape(results)[1],"events selected"
   theta_range = [1e-5,90]
   time_range=[1e0,1e9]#max(time)]
   
   PSF  = simu.getAttribute("PSF")
   if PSF == "Taylor2011":
      print " PSF:", PSF
   else:
      PSF = float(PSF)
      print " PSF:", PSF, "degre"

   #  SPECTRUM (SOURCE) =========================================================#
   if "simple case" not in fileId:
      powerlaw_index = simu.getAttribute("powerlaw_index") 
      print " Applying source spectrum:", powerlaw_index
      weightini = weight
      Es=array(list(set(Esource)))
      if powerlaw_index == "BL Lacertae*":
         Eth = 300
         gamaVHE = 2.94
         gamaHE  = 1.86
         condlist = [Esource>Eth,Esource<=Eth]
         spect = [(Esource/Eth)**(1-gamaVHE),(Esource/Eth)**(1-gamaHE)]
         weight = weightini*select(condlist,spect)
         Ws = select([Es>Eth,Es<=Eth],[(Es/Eth)**(1-gamaVHE),(Es/Eth)**(1-gamaVHE)])
      elif powerlaw_index == "1ES 0229+200":
         powerlaw_index = 5/3.
         weight = weightini*(Esource/min(Esource))**(1-powerlaw_index) 
         Ws= (Es/min(Es))**(1-powerlaw_index)  / n_phot
      else: 
         powerlaw_index = float(simu.getAttribute("powerlaw_index")) 
         weight = weightini*(Esource/min(Esource))**(1-powerlaw_index) 
         Ws= (Es/min(Es))**(1-powerlaw_index)  / n_phot

      nbBins = 100
      Es,Fs = spectrum(Es,Ws,nbBins=nbBins,E_range=[Emin,Emax])
      Source = Es[:,newaxis]
      Source = append(Source,Fs[:,newaxis],axis=1)
      if shape(Ws!=0)[0]>1:
          weight /= (max(Es)-min(Es))*(max(Fs)+min(Fs))/2
      savetxt(output_dir+"/Source_spectrum.txt",Source)

   NbTotEvents = sum(weight)
   print " ----------------------------------------------------------------------------- "

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

   #  NO SELECTION

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

   print "   > No selection of events ..." 

   #  IMAGING ===================================================================#

   nbBins = 300
   computeMap(theta,phi,weight,energy,output_dir,nbBins,borne=[theta_range[1],theta_range[1]])

   #  PARAMETER SPACE (theta, dt, E) ============================================#
   nbBins = 50

   ener,angle,dt = observables_vs_energy(energy,theta_arrival,time/yr,weight)

   Angle_Energy = ener[:,newaxis]
   Angle_Energy = append(Angle_Energy,angle[:,newaxis],axis=1)

   Delay_Energy = ener[:,newaxis]
   Delay_Energy = append(Delay_Energy,dt[:,newaxis],axis=1)

   dt,angle,ener = observables_vs_delay(energy,theta_arrival,time/yr,weight)

   Delay_vs_angle = angle[:,newaxis]
   Delay_vs_angle = append(Delay_vs_angle,dt[:,newaxis],axis=1)
   if PSF == "Taylor2011":
      cond = (theta_arrival<PSF_Taylor2011(energy))
   else:
      cond = (theta_arrival<PSF)

   ener,angle,dt = observables_vs_energy(energy[cond],theta_arrival[cond],time[cond]/yr,weight[cond])

   Angle_Energy = append(Angle_Energy,angle[:,newaxis],axis=1)

   Delay_Energy = append(Delay_Energy,dt[:,newaxis],axis=1)

   dt,angle,ener = observables_vs_delay(energy[cond],theta_arrival[cond],time[cond]/yr,weight[cond])

   Delay_vs_angle = append(Delay_vs_angle,dt[:,newaxis],axis=1)

   #  TIME AND ARRIVAL ANGLE DISTRIBUTION =======================================#
   nbBins = 200
   theta2,dndtheta = arrivalAngle(theta_arrival,weight,nbBins,theta_range)

   arrival_Angle = theta2[:,newaxis]
   arrival_Angle = append(arrival_Angle,dndtheta[:,newaxis],axis=1)

   delta_t,dNdt = timing(time,weight,nbBins)#,dt_range=[])

   Timing = delta_t[:,newaxis]
   Timing = append(Timing,dNdt[:,newaxis],axis=1)

   #  SPECTRUM (MEASURED) =======================================================#
   ener,flux = spectrum(energy,weight,nbBins,E_range=[Emin,Emax])
   Spectrum = ener[:,newaxis]
   Spectrum = append(Spectrum,flux[:,newaxis],axis=1)
   nbBins = 200
   if PSF == "Taylor2011":
      cond = (theta_arrival<PSF_Taylor2011(energy))
   else:
      cond = (theta_arrival<PSF)
   ener,flux = spectrum(energy[cond],weight[cond],nbBins,E_range=[Emin,Emax])
   Spectrum = append(Spectrum,flux[:,newaxis],axis=1)

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

   #  BY TIME RANGE
   #=============================================================================#
   print "   > Selection by time range ..." 

   tmax = [day,10*day,30*day] # seconds (1 month, 1 year, 5 years)

   for n in arange(0,3,1):
      cond= (time<tmax[n])

      contrib =sum(weight[cond])/NbTotEvents*100 
      print "     ... integration time =",float(tmax[n]),"s\t-> contribution:",int(contrib),"%"

      #  ARRIVAL ANGLE DISTRIBUTION  ============================================#
      nbBins = 200
      theta2,dndtheta = arrivalAngle(theta_arrival[cond],weight[cond],nbBins,theta_range)

      arrival_Angle = append(arrival_Angle,dndtheta[:,newaxis],axis=1)

      #  SPECTRUM (MEASURED) ====================================================#
      nbBins = 200
      if PSF == "Taylor2011":
         cond = (theta_arrival<PSF_Taylor2011(energy)) & (time<tmax[n]) 
      else:
         cond = (theta_arrival<PSF) & (time<tmax[n]) 
      ener,flux = spectrum(energy[cond],weight[cond],nbBins,E_range=[Emin,Emax])
      Spectrum = append(Spectrum,flux[:,newaxis],axis=1)

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

   #  BY GENERATION
   #=============================================================================#

   print "   > Selection by generation ..." 

   gen_tab =[0,2,4,6,8]#list(set(generation))

   Gen_cont = zeros([shape(gen_tab)[0],2])

   for gen in gen_tab:

      cond = (generation==gen) 

      contrib =sum(weight[cond])/NbTotEvents*100 
      i = gen_tab.index(gen)
      Gen_cont[i,0]=int(gen)
      Gen_cont[i,1]=contrib
      print "     ... gen=",int(gen),"-> contribution:",int(contrib),"%"

      #  PARAMETER SPACE (theta, dt, E) =========================================#

      ener,angle,dt = observables_vs_energy(energy[cond],theta_arrival[cond],time[cond]/yr,weight[cond])

      Angle_Energy = append(Angle_Energy,angle[:,newaxis],axis=1)

      Delay_Energy = append(Delay_Energy,dt[:,newaxis],axis=1)

      dt,angle,ener = observables_vs_delay(energy[cond],theta_arrival[cond],time[cond]/yr,weight[cond])

      Delay_vs_angle = append(Delay_vs_angle,dt[:,newaxis],axis=1)

      #  TIME AND ARRIVAL ANGLE DISTRIBUTION ==================================#

      nbBins = 200

      theta2,dndtheta = arrivalAngle(theta_arrival[cond],weight[cond],nbBins,theta_range)
      arrival_Angle = append(arrival_Angle,dndtheta[:,newaxis],axis=1)

      delta_t,dNdt = timing(time[cond],weight[cond],nbBins)

      Timing = append(Timing,dNdt[:,newaxis],axis=1)

      #  SPECTRUM (MEASURED) ====================================================#
      nbBins = 200
      if PSF == "Taylor2011":
         cond = (theta_arrival<PSF_Taylor2011(energy)) & (generation==gen) 
      else:
         cond = (theta_arrival<PSF) & (generation==gen) 
      ener,flux = spectrum(energy[cond],weight[cond],nbBins,E_range=[Emin,Emax])
      Spectrum = append(Spectrum,flux[:,newaxis],axis=1)
      
    
      print shape(Spectrum)

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

   print "   > writing files" 

   savetxt(output_dir+"/Spectrum.txt",Spectrum)

   savetxt(output_dir+"/arrival_Angle_distribution.txt",arrival_Angle)

   savetxt(output_dir+"/Timing.txt",Timing)

   savetxt(output_dir+"/Angle_versus_Energy.txt",Angle_Energy)
   savetxt(output_dir+"/Delay_versus_Angle.txt",Delay_vs_angle)
   savetxt(output_dir+"/Delay_versus_Energy.txt",Delay_Energy)

   savetxt(output_dir+"/Generation.txt",Gen_cont)
   print "#=============================================================================#"