#!/usr/bin/python
# Copyright (C) 2012  Alan Weinstein
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# A copy of the GNU General Public License may be found at
# http://www.gnu.org/copyleft/gpl.html
# or write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

"""
asd_ranges.py
This is example python code to read in a bunch of 
detector noise amplitude spectral densities (ASDs),
and compute the binary neutron star (BNS) detectability range.
"""

##############################################
# import everything we need, up front:
import numpy as np
import matplotlib
matplotlib.use('Agg')     # This lets us make plots without a display.
from matplotlib import pyplot as plt
import h5py

# First, generate and read in different noise spectra.

# The frequency vector over the detector band, in Hz:
fmin = 10.
fmax = 4096.
df = 0.1
f = np.arange(fmin,fmax,df)

# make a big array to hold a bunch of different noise spectra (well, 5)
N = 5
nspec = np.zeros([f.size, N])
fmin = np.zeros([N])
Fav = np.zeros([N])
tspec = [None]*N

# Official Initial LIGO Science Requirements Document noise curve:
# https://dcc.ligo.org/cgi-bin/DocDB/ShowDocument?docid=23947
# http://www.ligo.caltech.edu/~lazz/distribution/LSC_Data/SRD_strain_4k.txt
fname = 'SRD_strain_4k.txt'
ff = np.loadtxt(fname).reshape(97,2)
# Interpolate it to our frequency vector, and sore in our big array.
# Because this is sparsely sampled, do the interpolation on log(A):
nspec[:,0] = np.exp(np.interp(f,ff[:,0],np.log(ff[:,1])))
# The minimum frequency for which the above noise curve is valid:
fmin[0] = 40.     # Hz
tspec[0] = "iLIGO SRD"
Fav[0] = 0        # antenna factor for GRB051103

# Advanced LIGO ASD model, from LIGO Document T0900288:
# https://dcc.ligo.org/cgi-bin/DocDB/ShowDocument?docid=2974
fname = 'ZERO_DET_high_P.txt'
ff = np.loadtxt(fname)
nspec[:,1] = np.exp(np.interp(f,ff[:,0],np.log(ff[:,1])))
fmin[1] = 10.     # Hz
tspec[1] = "aLIGO HP 0-detune"
Fav[1] = 0        # antenna factor for GRB051103

# S6 H1 ASD (mode), from LIGO Document T1100338:
# https://dcc.ligo.org/cgi-bin/DocDB/ShowDocument?docid=63432
fname = 'H1-SPECTRA-957474700-MODE.txt'
ff = np.loadtxt(fname)
# note that we stored the PSD; the ASD is the sqrt
nspec[:,2] = np.exp(np.interp(f,ff[:,0],np.log(np.sqrt(ff[:,1]))))
fmin[2] = 40.     # Hz
tspec[2] = "S6 H1"
Fav[2] = 0        # antenna factor for GRB051103

# Noise spectrum from 2190-s of data from H2, near GRB051103 trigger:
fname = 'data_products/H2-SPEC_815043278-2190.hdf5'
fd = h5py.File(fname,'r')
ff = fd['H2:SPECTRUM']
nspec[:,3] = np.exp(np.interp(f,ff[:,0],np.log(ff[:,1])))
fd.close()
fmin[3] = 40.     # Hz
tspec[3] = "H2 SPEC"
Fav[3] = 0.5128   # antenna factor for GRB051103

# Noise spectrum from 2190-s of data from L1, near GRB051103 trigger:
fname = 'data_products/L1-SPEC_815043278-2190.hdf5'
fd = h5py.File(fname,'r')
ff = fd['L1:SPECTRUM']
nspec[:,4] = np.exp(np.interp(f,ff[:,0],np.log(ff[:,1])))
fd.close()
fmin[4] = 40.     # Hz
tspec[4] = "L1 SPEC"
Fav[4] = 0.4628   # antenna factor for GRB051103

# These are all amplitude spectral densities. Square to get power spectra:
sspec = np.power(nspec,2)

# Now compute the binary neutron star (BNS)
# detectability range for each of these spectra:

## constants
clight = 2.99792458e8	             # m/s
G = 6.67259e-11	                     # m^3/kg/s^2
parsec = 3.08568025e16               # m
MSol = 1.989e30                      # kg
tSol = MSol*G/np.power(clight,3)     # s

## Masses in solar masses
m1 = 1.4
m2 = 1.4
mtot = m1+m2
mred = (m1*m2)/mtot
mchirp = np.power(mred,3./5.)*np.power(mtot,2./5.)

# distance to a fiducial BNS source:
dist = 1.0                           # in Mpc
Dist =  dist * 1.0e6 * parsec /clight # from Mpc to seconds

# In stationary phase approx, this is htilde(f):
# See FINDCHIRP Eqns 3.4, or 8.4-8.5 ( http://arxiv.org/pdf/gr-qc/0509116.pdf )
hfe = (2.*tSol/Dist)*np.power(mchirp,5./6.)*np.sqrt(5./96./np.pi)*(np.pi*tSol)
hfe *= np.power(np.pi*tSol*f,-7./6.)
hfe2 = np.power(hfe,2)

# Innermost stable circular orbit (ISCO) parameter
Risco = 6.      # 6M for PN ISCO; 2.8M for EOB
# frequency at ISCO (end the chirp here; the merger and ringdown follow)
fisco = 1./(np.power(Risco,1.5)*np.pi*tSol*mtot)
# Single-detector SNR for detection above noise background:
SNR = 8.

# compute "inspiral horizon distance" for optimally oriented binary
# FINDCHIRP Eqn D2:
D = np.zeros(5)
for I in np.arange(0,5):
    fr = np.nonzero(np.logical_and(f > fmin[I] , f < fisco))
    D[I] = np.sqrt(4.*np.sum(hfe2[fr]/sspec[fr,I])*df)/SNR
    # print out the "inspiral horizon distance", 
    # the inspiral distance averaged over source location and orientation
    # (FINDCHIRP Appendix D), and the inspiral distance for GRB051103
    print(tspec[I]+': '+str(D[I])+', '+str(D[I]*0.4425)+', '+str(D[I]*Fav[I]))

# A BNS waveform amplitude (h(f)*sqrt(f)) at distance Dist from fmin=10 to fisco
# for overlay onto the ASD curves:
Dist = 3.6  # Mpc
fr = np.nonzero(np.logical_and(f > 10 , f < fisco))
ffr = f[fr]
hfr = 2.*hfe[fr]/Dist*f[fr]**0.5     # factor of 2 because of matched filter.
hfr *= 0.49                          # antenna response averaged over H2 and L1

# now plot all the noise ASDs, with a BNS waveform overlaid
plt.figure(1)
# These are our "official" colors for H1 (r), H2 (b), L1 (b)
plt.loglog(f,nspec[:,0],'k')
plt.loglog(f,nspec[:,1],'c')
plt.loglog(f,nspec[:,2],'r')
plt.loglog(f,nspec[:,3],'b')
plt.loglog(f,nspec[:,4],'g')
# Overlay a BNS waveform
plt.loglog(ffr,hfr,'k--')  # factor of 2 because of matched filter.
plt.xlim([10.,4096.])
plt.ylim([1.e-24,1.e-18])
plt.xlabel('frequency (Hz)')
plt.ylabel('Amplitude strain noise (1/rtHz), and BNS signal')
plt.grid(1)
#plt.legend(('iLIGO SRD: %.1f Mpc' % D[0],'aLIGO HP 0-detune: %3d Mpc' % D[1], \
#		   'S6 H1: %.1f Mpc' % D[2], \
#		   'H2 SPEC: %.1f Mpc' % D[3],'L1 SPEC: %.1f Mpc' % D[4], \
#		   'BNS at %.1f Mpc' % Dist ))
tspec.append('BNS at %.1f Mpc' % Dist)
plt.legend(tspec)
plt.savefig('asd_ranges.png')