#!/usr/stsci/pyssg/Python-2.5.4/bin/python
#
#NAME:
#  SkyCoverageAndSN.py
#
#PURPOSE:
#  To simulate the sky coverage seen by the CXBN over the mission lifetime by 
#  carrying out the orbit of the CXBN over that lifetime.
#
#KEYWORDS:
#  DateTime: str
#    The date and time for at which the simulation will start.  This should be 
#    in the pyephem prefered format of "Year/Month/Day Hour:Minute:Second" in UT.
#    The date and time for which we want to calculate the observation.  This should be 
#  TLE_File: str
#    The path to the TLE file that will be used as input for creating the PyEphem 
#    satellite object
#  FOV_Deg: float
#    The collimating angle of the instrument.  This will determine the threshold for 
#    angles at which a source is considered detected.
#  NumAngles: int
#    The number of angles at which to record the hits in the detector
#  CalculateSN: int
#    0 - Don't calculate SNR with the source catalog
#    1 - Calculate SNR with the source catalog
#  CalculateSkyCoverage: int
#    0 - Don't make an exposure map of the sky
#    1 - Make an exposure map of the sky
#
#NOTES:
#  The pyephem package actually checks the CheckSum on each line of the TLE file.  One must make 
#  sure that the last digit of the line equals the sum of all digits (modulo 10).  One easy way 
#  to make this true is just to play with the last digit until it works.
#
#  There are three reference coordinate systems used here:
#    e = Earth Centered Inertial Coordinate System
#    r = Satellite Reference Frame (-z points toward center of earth, +x points in direction of velocity)
#    b = Body Coordinate System (axes point along principal axes)
#
#EXAMPLES:
#  run SkyCoverage 
#
###############################################################################################################

import os
execfile(os.path.expanduser('~')+'/Desktop/Software/Python/CZTCubeSat/CZTCubeSatSetup.py')

import ephem
from datetime import *
from CXBN_Utilities import *

# Main routine if we are called as a script
if __name__ == "__main__":

  ##PARSE KEYWORDS
  keywords    = ['DateTime=', 'TLE_File=', 'FOV_Deg=', 'NumAngles=', 'CalculateSN=', 'CalculateSkyCoverage=']
  #DEFAULT VALUES FOR KEYWORDS
  DateTime        = '2011/01/01 00:00:00'
  TLE_File        = './TLE/CXBNTLE.txt'
  FOV_Deg         = 18
  NumAngles       = 360
  CalculateSN     = 1
  CalculateSkyCoverage = 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 ['--TLE_File']:
      TLE_File = p
    elif o in ['--FOV_Deg']:
      FOV_Deg = float(p)
    elif o in ['--NumAngles']:
      NumAngles = float(p)
    elif o in ['--CalculateSN']:
      CalculateSN = int(p)
    elif o in ['--CalculateSkyCoverage']:
      CalculateSkyCoverage = int(p)

  #############################################################################################################
  ################################ SATELITTE AND CONSTANT OBJECTS #############################################

  #Create an instance of the constant object
  c = c()

  #Create a satellite object
  cxbn = CubeSat()

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

  #############################################################################################################
  ################################ BAT-SWIFT SOURCES ##########################################################
  #Read in the A4 Catalog first of all 
  SourceNums, Names, RA_Sources, DEC_Sources, BAT_S_N, CounterpartName, OtherNames, CtPt_RAs, CtPt_DECs, Flux1s, \
    Flux2s, Flux3s, Contaminations, Gamma1s, Gamma2s, Gamma3s, ChiSqs, Redshifts, Luminosities, Classes, Types= \
    readcol(os.path.expanduser('~')+'/Desktop/Software/Python/CZTCubeSat/BAT_Catalog.txt', \
            fsep="|",twod=False)

  #Get the fluxes as an array
  Flux1Arr = array(Flux1s)
  Flux2Arr = array(Flux2s)
  Flux3Arr = array(Flux3s)

  #Convert the RAs and DECs to decimal
  Cat_Size = size(RA_Sources)
  RA_Sources_Deg  = zeros(Cat_Size, dtype=float64)
  DEC_Sources_Deg = zeros(Cat_Size, dtype=float64)
  RA_Sources_Rad  = zeros(Cat_Size, dtype=float64)
  DEC_Sources_Rad = zeros(Cat_Size, dtype=float64)

  for i in arange(len(RA_Sources)):
    RA_Sources_Deg[i]   = float(RA_Sources[i])
    if (RA_Sources_Deg[i] > 180) : RA_Sources_Deg[i] = RA_Sources_Deg[i]-360
    DEC_Sources_Deg[i]  = float(DEC_Sources[i])
    RA_Sources_Rad[i]   = (pi/180)*RA_Sources_Deg[i]
    DEC_Sources_Rad[i]  = (pi/180)*DEC_Sources_Deg[i]

  #############################################################################################################
  ################################ DATE AND TIME ##############################################################
  #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])
  DateStr  = DateTime[0:10]
  DateTime = datetime(Year, Month, Day, Hours, Mins, Secs)

  #Convert this into a PyEphem Date (# days since noon on 1899 December 31) 
  EpDate_0 = ephem.date(DateTime)
  EpDate   = EpDate_0

  #############################################################################################################
  ################################ PyEphem OBJECTS for CXBN ## ################################################

  #Call getsat to return an ephemeris satellite object based on the TLE
  CXBN  = getsat(TLE_File, EpDate)

  #Create the observer object and assign date and location
  ObsTech = ephem.Observer()

  #Get the initial position vector of the satellite
  CXBN.compute(EpDate)
  r_eci      = getSatECI(CXBN, c)
  r_eci_prev = r_eci

  #############################################################################################################
  ############################### TIME VECTORS ################################################################

  #Set the time increment for the propogation
  dt       = double(3600*24)                                #Time difference in seconds (Make it a day)
  t_final  = 3600*24*366.                                   #Final time in seconds
  NumTimes = ceil(t_final/dt)
  Times    = arange(NumTimes)

  #Make an array of the CZT Pointings
  AngularInterval    = 360./NumAngles                       #Anglular interval during a full sweep
  CZTPointingAngles  = AngularInterval*arange(NumAngles)    #Array that holds these angles
  FOV_Rad            = FOV_Deg*pi/180                       #The field of view in radians       

  #Figure out the time spent scanning a given portion of the sky 
  SpinPeriod              = 6                               #Spin period for sky sweep in seconds
  NumSpinsInTimeStep      = dt/SpinPeriod                   #Number of spins during a day
  TimePerAnglePerTimeStep = dt/NumAngles                    #TimePerAnglePerTimeStep
  f_EarthOccult           = 0.5                             #Fraction of time the earth will be blocking

  #Convert the Flues into a count rate for the CZT Detector on CXBN (assume average energy is 100 keV)
  Flux1CXBN = Flux1Arr*1e-12*(6.241509e11)*cxbn.CZT_EffectiveArea*TimePerAnglePerTimeStep*f_EarthOccult/100e3
  Flux2CXBN = Flux2Arr*1e-12*(6.241509e11)*cxbn.CZT_EffectiveArea*TimePerAnglePerTimeStep*f_EarthOccult/100e3
  Flux3CXBN = Flux3Arr*1e-12*(6.241509e11)*cxbn.CZT_EffectiveArea*TimePerAnglePerTimeStep*f_EarthOccult/100e3

  #Add the particle background rate
  ParticleRate = 0.5                                        #The particle background rate in units of counts/s

  #Set the initial time to t=0
  Time=0

  #############################################################################################################
  ############################### TRACKING ARRAYS  ############################################################
  #Arrays for tracking variables (not necessary for functioning)
  X_Arr       = zeros(NumTimes, dtype=float32)
  B_Arr       = zeros(NumTimes, dtype=float32)
  Sun_RA_Arr  = zeros(NumTimes, dtype=float32)
  Sun_DEC_Arr = zeros(NumTimes, dtype=float32)

  #Make a grid of points for the RA and DEC
  NumGridPoints     = 61
  RA_Grid_x         = frange(NumGridPoints-1)*2*pi/(NumGridPoints-1)-pi
  DEC_Grid_y        = frange(NumGridPoints-1)*pi/(NumGridPoints-1)-pi/2
  RA_Grid, DEC_Grid = meshgrid(RA_Grid_x, DEC_Grid_y)
  SkyCoverageGrid   = zeros([NumGridPoints, NumGridPoints], dtype=int32)
  SignalGrid        = zeros([NumGridPoints, NumGridPoints], dtype=float32)
  NoiseGrid         = zeros([NumGridPoints, NumGridPoints], dtype=float32)
  NoiseGrid[:]      = 1

  #############################################################################################################
  ############################### BASIS VECTORS AND COORDINATE SYSTEMS#########################################
  #Set up the basis vectors for the inertial frame
  i   = vec3(1,0,0)   ; j   = vec3(0,1,0)    ; k   = vec3(0,0,1)
  e_i = vec3(1,0,0)   ; e_j = vec3(0,1,0)    ; e_k = vec3(0,0,1)
  s_i = vec3(1,0,0)   ; s_j = vec3(0,1,0)    ; s_k = vec3(0,0,1)  

  #############################################################################################################
  #PROPAGATE THE SATELLITE IN ITS ORBIT
  #############################################################################################################
  for t_i in Times:  

    #Inrement the time 
    EpDate = double(EpDate) + ephem.second*dt
    Time  += dt

    #Recompute the solution with the SGP4
    CXBN.compute(EpDate)

    #Now update the observer object
    ObsTech.long      = float(CXBN.sublat)
    ObsTech.lat       = float(CXBN.sublong)
    ObsTech.elevation = float(CXBN.elevation)
    ObsTech.date      = EpDate
    print ObsTech.date

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

    #Assign angles of the Sun for testing purposes
    sun_DEC  = float(sun.dec)
    sun_RA   = float(sun.ra)
    theta_s  = pi/2-sun_DEC
    phi_s    = sun_RA

    #Keep track of the RA and DEC
    Sun_DEC_Arr[t_i] = sun_DEC
    Sun_RA_Arr[t_i]  = sun_RA-pi
    print "Sun RA and DEC" + str(sun_RA) + ' , ' + str(sun_DEC)

    #Now find the rotation matrix that brings the +z axis to the sun pointing vector
    sun_rot_mat = mat3(cos(theta_s)*cos(phi_s)  , -sin(phi_s) , sin(theta_s)*cos(phi_s), \
                       cos(theta_s)*sin(phi_s)  ,  cos(phi_s) , sin(theta_s)*sin(phi_s), \
                      -sin(theta_s)             ,  0          , cos(theta_s))

    #Form the basis vectors in this rotated frame
    x_hat_s = sun_rot_mat*i
    y_hat_s = sun_rot_mat*j
    z_hat_s = sun_rot_mat*k

    #Get the position of the satellite in Earth-Centered Inertial Coordinates
    r_eci = getSatECI(CXBN, c) 
    v_eci = (1/dt)*(r_eci - r_eci_prev)
    r_eci_prev = r_eci

    #Now go through and calculate these angles rotated to the sun pointed frame  
    for (AngleNum,CZTPointingAngle) in zip(arange(NumAngles), CZTPointingAngles):

      #Convert to radians
      CZTPointingAngleRad = CZTPointingAngle*pi/180

      #Form the rotating vector in the perpendicular plane
      ThisPerpVec = cos(CZTPointingAngleRad)*x_hat_s+sin(CZTPointingAngleRad)*y_hat_s

      #Account for poor handling of cosine and sine
      if (abs(ThisPerpVec.x) < 1e-10) : ThisPerpVec.x = 0
      if (abs(ThisPerpVec.y) < 1e-10) : ThisPerpVec.y = 0
      if (abs(ThisPerpVec.z) < 1e-10) : ThisPerpVec.z = 0

      #Now project it onto the various axes
      ThisPerpVecBasisX = ThisPerpVec*vec3(1,0,0)
      ThisPerpVecBasisY = ThisPerpVec*vec3(0,1,0)
      ThisPerpVecBasisZ = ThisPerpVec*vec3(0,0,1)

      #Now get the angles corresponding to the vector in the original frame
      DEC_perp = acos(ThisPerpVecBasisZ)-pi/2
      RA_perp  = atan2(ThisPerpVecBasisY,ThisPerpVecBasisX)
      #print RA_perp
      #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

      #Now Try to figure out all the sources within the Field Of View
      if (CalculateSN != 0):
        ThisNoise  = ParticleRate  
        ThisSignal = TimePerAnglePerTimeStep*cxbn.CountRateDXRB*f_EarthOccult
        SourcesWithDECInFOVInds = (where(abs(DEC_perp-DEC_Sources_Rad) < FOV_Rad))[0]
        for ThisSourceWithDECInFOVInd in SourcesWithDECInFOVInds:
          #Get the RA and DEC of the source we're looking at
          ThisRA  = RA_Sources_Rad[ThisSourceWithDECInFOVInd]
          ThisDEC = DEC_Sources_Rad[ThisSourceWithDECInFOVInd]

          #Determine separation between CZT Pointing and Object
          ThisAngSep  = Astro_Angular_Separation(RA_perp, DEC_perp, ThisRA, ThisDEC)
          if (abs(ThisAngSep) < FOV_Rad):
            ThisNoise  += Flux1CXBN[ThisSourceWithDECInFOVInd]*cxbn.CZT_EffectiveArea*abs(cos(ThisAngSep))

        #Find out what part of the grid this is closest to
        DEC_Diffs = abs(DEC_perp-DEC_Grid_y)
        ThisDECInd = (where(DEC_Diffs == min(DEC_Diffs)))[0]
        ThisDEC = DEC_Grid_y[ThisDECInd]
        for ThisRAInd in arange(NumGridPoints):
          ThisRA = RA_Grid_x[ThisRAInd]
          ThisAngSep  = Astro_Angular_Separation(RA_perp, DEC_perp, ThisRA, ThisDEC)
          if (abs(ThisAngSep) < 0.05):
            #RA_Diffs  = abs(RA_perp-RA_Grid_x)
            #ThisRAInd = (where(RA_Diffs == min(RA_Diffs)))[0]
            SignalGrid[ThisDECInd, ThisRAInd] += ThisSignal
            NoiseGrid[ThisDECInd, ThisRAInd]  += ThisNoise

      #Figure out what areas of the sky are in the FOV
      if (CalculateSkyCoverage != 0):
        DECIndsInFOV = (where(abs(DEC_perp-DEC_Grid_y) < FOV_Rad))[0]
        for ThisDECInd in DECIndsInFOV:
          ThisDEC = DEC_Grid_y[ThisDECInd]
          for ThisRAInd in arange(NumGridPoints):
            ThisRA = RA_Grid_x[ThisRAInd]

            #Determine separation between CZT Pointing and Object
            ThisAngSep  = Astro_Angular_Separation(RA_perp, DEC_perp, ThisRA, ThisDEC)
            if (abs(ThisAngSep) < FOV_Rad): SkyCoverageGrid[ThisDECInd, ThisRAInd]+=1*TimePerAnglePerTimeStep*f_EarthOccult

#######PLOT SIGNAL TO NOISE
if (CalculateSN != 0):
  #Now plot all the angles
  Fig1=mplot.figure(1, figsize=[16,8])
  mplot.clf()
  mplot.subplot(111,projection="mollweide")
  mplot.grid(True)

  #Set up an array
  image=mplot.imshow(SignalGrid, extent=(-pi,pi,-pi/2,pi/2), aspect=0.5)
  mplot.xlabel('Right Ascension (Degrees)')
  mplot.ylabel('Declination (Degrees)')
  cb = Fig1.colorbar(image, orientation='horizontal')
  
  #Now plot all the angles
  Fig2=mplot.figure(2, figsize=[16,8])
  mplot.clf()
  SP=mplot.subplot(111)#,projection="aitoff")
  mplot.grid(True)

  #Set up an array
  image=mplot.imshow(NoiseGrid, extent=(-180,180,-180/2,180/2))
  mplot.xlabel('Right Ascension (Degrees)')
  mplot.ylabel('Declination (Degrees)')
  cb = Fig2.colorbar(image, orientation='horizontal')

  mplot.scatter(RA_Sources_Deg, DEC_Sources_Deg, s=200*sqrt(Flux1Arr)/max(sqrt(Flux1Arr)), c='c')
  xlim([-180,180])
  ylim([-90,90])

  #Now plot all the angles
  Fig3=mplot.figure(3, figsize=[16,8])
  mplot.clf()
  FirstAx=Fig3.add_axes([0.05, 0.25, .95, 0.70])
  #SP=mplot.subplot(111)#,projection="aitoff")
  mplot.grid(True)

  #Set up an array
  image=mplot.imshow(SignalGrid/NoiseGrid, extent=(-180,180,-180/2,180/2))
  mplot.scatter(RA_Sources_Deg, DEC_Sources_Deg, s=200*sqrt(Flux1Arr)/max(sqrt(Flux1Arr)), c='c')
  mplot.xlabel('Right Ascension (Degrees)')
  mplot.ylabel('Declination (Degrees)')
  xlim([-180,180])
  ylim([-90,90])
  draw()
  cax = axes([0.125, 0.10, 0.775, 0.075]) 
  cb  = Fig3.colorbar(image, orientation='horizontal', cax=cax)
  cb.set_label('Signal to Noise')
  mplot.savefig('CXBN_SNR_FullSkyMap_OneYear.png')

#######PLOT SKY COVERAGE
if (CalculateSkyCoverage != 0):
  #Now plot all the angles
  Fig=mplot.figure(0, figsize=[10,5])
  mplot.clf()
  mplot.subplot(111,projection="mollweide")
  mplot.grid(True)
 
  #Set up an array
  image=mplot.imshow(SkyCoverageGrid, extent=(-pi,pi,-pi/2,pi/2), aspect=0.5)
  mplot.xlabel('Right Ascension (Degrees)')
  mplot.ylabel('Declination (Degrees)')
  cb = Fig.colorbar(image, orientation='horizontal')
  cb.set_label('Exposure time')
  draw()
  #Plot the sun positions
  mplot.plot(Sun_RA_Arr, Sun_DEC_Arr, 'yo')
  mplot.scatter(RA_Sources_Rad, DEC_Sources_Rad, s=200*sqrt(Flux1Arr)/max(sqrt(Flux1Arr)), c='c')
