#!/usr/bin/env python
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Print or plot info about instrument response function (IRF) files.
"""
from __future__ import print_function, division
import logging
logging.basicConfig(level=logging.INFO, format='%(levelname)s - %(message)s')
import os

import numpy as np
from astropy.units import Quantity
from astropy.io import fits


def retrieve_psf_info(hdu_list, energies, theta, fractions, plot=False):
    """
    Retrieve psf information from psf fits file.
    
    Parameters
    ----------
    hdu_list : `~astropy.io.fits.HDUList`
            HDU list with ``POINT SPREAD FUNCTION``extension.
    plot : bool
        Make a containment radius vs. theta and energy plot.
    """
    from gammapy.irf import EnergyDependentMultiGaussPSF
    psf = EnergyDependentMultiGaussPSF.from_fits(hdu_list)
    energies = Quantity(energies, 'TeV')
    thetas = Quantity(theta, 'deg')
    print(psf.info(fractions=fractions, energies=energies, thetas=thetas))
    
    if plot:
        for fraction in fractions:
            filename = 'containment_R{0:.0f}_energy_theta.png'.format(100 * fraction)
            psf.plot_containment(fraction, filename)
        
    
def retrieve_arf_info(hdu_list, energies, plot=False):
    """
    Retrieve effective area information from arf fits file.
    
    Parameters
    ----------
    hdu_list : `~astropy.io.fits.HDUList`
            HDU list with ``SPECRESP``extension.
    plot : bool
        Make an effective area vs. energy plot.
    """
    from gammapy.irf import EnergyDependentTableARF
    arf = EnergyDependentTableARF.from_fits(hdu_list)
    energies = Quantity(energies, 'TeV')
    print(arf.info(energies=energies))
    
    if plot:
        arf.plot_area_vs_energy('effective_area.png')
    

from gammapy.utils.scripts import argparse, GammapyFormatter
parser = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=GammapyFormatter)
parser.add_argument('infiles', type=str, nargs='+',
                    help='Input FITS file name')
parser.add_argument('--plot', action='store_true', 
                    help='Make info plot. Containment vs. theta and '
                    'energy for PSF or effective area vs. energy for ARF.')
parser.add_argument('--thetas', type=float, default=[0.], nargs='+',
                    help='Thetas where to evaluate PSF info.')
parser.add_argument('--energies', type=float, default=[1., 10.], nargs='+',
                    help='Energies where to evaluate PSF and ARF info.')
parser.add_argument('--fractions', type=float, default=[0.68, 0.95], nargs='+',
                    help='Containment fractions to compute for the PSF info.')
args = parser.parse_args()

for infile in args.infiles:
    hdu_list = fits.open(infile)
    hdu_names = [hdu.name for hdu in hdu_list]
    if 'POINT SPREAD FUNCTION' in hdu_names:
        logging.info('Auto detected PSF FITS file.')
        logging.info('Retrieving PSF info for {0}'.format(os.path.split(infile)[1]))
        retrieve_psf_info(hdu_list, args.energies, args.thetas, args.fractions, args.plot)
    elif 'SPECRESP' in hdu_names:
        logging.info('Auto detected ARF FITS file.')
        logging.info('Retrieving ARF info for {0}'.format(os.path.split(infile)[1]))
        retrieve_arf_info(hdu_list, args.energies, args.plot)
    else:
        logging.error('No valid FITS file found.')
