#
#  Copyright (C) 2017  Smithsonian Astrophysical Observatory
#
#
#  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.
#

"""Test handling of ROSAT [1]_ data.

At the moment the tests are of routines in the sherpa.astro.io module,
since the high-level routins like sherpa.astro.ui.unpack_pha are a
very-thin wrapper around the io routines.

References
----------

.. [1] The ROSAT Mission https://heasarc.gsfc.nasa.gov/docs/rosat/rosat3.html

"""

import numpy as np
from numpy.testing import assert_allclose

import pytest
import six

from sherpa.utils.testing import requires_data, requires_fits
from sherpa.utils.err import IOErr
from sherpa.astro.data import DataPHA, DataRosatRMF

# Should each test import io instead of this? Also, do we have a
# better way of determining what the backend is?
#
try:
    from sherpa.astro import io
    if io.backend.__name__ == "sherpa.astro.io.pyfits_backend":
        backend = "pyfits"
    elif io.backend.__name__ == "sherpa.astro.io.crates_backend":
        backend = "crates"
    else:
        # Should not happen, but do not want to error out here.
        # Leave io as whatever was loaded, which will likely cause
        # the tests to fail.
        backend = None

except ImportError:
    io = None
    backend = None

from sherpa.astro import ui


@pytest.fixture(autouse=True)
def cleanup_astro_session(request):
    """Ensure sherpa.astro.ui is cleaned before and after the test."""
    ui.clean()

    def fin():
        ui.clean()

    request.addfinalizer(fin)


# The RMF includes the ARF (so is a "RSP" file).
#
PHAFILE = 'xrbg_xspec.pi'
RMFFILE = 'pspcc_gain1_256.rsp'

"""
Plotting the data - once the response has been read in - into
XSPEC 12.9.1p returns the following (this is the output of WDATA and
represents X value, X error, Y value, Y error).

First for X=channel, Y=normalized count/s/channel

    1 0.5 0 1
    2 0.5 2.80112115E-4 4.72556394E-6
    3 0.5 4.47507802E-4 5.81684435E-6
    4 0.5 4.62442949E-5 2.01607168E-6
    5 0.5 7.15584611E-5 2.57730176E-6
    6 0.5 8.68857605E-5 2.57510169E-6
    7 0.5 1.03978207E-4 2.74851345E-6
    8 0.5 5.3703181E-5 2.34389222E-6
    9 0.5 0 1
    ...
    63 0.5 0 1

Now for X=keV, Y=normalized count/s/keV

    4.96509969E-2 3.46949995E-2 0 14.4112997
    0.143822998 5.94770014E-2 2.35479348E-3 3.97259755E-5
    0.312340498 0.109040499 2.05202564E-3 2.66728621E-5
    0.470945001 4.95640039E-2 4.66510886E-4 2.03380623E-5
    0.609724522 8.92154872E-2 4.01042809E-4 1.44442511E-5
    0.803023994 0.104084015 4.17382835E-4 1.2370303E-5
    1.11032045 0.20321247 2.55836174E-4 6.76265927E-6
    1.66048098 0.346947968 7.73937136E-5 3.37787287E-6
    2.01238537 4.95660305E-3 0 100.875542
    ...
    2.54767656 4.95648384E-3 0 100.877968

"""


def validate_pha(pha, errors=False):
    """Has the PSPC PHA file been read in correctly?

    Parameters
    ----------
    pha : sherpa.astro.data.DataPHA
    errors : bool, optional
        Have the systematic errors been read in from the file?

    Notes
    ------
    XSPEC 12.9.1p reports this message when reading in this file:

    ***Warning: Unrecognized grouping for channel(s). It/they will be reset to 1.
    HOWEVER NOTE: Resets only occurred for channels of bad quality

    """

    assert isinstance(pha, DataPHA)

    # assert_allclose defaults to a relative tolerance of 1e-7.
    assert_allclose(pha.exposure, 1.0)
    assert_allclose(pha.backscal, 1.0)
    assert_allclose(pha.areascal, 1.0)

    # Although channel is an integer, it's treated as a
    # float.
    #
    nchan = 256
    assert_allclose(pha.channel, np.arange(1, nchan + 1))

    # The file has a RATE rather than COUNTS column; this gets
    # converted into a counts value, and the rate flag is set.
    # The exposure time here is 1.0, so the values are the
    # same (rate and counts) even if the (implied) units are not.
    #
    assert len(pha.counts) == nchan
    assert_allclose(pha.counts.min(), 0.0)
    assert_allclose(pha.counts.max(), 4.4750780216e-4)
    assert_allclose(pha.counts.sum(), 1.0899898225e-3)
    assert np.argmax(pha.counts) == 29

    idx = np.where(pha.counts > 0)[0]

    nonzerobins = [13, 29, 44, 59, 79, 110, 154]
    assert np.all(idx == nonzerobins)

    assert_allclose(pha.counts[nonzerobins],
                    [2.80112115e-4,
                     4.47507802e-4,
                     4.62442949e-5,
                     7.15584611e-5,
                     8.68857605e-5,
                     1.03978207e-4,
                     5.3703181e-5])

    # Simple statistics on the GROUPING and QUALITY channels
    #
    assert len(pha.quality) == nchan
    assert_allclose(pha.quality.min(), 0)
    assert_allclose(pha.quality.max(), 2)
    assert_allclose(pha.quality.sum(), 124)
    assert np.argmin(pha.quality) == 7
    assert np.argmax(pha.quality) == 0

    assert len(pha.grouping) == nchan
    assert_allclose(pha.grouping.min(), -1)
    assert_allclose(pha.grouping.max(), 1)
    assert_allclose(pha.grouping.sum(), -185)
    assert np.argmin(pha.grouping) == 1
    assert np.argmax(pha.grouping) == 0

    # TODO: check that grouping/quality is applied correctly

    emptyfields = ['syserror', 'bin_lo', 'bin_hi']
    if errors:
        assert len(pha.staterror) == nchan
        assert_allclose(pha.staterror.min(), 0.0)
        assert_allclose(pha.staterror.max(), 5.8168443502e-06)
        assert_allclose(pha.staterror.sum(), 2.2803289085e-05)
        assert np.argmin(pha.staterror) == 0
        assert np.argmax(pha.staterror) == 29

        idx = np.where(pha.counts > 0)[0]
        assert np.all(idx == nonzerobins)

        assert_allclose(pha.staterror[idx],
                        [4.72556394E-6,
                         5.81684435E-6,
                         2.01607168E-6,
                         2.57730176E-6,
                         2.57510169E-6,
                         2.74851345E-6,
                         2.34389222E-6])

    else:
        emptyfields.insert(0, 'staterror')

    for field in emptyfields:
        assert getattr(pha, field) is None

    assert pha.grouped is True
    assert pha.subtracted is False
    assert pha.units == 'channel'
    assert pha.rate is True
    assert pha.plot_fac == 0
    assert pha.response_ids == []
    assert pha.background_ids == []
    assert pha.get_background() is None


def validate_rmf(rmf):
    """Has the PSPC RMF file been read in correctly?

    Parameters
    ----------
    rmf : sherpa.astro.data.DataRMF

    Notes
    -----
    This is based on the values returned by Crates in Python 2.7,
    from CIAO 4.9, since it produces the same values as XSPEC 12.9.1
    when used to evaluate a power-law model.

    """

    assert isinstance(rmf, DataRosatRMF)

    # There are 10 bins which have N_GRP set to 0. These rows are filtered
    # out from the F_CHAN and N_CHAN arrays.
    #
    nchans = 256
    nbins_all = 729
    nbins_nonzero = nbins_all - 10

    assert rmf.detchans == nchans
    assert rmf.offset == 1

    for field in ['energ_lo', 'energ_hi', 'n_grp']:
        assert getattr(rmf, field).size == nbins_all

    for field in ['f_chan', 'n_chan']:
        assert getattr(rmf, field).size == nbins_nonzero

    nmatrix = 115269
    assert rmf.matrix.size == nmatrix

    for field in ['e_min', 'e_max']:
        assert getattr(rmf, field).size == nchans

    # Where are the N_GRP=0 bins?
    zeroidx = [0, 1, 2, 3, 4, 5, 6, 7, 8, 728]
    assert np.all(np.where(rmf.n_grp == 0)[0] == zeroidx)

    # All n_grp values are 0 or 1
    assert rmf.n_grp.sum() == nbins_nonzero
    assert (rmf.n_grp[rmf.n_grp > 0] == 1).all()

    assert rmf.f_chan.min() == 8
    assert rmf.f_chan.max() == 194
    assert rmf.f_chan.sum() == 7732
    assert np.argmax(rmf.f_chan) == 718

    assert rmf.n_chan.min() == 2
    assert rmf.n_chan.max() == 245
    assert rmf.n_chan.sum() == nmatrix
    assert np.argmax(rmf.n_chan) == 410

    # Rather than check each element, use some simple summary
    # statistics.
    assert_allclose(rmf.matrix.min(), 1.0e-5)
    assert_allclose(rmf.matrix.max(), 11.146280288696289)
    assert_allclose(rmf.matrix.sum(), 42994.357679666558)
    assert np.argmin(rmf.matrix) == 0
    assert np.argmax(rmf.matrix) == 7561

    # Expect the upper edge of bin j to equal the lower
    # edge of bin j + 1.
    #
    assert_allclose(rmf.energ_lo[1:], rmf.energ_hi[:-1])
    assert_allclose(rmf.e_min[1:], rmf.e_max[:-1])

    assert_allclose(rmf.energ_lo[0], 0.054607998579740524)
    assert_allclose(rmf.energ_hi[0], 0.071539998054504395)
    assert_allclose(rmf.energ_lo[-1], 3.0)
    assert_allclose(rmf.energ_hi[-1], 3.0099999904632568)

    assert_allclose(rmf.e_min[0], 0.014956000261008739)
    assert_allclose(rmf.e_max[0], 0.024869000539183617)
    assert_allclose(rmf.e_min[-1], 2.5427200794219971)
    assert_allclose(rmf.e_max[-1], 2.5526330471038818)


@requires_data
@requires_fits
def test_read_pha(make_data_path):
    """Can we read in a ROSAT PHA file."""

    infile = make_data_path(PHAFILE)
    pha = io.read_pha(infile)
    validate_pha(pha, errors=False)


@requires_data
@requires_fits
def test_read_pha_errors(make_data_path):
    """Can we read in a ROSAT PHA file."""

    infile = make_data_path(PHAFILE)
    pha = io.read_pha(infile, use_errors=True)
    validate_pha(pha, errors=True)


# In the error checks (i.e. checking that invalid files fail),
# the error messages and types depend on the backend. An attempt
# is made to be specific when testing these, primarily to point
# out when any refactoring changes this (since in many ways it
# would be nice to have the same error messages between the
# backends when possible).
#
@requires_data
@requires_fits
def test_read_pha_fails_rmf(make_data_path):
    """Just check in we can't read in a RMF as a PHA file."""

    if backend == 'pyfits':
        emsg = " does not appear to be a PHA spectrum"
        etype = IOErr
    elif backend == 'crates':
        emsg = "dmKeyRead() could not find key. 'HDUCLAS1'"
        etype = ValueError
    else:
        assert False, "Internal error: unknown backend {}".format(backend)

    infile = make_data_path(RMFFILE)
    with pytest.raises(etype) as excinfo:
        io.read_pha(infile)

    assert emsg in str(excinfo.value)


@requires_data
@requires_fits
def test_read_arf_fails_pha(make_data_path):
    """Just check in we can't read in a PHA file as an ARF."""

    if backend == 'pyfits':
        emsg = ' does not appear to be an ARF'
    elif backend == 'crates':
        emsg = " is not a filename or ARFCrate obj"
    else:
        assert False, "Internal error: unknown backend {}".format(backend)

    infile = make_data_path(PHAFILE)
    with pytest.raises(IOErr) as excinfo:
        io.read_arf(infile)

    assert emsg in str(excinfo.value)


@requires_data
@requires_fits
def test_read_arf_fails_rmf(make_data_path):
    """Just check in we can't read in a RNF as an ARF."""

    if backend == 'pyfits':
        emsg = ' does not appear to be an ARF'
    elif backend == 'crates':
        emsg = "Required column 'SPECRESP' not found in "
    else:
        assert False, "Internal error: unknown backend {}".format(backend)

    infile = make_data_path(RMFFILE)
    with pytest.raises(IOErr) as excinfo:
        io.read_arf(infile)

    assert emsg in str(excinfo.value)


@requires_data
@requires_fits
def test_read_rmf(make_data_path):

    if not six.PY2 and (backend == "crates"):
        pytest.skip('Python3 and Crates: known to fail')

    infile = make_data_path(RMFFILE)
    rmf = io.read_rmf(infile)
    validate_rmf(rmf)


@requires_data
@requires_fits
def test_read_rmf_crates_py3(make_data_path):
    """At present the RMF can not be read with the crates backend with
    Python 3.5. This test will let us know when Crates has been fixed.
    """

    if six.PY2 or (backend != "crates"):
        pytest.skip('Test only needed for Crates with Python 3')

    infile = make_data_path(RMFFILE)
    with pytest.raises(ValueError) as excinfo:
        io.read_rmf(infile)

    assert str(excinfo.value) == "dmKeyRead() cound not find key. " + \
        "'HDUCLAS1'"


@requires_data
@requires_fits
def test_read_rmf_fails_pha(make_data_path):
    """Just check in we can't read in a PHA file as a RMF."""

    if backend == 'pyfits':
        emsg = ' does not appear to be an RMF'
    elif backend == 'crates':
        emsg = "Required column 'ENERG_LO' not found in "
    else:
        assert False, "Internal error: unknown backend {}".format(backend)

    infile = make_data_path(PHAFILE)
    with pytest.raises(IOErr) as excinfo:
        io.read_rmf(infile)

    assert emsg in str(excinfo.value)


@requires_data
@requires_fits
def test_can_use_pspc_data(make_data_path):
    """A basic check that we can read in and use the ROSAT PSPC data.

    Unlike the previous tests, that directly access the io module,
    this uses the ui interface.
    """

    if not six.PY2 and (backend == "crates"):
        pytest.skip('Python3 and Crates: known to fail')

    # The PSPC PHA file does not have the ANCRFILE/RESPFILE keywords
    # set up, so the responses has to be manually added.
    #
    ui.load_pha(make_data_path(PHAFILE), use_errors=True)
    assert ui.get_analysis() == 'channel'

    ui.load_rmf(make_data_path(RMFFILE))
    assert ui.get_analysis() == 'energy'

    ui.set_source(ui.powlaw1d.pl)
    ui.set_par('pl.gamma', 1.7)
    ui.set_par('pl.ampl', 2e-6)

    s = ui.get_stat_info()[0]
    assert s.numpoints == 63
    assert s.dof == 61

    # Value obtained from XSPEC 12.9.1p; Sherpa returns
    # sexpected = 973.2270845920297
    sexpected = 973.23
    assert_allclose(s.statval, sexpected, rtol=0, atol=0.005)

    # apply an energy filter to remove the "bogus" points
    ui.ignore(None, 0.05)

    s = ui.get_stat_info()[0]
    assert s.numpoints == 62
    assert s.dof == 60
    assert_allclose(s.statval, sexpected, rtol=0, atol=0.005)

    ui.ignore(2.01, None)

    s = ui.get_stat_info()[0]
    assert s.numpoints == 7
    assert s.dof == 5

    assert_allclose(s.statval, sexpected, rtol=0, atol=0.005)
