#!/usr/bin/env python
#
# Copyright (C) 2019 Gino Contestabile, Francesco Pannarale
#
# 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 3 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.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

"""Produce the sky grid plot for the triggered search (PyGRB)."""

# =============================================================================
# Preamble
# =============================================================================

import sys
import os
import logging
import numpy
import h5py
from matplotlib import pyplot as plt
from matplotlib import rc
import pycbc.version
from pycbc import init_logging
from pycbc.results import save_fig_with_metadata
from pycbc.results import pygrb_postprocessing_utils as ppu

plt.switch_backend('Agg')
rc('font', size=14)

__author__  = "Francesco Pannarale <francesco.pannarale@ligo.org>"
__version__ = pycbc.version.git_verbose_msg
__date__ = pycbc.version.date
__program__ = "pycbc_pygrb_plot_skygrid"

# =============================================================================
# Functions
# =============================================================================
# Load trigger data
def load_data(input_file, vetoes, injections=False):
    """Load data from a trigger/injection file"""
    
    logging.info("Loading triggers...")
    
    trigs = ppu.load_triggers(input_file, vetoes)
        
    return trigs

# =============================================================================
# Main script starts here
# =============================================================================
parser = ppu.pygrb_initialize_plot_parser(description=__doc__,
                                          version=__version__)
parser.add_argument("-t", "--trig-file", action="store",
                    default=None, required=True,
                    help="The location of the trigger file")
opts = parser.parse_args()

init_logging(opts.verbose, format="%(asctime)s:%(levelname)s : %(message)s")

trig_file = os.path.abspath(opts.trig_file)
outfile = opts.output_file
if opts.plot_title is None:
    opts.plot_title = 'PyGRB sky grid'

logging.info("Imported and ready to go.")

# Set output directories
outdirs = [os.path.split(os.path.abspath(outfile))[0]]
for outdir in outdirs:
    if not os.path.isdir(outdir):
        os.makedirs(outdir)

# Extract IFOs and vetoes
ifos, vetoes = ppu.extract_ifos_and_vetoes(trig_file, opts.veto_files,
                                           opts.veto_category)

# Load trigger data
trig_data = load_data(trig_file, vetoes)

#
# Generate sky grid plot
#

xlabel = "Longitude (Degrees)"
ylabel = "Latitude (Degrees)"

if opts.verbose:
    sys.stdout.write("\nPlotting...\n")
    fig_name = os.path.split(os.path.abspath(outfile))[1]
    sys.stdout.write(" * %s (%s vs %s)...\n" % (fig_name, xlabel, ylabel))

fig = plt.figure()
ax = fig.gca()
ax.set_xlabel(xlabel)#, fontsize=16)
ax.set_ylabel(ylabel)#, fontsize=16)
ax.plot(trig_data['network/longitude'][:], trig_data['network/latitude'][:],
        'ko', markerfacecolor='blue')
# Wrap up
save_fig_with_metadata(fig, outfile, cmd=' '.join(sys.argv),
                       title=opts.plot_title, caption=opts.plot_caption)
                       #fig_kwds=fig_kwds,
plt.close()
