#!/bin/python

import os, shutil
from xml.dom               import minidom
from numpy                 import append, savetxt, shape, array, newaxis, zeros, arange, select
from numpy                 import where, size, log10, ones
from src.read              import select_events, ReadProfile, resultdir
from src.distribution      import distribution
from src.image             import image
from src.observables       import observables_vs_delay, observables_vs_energy
from src.constants         import degre, day, yr

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

for simu in simulations:
   print "#=============================================================================#"
   fileId = simu.getAttribute("simulation_dir")
   print " Reading data from ", fileId 

   output_dir = resultdir+simu.getAttribute("id")+"/"
   if not os.path.exists(output_dir):
      os.makedirs(output_dir)
   # copy profile
   shutil.copy(fileId+"/profile.dat",output_dir+"profile.dat")
   print " Output directory: ", output_dir 

   print " Applying selection:"
   Emin = float(simu.getAttribute("Emin")) 
   Emax = float(simu.getAttribute("Emax"))
   NbinsE=int(log10(Emax)-log10(Emin))*15
   print "    > Energy range: [", Emin,"GeV,",Emax*1e-3,"TeV]" 
   powerlaw_index = simu.getAttribute("powerlaw_index") 
   if "simple case" not in fileId:
      print "    > Source spectrum:", powerlaw_index

   tjet = float(simu.getAttribute("tjet"))
   print "    > jet opening angle:", tjet, "degre" 
   tobs = float(simu.getAttribute("tobs"))
   print "    > observation angle:", tobs, "degre" 

   theta_range = [5e-7*degre,90*degre] # degre
   dt_range = [1e-4,1e17] # s
   E_range= [Emin,Emax] # GeV

   print "# Simulation parameters "
   EGMF,E0,Dsource,redshift,n_phot,n_lept  = ReadProfile(fileId) 
   print "    > Source redshift: ", redshift, "(eq.: ",Dsource,"Mpc)"
   print "    > EGMF:            ", EGMF, "G"
   print "    > run over         ", int(n_phot), " photons"

   # read files
   weight, energy, time, theta_d, phi_d, Esource, generation =  select_events(fileId,
         Erange=[Emin,Emax],powerlaw_index=powerlaw_index,tjet=tjet,tobs=tobs)
   NbTotEvents = sum(weight)
   print " ----------------------------------------------------------------------------- "

   if "simple case" not in fileId:
      print "   > Source spectrum ..." 
      Es=array(list(set(Esource)))
      Ws= (Es/min(Es))**(1-float(powerlaw_index))  
      ener,flux = distribution(Es,Ws,NbinsE,E_range)
      Spectrum = ener[:,newaxis]
      Spectrum = append(Spectrum,flux[:,newaxis],axis=1)
      savetxt(output_dir+"/Source_spectrum.txt",Spectrum)

   #=============================================================================#
   #  NO SELECTION
   #=============================================================================#
   print "   > No selection of events ..." 
   #  IMAGING ===================================================================#
   step = 6
   PSF = 20
   image(theta_d,phi_d,weight,output_dir,step,borne=[PSF,PSF])

   #  PARAMETER SPACE (theta, dt, E) ============================================#
   nbBins = 50
   ener,angle,dt  = observables_vs_energy(energy,theta_d,time,weight)
   Angle_Energy   = ener[:,newaxis]
   Angle_Energy   = append(Angle_Energy,angle[:,newaxis]/degre,axis=1)
   Delay_Energy   = ener[:,newaxis]
   Delay_Energy   = append(Delay_Energy,dt[:,newaxis],axis=1)
   dt,angle,ener  = observables_vs_delay(energy,theta_d,time,weight)
   Delay_vs_angle = angle[:,newaxis]/degre
   Delay_vs_angle = append(Delay_vs_angle,dt[:,newaxis],axis=1)

   #  DISTRIBUTIONS ============================================================#
   nbBins = 80
   theta2,dndtheta = distribution(theta_d,weight,nbBins,theta_range)
   arrival_Angle = theta2[:,newaxis]
   arrival_Angle = append(arrival_Angle,dndtheta[:,newaxis],axis=1)
   delta_t,dNdt = distribution(time,weight,nbBins,dt_range)
   Timing = delta_t[:,newaxis]
   Timing = append(Timing,dNdt[:,newaxis],axis=1)
   ener,flux = distribution(energy,weight,NbinsE,E_range)
   Spectrum = ener[:,newaxis]
   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_d[cond],time[cond],weight[cond])
      Angle_Energy = append(Angle_Energy,angle[:,newaxis]/degre,axis=1)
      Delay_Energy = append(Delay_Energy,dt[:,newaxis],axis=1)
      dt,angle,ener = observables_vs_delay(energy[cond],theta_d[cond],time[cond],weight[cond])
      Delay_vs_angle = append(Delay_vs_angle,dt[:,newaxis],axis=1)

       #  DISTRIBUTIONS =========================================================#
      theta2,dndtheta = distribution(theta_d[cond],weight[cond],nbBins,theta_range)
      arrival_Angle = append(arrival_Angle,dndtheta[:,newaxis],axis=1)
      delta_t,dNdt = distribution(time[cond],weight[cond],nbBins,dt_range)
      Timing = append(Timing,dNdt[:,newaxis],axis=1)
      ener,flux = distribution(energy[cond],weight[cond],NbinsE,E_range)
      Spectrum = append(Spectrum,flux[:,newaxis],axis=1)

   #=============================================================================#
   #  BY TIME RANGE
   #=============================================================================#
   print "   > Selection by time range ..." 
   tmax = [day, 30*day, yr, 1e3*yr, 1e6*yr] # yr 

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

      theta2,dndtheta = distribution(theta_d[cond],weight[cond],nbBins,theta_range)
      arrival_Angle = append(arrival_Angle,dndtheta[:,newaxis],axis=1)
      ener,flux = distribution(energy[cond],weight[cond],NbinsE,E_range)
      Spectrum = append(Spectrum,flux[:,newaxis],axis=1)

   #=============================================================================#
   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 "#=============================================================================#"