# --------------------------------------------------------------------------------------------
# This script gives an example of reading the frequency-dependent, complex-valued calibration
# error and uncertainty for a given detector at a given GPS time and plotting the results in
# magnitude and phase.
#
# The O3 calibration error and uncertainty estimates for LIGO Hanford, LIGO Livingston, and Virgo
# at all times are packed in the following files:
# (Download: https://dcc.ligo.org/T2100313/public)
#     LIGO (Hanford and Livingston): LIGO_O3_cal_uncertainty.tgz
#     Virgo: Virgo_O3_cal_uncertainty.tgz
#
# Similar files can be found for O2 (LIGO and Virgo) and O1 (LIGO) within the same link.
#
# Please see README for detailed description of the data.
# --------------------------------------------------------------------------------------------

import numpy as np
from glob import glob

# --------------------------------------------------------------------------------------------
# Set parameters and read data
#
# Detector: 'H1' for LIGO Hanford
#           'L1' for LIGO Livingston
#           'V1' for Virgo
#
# GPS time: At the two LIGO detectors, the calibration error and uncertainty are estimated
#           at discrete times on an hourly cadence. Set the GPS time we are interested in
#           (e.g., a time when a gravitational-wave event happened), and find the data
#           corresponding to a discrete time closest to the interested time.
#           At Virgo, the single set of calibration error and uncertainty are applicable to
#           the whole observing run.
# --------------------------------------------------------------------------------------------

# Set detector to LIGO Hanford
IFO = 'H1'
# Set GPS time to a fake event time (randomly chosen)
event_time = 1256652999

# Create a base dir, unzip the tarballs in the base dir
# In this example, the base dir is set to O3CalUncertainty under the current dir
base_dir = './O3CalUncertainty/'

# Collect all hourly discrete GPS times for this detector
all_times = []
for f in glob(base_dir + IFO + '/*FinalResults.txt'):
    name_str = f.split('_')
    gps = np.int(name_str[4]) 
    all_times.append(gps)

# Find the closest GPS time available in the datasets to the event time we are interested in
data_time = all_times[np.argmin(np.abs(np.array(all_times) - event_time))]

# Load data for this detector at this closest GPS time
data_file = glob(base_dir + IFO + '/*' + str(data_time) + '*FinalResults.txt')[0]
data = np.loadtxt(data_file)

# Read the corresponding columns:
# Column 0 - Frequency
# Column 1 - Median systematic error in magnitude
# Column 2 - Median systematic error in phase
# Column 3 - Lower (-1 sigma) uncertainty boundary in magnitude
# Coulmn 4 - Lower (-1 sigma) uncertainty boundary in phase
# Column 5 - Upper (+1 sigma) uncertainty boundary in magnitude
# Coulmn 6 - Upper (+1 sigma) uncertainty boundary in phase

ff = data[:,0]
median_mag = data[:,1]
median_phase = data[:,2]
min_1sigma_mag = data[:,3]
min_1sigma_phase = data[:,4]
max_1sigma_mag = data[:,5]
max_1sigma_phase = data[:,6]

# --------------------------------------------------------------------------------------------
# Plot the frequency-dependent, complex-valued calibration error and uncertainty
# --------------------------------------------------------------------------------------------

import matplotlib as mpl
import matplotlib.pyplot as plt
mpl.rcParams.update({'font.family': 'DejaVu Sans',
                     'font.size': 22})

# Set color convention for different detectors
if IFO == 'H1':
    ifo_color = 'r'
elif IFO == 'L1':
    ifo_color = 'b'
else:
    ifo_color = 'm'

# Set frequency, magnitude error and phase error display ranges
freqRange = [10.0, 5000.0]
yMagRange   = [-8.0,8.0]
yPhaseRange = [-6.0,6.0]

# Create plot with two subplots for magnitude and phase errors respectively
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(15, 15))
ax0, ax1 = axes.flat

# Plot median magnitude error as a solid curve
ax0.plot(ff, 100 * (median_mag - 1.0), color=ifo_color, label =r'Median')
# Plot +/- 1-sigma uncertainty boundaries as dashed curves
ax0.plot(ff, 100 * (min_1sigma_mag - 1.0), linestyle='--', color=ifo_color, label =r'$1 \sigma$ Uncertainty')
ax0.plot(ff, 100 * (max_1sigma_mag - 1.0), linestyle='--',color=ifo_color)
# Fill the whole 1-sigma uncertainty range with shaded color
ax0.fill_between(ff, 100 * (min_1sigma_mag - 1.0), 100 * (max_1sigma_mag - 1.0),color=ifo_color,alpha=0.3)
# Set axes ranges, ticks, and grids
ax0.set_xlim(freqRange)
ax0.set_ylim(yMagRange)
ax0.minorticks_on()
ax0.set_yticks(np.arange(yMagRange[0],yMagRange[1]+1,2))
ax0.grid(ls='-', which='major',lw=1.1)
ax0.grid(ls='--', which='minor')

# Plot median phase error as a solid curve
ax1.plot(ff, 180/np.pi*median_phase, color=ifo_color,label =r'Median')
# Plot +/- 1-sigma uncertainty boundaries as dashed curves
ax1.plot(ff, 180/np.pi*min_1sigma_phase, linestyle='--',color=ifo_color,label =r'$1 \sigma$ Uncertainty')
ax1.plot(ff, 180/np.pi*max_1sigma_phase, linestyle='--',color=ifo_color)
# Fill the whole 1-sigma uncertainty range with shaded color
ax1.fill_between(ff, 180/np.pi*min_1sigma_phase, 180/np.pi*max_1sigma_phase,color=ifo_color,alpha=0.3)
# Set axes ranges, ticks, and grids
ax1.set_xlim(freqRange)
ax1.set_ylim(yPhaseRange)
ax1.minorticks_on()
ax1.set_yticks(np.arange(yPhaseRange[0],yPhaseRange[1]+1,2))
ax1.grid(ls='-', which='major',lw=1.1)
ax1.grid(ls='--', which='minor')

for ax in axes.flat:
    # Set frequency ticks
    ax.set_xscale("log")
    ax.set_xticks(np.array([10, 20, 50, 100, 200, 500, 1000, 2000, 5000]))
    ax.set_xticklabels(['10', '20', '50', '100', '200', '500', '1000', '2000', '5000'])
    
    # Set legend
    legend = ax.legend(loc='lower center', ncol=2, handlelength=3)
    for i in legend.legendHandles:
        i.set_linewidth(2)

# Set axes labels
ax1.set_xlabel(r'Frequency [Hz]')
ax0.set_ylabel(r'Magnitude error [%]')
ax1.set_ylabel(r'Phase error [deg]')

# Save figure
figfilename='./%s-%s.png' % (IFO,data_time)
plt.savefig(figfilename, bbox_inches='tight', pad_inches=0.2)