#! /usr/bin/env python
# ==========================================================================
# Display lightcurve generated by cslightcrv
#
# Copyright (C) 2017-2020 Juergen Knoedlseder
#
# 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, see <http://www.gnu.org/licenses/>.
#
# ==========================================================================
import sys
try:
    import matplotlib.pyplot as plt
    plt.figure()
    plt.close()
except (ImportError, RuntimeError):
    print('This script needs the "matplotlib" module')
    sys.exit()
import gammalib
import cscripts


# =============== #
# Plot lightcurve #
# =============== #
def plot_lightcurve(filename, plotfile):
    """
    Plot lightcurve

    Parameters
    ----------
    filename : str
        Name of lightcurve FITS file
    plotfile : str
        Plot file name
    """
    # Read spectrum file    
    fits    = gammalib.GFits(filename)
    table   = fits.table(1)
    
    # Extract standard columns
    c_mjd   = table['MJD']
    c_emjd  = table['e_MJD']
    c_ts    = table['TS']

    # Extract columns dependent on flux type
    if table.contains('EnergyFlux'):
        c_flux  = table['EnergyFlux']
        c_eflux = table['e_EnergyFlux']
        c_upper = table['EFluxUpperLimit']
        ylabel  = r'E $\times$ dN/dE (erg cm$^{-2}$ s$^{-1}$)'
    elif table.contains('PhotonFlux'):
        c_flux  = table['PhotonFlux']
        c_eflux = table['e_PhotonFlux']
        c_upper = table['FluxUpperLimit']
        ylabel  = r'N(E) (ph cm$^{-2}$ s$^{-1}$)'
    else:
        c_flux  = table['Prefactor']
        c_eflux = table['e_Prefactor']
        c_upper = table['DiffUpperLimit']
        ylabel  = r'dN/dE (cm$^{-2}$ s$^{-1}$ MeV$^{-1}$)'

    # Initialise arrays to be filled
    mjd       = []
    e_mjd     = []
    flux      = []
    e_flux    = []
    ul_mjd    = []
    ul_e_mjd  = []
    ul_flux   = []
    ul_e_flux = []

    # Loop over rows of the file
    nrows = table.nrows()
    for row in range(nrows):

        # Get Test Statistic, flux and flux error
        ts    = c_ts.real(row)
        flx   = c_flux.real(row)
        e_flx = c_eflux.real(row)

        # If Test Statistic is larger than 9 and twice the flux error is
        # smaller than the flux, then append flux point ...
        if ts > 9.0 and 2.0*e_flx < flx:
            mjd.append(c_mjd.real(row))
            e_mjd.append(c_emjd.real(row))
            flux.append(c_flux.real(row))
            e_flux.append(c_eflux.real(row))

        # ... otherwise append upper limit
        else:
            ul_mjd.append(c_mjd.real(row))
            ul_e_mjd.append(c_emjd.real(row))
            ul_flux.append(c_upper.real(row))
            ul_e_flux.append(0.5*c_upper.real(row))

    # Plot the spectrum 
    plt.figure()
    plt.semilogy()
    plt.grid()
    plt.errorbar(mjd, flux, yerr=e_flux, xerr=[e_mjd, e_mjd],
                 fmt='ro')
    plt.errorbar(ul_mjd, ul_flux, xerr=[ul_e_mjd, ul_e_mjd],
                 yerr=ul_e_flux, uplims=True, fmt='ro')
    plt.xlabel('MJD (days)')
    plt.ylabel(ylabel)

    # Show figure
    if len(plotfile) > 0:
        plt.savefig(plotfile)
    else:
        plt.show()

    # Return
    return


# =============== #
# Show lightcurve #
# =============== #
def show_lightcurve():
    """
    Show lightcurve
    """
    # Set usage string
    usage = 'show_lightcurve.py [-p plotfile] [file]'

    # Set default options
    options = [{'option': '-p', 'value': ''}]

    # Get arguments and options from command line arguments
    args, options = cscripts.ioutils.get_args_options(options, usage)

    # Extract script parameters from options
    plotfile = options[0]['value']

    # Plot lightcurve
    plot_lightcurve(args[0], plotfile)

    # Return
    return


# ======================== #
# Main routine entry point #
# ======================== #
if __name__ == '__main__':

    # Show lightcurve
    show_lightcurve()