#SNRvsCollimatorOpeningAngle.py                  Lance Simms, LLNL 2010
#
#PURPOSE:
#  To simulate a full 360 revolution of the CZT detector aboard the 
#  DXRB nanosat.  The satellite points at the sun, and the CZT detector
#  spins in a plane perpendicular to the sun pointing axis.
#
#KEYWORDS:
#  DateTime: str
#    The date and time for which we want to calculate the observation.  This should be 
#    in the pyephem prefered format of "Year/Month/Day Hour:Minute:Second" in UT.
#  Lat: str
#    The latitude of the satellite when it makes the observation
#  Long: str
#    The longitude of the satellite when it makes the observation
#  SpinPeriod: float
#    The number of seconds per full 360 degree rotation of the detector
#  NumAngles: int
#    The number of angles at which to record the hits in the detector
#  NumOrbits: int
#    The number of orbits to use for summing the angular information
#  Sources100TimesBrighter: bool
#    0 - Keep source fluxes normal; 1 - Increase source flux to show maximum in signal to 
#    noise plot
#
#EXAMPLES
#  run SNRvsCollimatorOpeningAngle --SpinPeriod=600
#  run SNRvsCollimatorOpeningAngle --SpinPeriod=1
#  run SNRvsCollimatorOpeningAngle --SpinPeriod=1 --Sources100TimesBrighter=1
#
########
import os
execfile(os.path.expanduser('~')+'/Desktop/Software/Python/CZTCubeSat/CZTCubeSatSetup.py')

import ephem
from datetime import *

##PARSE KEYWORDS
keywords    = ['DateTime=', 'Lat=', 'Long=', 'FOV_Deg=', 'SpinPeriod=', \
               'NumAngles=', 'NumOrbits=', 'Sources100TimesBrighter=']

#DEFAULT VALUES FOR KEYWORDS
DateTime    = '2012/01/01 20:00:00'
Lat         = '37.00'
Long        = '-122.00'
SpinPeriod  = 1
NumAngles   = 360
NumOrbits   = 1
Sources100TimesBrighter = 0

#NOW PARSE KEYWORDS
opts, extraparams = getopt.getopt(sys.argv[1:],'',keywords)
for o,p in opts:
  if o in ['--DateTime']:
    DateTime = p
  elif o in ['--Lat']:
    Lat = p
  elif o in ['--Long']:
    Long = p
  elif o in ['--SpinPeriod']:
    SpinPeriod = float(p)
  elif o in ['--NumAngles']:
    NumAngles = float(p)
  elif o in ['--NumOrbits']:
    NumOrbits = int(p)
  elif o in ['--Sources100TimesBrighter']:
    Sources100TimesBrighter = int(p)

#Read in the A4 Catalog first of all 
Names, RAs, DECs, CountRates, CountRateErrs, Fluxes, DataTabs, Observatories = \
  readcol(os.path.expanduser('~')+'/Desktop/Software/Python/CZTCubeSat/A4_Catalog_NoHeader.txt', \
          fsep="|",twod=False)

#Convert the RAs and DECs to decimal
Cat_Size = size(RAs)
RAs_Deg  = zeros(Cat_Size, dtype=float64)
DECs_Deg = zeros(Cat_Size, dtype=float64)
RAs_Rad  = zeros(Cat_Size, dtype=float64)
DECs_Rad = zeros(Cat_Size, dtype=float64)

for i in arange(len(RAs)):
  RAs_Deg[i]  = hms2str_deg_space(RAs[i])
  DECs_Deg[i] = dms2str_deg_space(DECs[i])
  RAs_Rad[i]  = (pi/180)*RAs_Deg[i]
  DECs_Rad[i] = (pi/180)*DECs_Deg[i]

#Convert the Count Rates to our 1cm x 2cm = 2cm^2 CZT Detector
EffectiveAreaCZT     = 2.   #Area of CZT in cm^2
EffectiveAreaHEAO_A4 = 80.  #Effective area of A4 (see Levine et al. 1984) 
CountRatesCZT        = CountRates*(EffectiveAreaCZT/EffectiveAreaHEAO_A4)

#Allow for brightness adjustment
if Sources100TimesBrighter == 1:
  CountRatesCZT = 1000000*CountRatesCZT

#This is the DXRB Flux calculated from the script CalculateBackgroundFlux.py (from 10-100keV)
FluxFrom10_100keV    = 5.25   #Flux in Counts/cm^2/sr/s

#Particle Background (non-rejected particles we estimate at 0.2 counts
ParticleBackground   = 1

#Format the time appropriately
Year  = int(DateTime[0:4])
Month = int(DateTime[5:7])
Day   = int(DateTime[8:10])
Hours = int(DateTime[11:13])
Mins  = int(DateTime[14:16])
Secs  = int(DateTime[17:19])
DateTime = datetime(Year, Month, Day, Hours, Mins, Secs)

#Create the observer object and assign date and location
ObsTech = ephem.Observer()
ObsTech.long, ObsTech.lat = float(Lat), float(Long)
ObsTech.date = DateTime

#Compute the position of the Sun and the Moon
sun, moon = ephem.Sun(), ephem.Moon()
sun.compute(ObsTech)
moon.compute(ObsTech)

#Print out the info
print "Sun's location :"
print sun.ra, sun.dec
print "Phi=", float(sun.ra), "Theta=", float(sun.dec)

#Assign angles of the Sun for testing purposes
theta = float(sun.dec)
phi   = float(sun.ra)
#theta = 0
#phi   = 0# pi/2#

#Make an array of the CZT Pointings
AngularInterval    = 360./NumAngles
CZTPointingAngles  = AngularInterval*arange(NumAngles)
TotHitsPerPointing = zeros(NumAngles, dtype=float32)
TimePerPointing    = float(SpinPeriod)/NumAngles

#Make an array for the Fields of View
NumFOVs       = 31
FOV_Int       = 3
FOV_Arr     = FOV_Int*(arange(NumFOVs))
SNR_Arr     = zeros(NumAngles, dtype=float32) 
Sig_Arr     = zeros(NumAngles, dtype=float32)
Noise_Arr   = zeros(NumAngles, dtype=float32)
DEC_Arr     = zeros(NumAngles, dtype=float32)
RA_Arr      = zeros(NumAngles, dtype=float32)
SNR_Avg_Arr = zeros(NumFOVs, dtype=float32)
Sig_Avg     = zeros(NumFOVs, dtype=float32)
Noise_Avg   = zeros(NumFOVs, dtype=float32)

for (FOV, FOV_Num) in zip(FOV_Arr, arange(NumFOVs)):

  #Print out the FOV
  print 'Field of View (Degrees) = ' + str(FOV)

  #Recalculate the DXRB Signal with the new opening angle
  CountRateDXRB = FluxFrom10_100keV*(EffectiveAreaCZT*2*pi*(1-cos(FOV*pi/180)))
 
  #Go through a number of orbits
  for OrbitNum in arange(NumOrbits):

    #Skip the FOV of zero
    if FOV == 0 : break
  
    #Print out the orbit number
    print 'Orbit Number : ' + str(OrbitNum)
  
    #Now go through and calculate these angles rotated to the sun pointed frame  
    for (AngleNum,CZTPointingAngle) in zip(arange(NumAngles), CZTPointingAngles):
    
      #Determine the angle in the unrotated plane
      x=cos(CZTPointingAngle*pi/180)
      y=sin(CZTPointingAngle*pi/180)
    
      #Rotate about the Declination axis with an angle of ra
      x_prime  = x*math.cos(phi)-y*math.sin(phi)
      y_prime  = x*math.sin(phi)+y*math.cos(phi)
    
      #Rotate about the Right Ascension=0 axis with an angle of dec
      #y-axis points to 0 RA, so rotate about x-axis
      x_dprime = x_prime
      y_dprime = y_prime*math.cos(theta)
      z_dprime = y_prime*math.sin(theta)
      #print str(x_dprime) + ' , ' + str(y_dprime) + ' , ' + str(z_dprime)
      #Account for poor handling of cosine and sine
      if (abs(y_dprime) < 1e-10) : y_dprime = 0
    
      #Now get the angles corresponding to the vector in the original frame
      DEC_perp = math.acos(z_dprime)-pi/2 
      RA_perp  = math.atan2(y_dprime,x_dprime)
      if (RA_perp < 0) : RA_perp = 2*pi+RA_perp
      if RA_perp == 0: 
        if ((AngleNum >= 90) and (AngleNum < 270)) : RA_perp=RA_perp+pi

      #Determine number of hits from background
      NumHitsFromDXRB  = CountRateDXRB*TimePerPointing
      NumHitsParticles = ParticleBackground*TimePerPointing
      NumHitsSources   = 0
      NumSources = 0
      #Now find all of the objects within FOV_Deg degrees of pointing
      for i in arange(Cat_Size):
    
        #Determine separation between CZT Pointing and Object
        ThisAngSep  = Astro_Angular_Separation(RA_perp, DEC_perp, RAs_Rad[i], DECs_Rad[i]) 
        ThisAngSep  = (180/pi)*ThisAngSep
        #If object is in FOV, calculate flux (no cos dependence yet)
        if (abs(ThisAngSep) < FOV):
    
          #See how many hits we expect 
          NumSources+=1
          NumHitsSources += CountRatesCZT[i]*TimePerPointing*cos((ThisAngSep/180.)*pi)

      #Now calculate the SNR at this angle
      Signal = NumHitsFromDXRB
      Noise  = sqrt(NumHitsFromDXRB)+NumHitsSources+NumHitsParticles
      SNR_Arr[AngleNum] = (Signal/Noise)
      Sig_Arr[AngleNum] = Signal
      Noise_Arr[AngleNum] = Noise
      RA_Arr[AngleNum]  = RA_perp
      DEC_Arr[AngleNum] = DEC_perp
      #print str(NumHitsFromDXRB)+' , '+str(NumHitsParticles)+' , '+str(NumHitsSources)
      #print str(Signal) + ' , ' + str(Noise)
      #print str(SNR_Arr[AngleNum]) 
 
  #Now get the average signal to noise over the orbit
  SNR_Avg_Arr[FOV_Num] = mean(SNR_Arr) 
  Sig_Avg[FOV_Num]     = mean(Sig_Arr)
  Noise_Avg[FOV_Num]   = mean(Noise_Arr)
  #mplot.clf()
  #mplot.plot(Noise_Arr)
  #mplot.plot(RA_Arr)
  #mplot.plot(DEC_Arr)
  #pdb.set_trace()

#Now plot all the angles
Fig=mplot.figure(0)
mplot.clf()
Fig.add_axes([0.150,0.100,0.80,0.6750])
mplot.plot(FOV_Arr, SNR_Avg_Arr, label=str(SpinPeriod))
mplot.xlabel("Collimator Half-Opening Angle (Degrees)")
mplot.ylabel("Signal To Noise Ratio")
mplot.figtext(0.08,0.94,r'Date      : '+str(DateTime))
mplot.figtext(0.08,0.89,r'Sun RA   : '+str(sun.ra))
mplot.figtext(0.08,0.84,r'Sun DEC : '+str(sun.dec))
mplot.figtext(0.53,0.94,r'Spin Period : '+str(SpinPeriod)+' seconds')
mplot.figtext(0.53,0.89,r'Num Orbits  : '+str(NumOrbits))

#Save the figure
SaveFig = 0
if SaveFig == 1:
  if Sources100TimesBrighter == 0:
    mplot.savefig('./Figures/SignalToNoiseRatioFor_'+str(NumOrbits)+'_Orbits'+'_SpinPeriod_'+\
                  str(SpinPeriod)+'_s.png')
  else:
    mplot.figtext(0.35,0.70,r'SOURCES 1,000,000 TIMES BRIGHTER')
    mplot.savefig('./Figures/SignalToNoiseRatioFor_'+str(NumOrbits)+'_Orbits'+'_SpinPeriod_'+\
                  str(SpinPeriod)+'_s_Sources100000TimesBrighter.png')
