#
#  Copyright (C) 2012, 2015-2016, 2018-2023, 2025
#  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.
#

from collections import namedtuple
import logging

import numpy as np

import pytest

from sherpa.utils.testing import requires_data
from sherpa.data import Data1D
from sherpa.fit import Fit
from sherpa.instrument import ConvolutionKernel, PSFModel
from sherpa.models import ArithmeticModel, Parameter
from sherpa.models.basic import PowLaw1D
from sherpa.models.parameter import hugeval
from sherpa.stats import LeastSq
from sherpa.utils.err import ArgumentErr, IdentifierErr, PSFErr, StatErr, SessionErr
from sherpa.utils.logging import SherpaVerbosity
from sherpa import ui

# As the user model is called UserModel, refer to the Sherpa version
# by its full module name
import sherpa.models.basic


@pytest.fixture(autouse=True)
def hide_logging():
    "hide INFO-level logging in all these tests"

    with SherpaVerbosity("ERROR"):
        yield


class UserModel(ArithmeticModel):
    def __init__(self, name='usermodel'):
        self.param1 = Parameter(name, 'param1', 1, min=0, max=100)
        self.param2 = Parameter(name, 'param2', 1, min=-100, max=100)

        ArithmeticModel.__init__(self, name, (self.param1,
                                              self.param2))

    def calc(self, p, x, *args, **kwargs):
        return p[0] * x + p[1]


# A function-based version of this
#
def um_line(pars, xlo, *args, **kwargs):
    """A test user model: straight line with slope and intercept"""
    return pars[0] * np.asarray(xlo) + pars[1]


"""
Tests for PR #155

TODO: Some test cases cannot be implemented (easily) without mock
In particular one cannot test that the correct matrix is used when one is not
provided and that an error is thrown when, for any reason, the computed covariance matrix
is None.
"""


WRONG_STAT_MSG = "Fit statistic must be cash, cstat or wstat, not {}"
WSTAT_ERR_MSG = "No background data has been supplied. Use cstat"
NO_COVAR_MSG = "covariance has not been performed"
RIGHT_STATS = {'cash', 'cstat', 'cstatnegativepenalty', 'wstat'}


@pytest.fixture
def setup_covar(make_data_path):

    ui.load_data(make_data_path('sim.poisson.1.dat'))
    ui.set_model(PowLaw1D("p1"))


@requires_data
@pytest.mark.parametrize('stat', sorted(set(ui.list_stats())
                                        - RIGHT_STATS
                                        # As of 4.18.0 chi2/... now error out
                                        # with unique errors, so skip testing
                                        # them
                                        - {"chi2", "leastsq", "userstat"}
                                        ))
def test_covar_wrong_stat(stat, clean_ui, setup_covar):
    "Test an exception is thrown is the proper stat is not set"

    # ui.covar()  this is not needed as of 4.18.0
    ui.set_stat(stat)
    with pytest.raises(ValueError,
                       match=WRONG_STAT_MSG.format(stat)):
        ui.get_draws()


@requires_data
def test_covar_wstat_no_background(clean_ui, setup_covar):
    "Test an exception is thrown when wstat is used without background"

    ui.covar()
    ui.set_stat("wstat")
    with pytest.raises(StatErr,
                       match=WSTAT_ERR_MSG):
        ui.get_draws()


SEED_VALUE = 12347


def check_draws(stats, accept, params):
    """Check the test_covar_as_none/argument output

    Ideally we'd be able to check that test_covar_as_none and
    test_covar_as_argument return the same answers, but

      a) numerical differences, since we only give the covar matrix
         to 8 dp
      b) the fits, and hence parameter choices, depend on the chosen
         statistic, leading to subtly different values

    So, all we check is

      - the shapes are as expected
      - the accept flag is as expected, as this happens to be the
        same (but doesn't have to be).

    """

    assert stats.size == 11
    assert accept.size == 11
    assert params.shape == (2, 11)

    # this depends on SEED_VALUE
    assert accept == pytest.approx([0, 1, 0, 1, 1, 1, 1, 0, 0, 0, 1])


@requires_data
@pytest.mark.parametrize('statname', sorted(RIGHT_STATS - {'wstat'}))
def test_covar_as_argument(statname, clean_ui, setup_covar):
    """covariance matrix is sent to get_draws"""

    ui.set_stat(statname)
    ui.fit()

    ui.set_rng(np.random.RandomState(SEED_VALUE))

    # This should be the same as
    #
    #     ui.covar()
    #     matrix = ui.get_covar_results().extra()
    #
    # but we don't check this as we don't want to have run covar to
    # match expected use of get_draws.
    #
    matrix = [[0.00064075, 0.01122127], [0.01122127, 0.20153251]]
    stats, accept, params = ui.get_draws(niter=10, covar_matrix=matrix)
    check_draws(stats, accept, params)


@requires_data
@pytest.mark.parametrize('statname', sorted(RIGHT_STATS - {'wstat'}))
def test_covar_as_none(statname, clean_ui, setup_covar):
    """covariance matrix is not sent to get_draws"""

    ui.set_stat(statname)
    ui.fit()
    # ui.covar() as of 4.18.0 this is no-longer needed

    ui.set_rng(np.random.RandomState(SEED_VALUE))

    stats, accept, params = ui.get_draws(niter=10)
    check_draws(stats, accept, params)


def identity(x):
    return x


@pytest.fixture
def setup_ui_full(make_data_path, setup_ui):

    out = namedtuple('setup_ui_data', ['ascii', 'single', 'double',
                                       'filter'])

    out.ascii = make_data_path('sim.poisson.1.dat')
    out.single = make_data_path('single.dat')
    out.double = make_data_path('double.dat')
    out.filter = make_data_path('filter_single_integer.dat')
    return out


@pytest.fixture
def setup_ui():
    """A version of setup_ui_full where we don't need the data."""

    ui.dataspace1d(1, 1000, dstype=ui.Data1D)


@pytest.fixture
def setup_ui_2d():
    """Basic 2D data set"""

    ui.dataspace2d((3, 4))


@requires_data
def test_ui_ascii(clean_ui, setup_ui_full):
    ui.load_data(1, setup_ui_full.ascii)
    ui.load_data(1, setup_ui_full.ascii, 2)
    ui.load_data(1, setup_ui_full.ascii, 2, ("col2", "col1"))


# Test table model
@requires_data
def test_ui_table_model_ascii_table(clean_ui, setup_ui_full):
    ui.load_table_model('tbl', setup_ui_full.single)
    ui.load_table_model('tbl', setup_ui_full.double)


# Test user model
@requires_data
def test_ui_user_model_ascii_table(clean_ui, setup_ui_full):
    ui.load_user_model(identity, 'mdl', setup_ui_full.single)
    ui.load_user_model(identity, 'mdl', setup_ui_full.double)


@requires_data
def test_ui_filter_ascii(clean_ui, setup_ui_full):
    ui.load_filter(setup_ui_full.filter)
    ui.load_filter(setup_ui_full.filter, ignore=True)


def test_ui_add_model(clean_ui, setup_ui):
    ui.add_model(UserModel)
    ui.set_model('usermodel.user1')


def test_ui_set_full_model(clean_ui, setup_ui):
    ui.load_psf('psf1', 'gauss1d.g1')
    ui.set_full_model('psf1(gauss1d.g2) + const1d.c1')
    ui.get_model()


def test_ui_set_full_model_2d(clean_ui, setup_ui_2d):
    ui.load_psf('psf1', 'gauss2d.g1')
    ui.set_full_model('psf1(gauss2d.g2 ) +const2d.c1')
    ui.get_model()


def test_ui_set_full_model_mismatch_2d(clean_ui, setup_ui):
    ui.load_psf('psf1', 'gauss2d.g1')
    ui.set_full_model('psf1(gauss2d.g2)+const2d.c1')
    with pytest.raises(PSFErr,
                       match="kernel 'gauss2d.g1' and data 'dataspace1d' do not match: 2D vs 1D"):
        ui.get_model()


def test_ui_set_full_model_2d_mismatch_1d(clean_ui, setup_ui_2d):
    ui.load_psf('psf1', 'gauss1d.g1')
    ui.set_full_model('psf1(gauss1d.g2 ) +const1d.c1')
    with pytest.raises(PSFErr,
                       match="kernel 'gauss1d.g1' and data 'dataspace2d' do not match: 1D vs 2D"):
        ui.get_model()


def test_ui_set_full_model_checks_dimensions_match(clean_ui, setup_ui_2d):
    ui.load_psf('psf1', 'gauss2d.g1')
    with pytest.raises(ArgumentErr,
                       match=r"invalid model expression: Models do not match: 2D \(psfmodel.psf1\) and 1D \(gauss1d.g2\)"):
        ui.set_full_model('psf1(gauss1d.g2)+const2d.c1')


# Bug 12644
@requires_data
def test_ui_source_methods_with_full_model(clean_ui, setup_ui_full, caplog):

    ui.load_data('full', setup_ui_full.ascii)
    ui.set_full_model('full', 'powlaw1d.p1')

    # Test Case 1
    with pytest.raises(IdentifierErr,
                       match="Convolved model\n'powlaw1d.p1'\n is set for dataset full. You should use get_model instead."):
        ui.get_source('full')

    with pytest.raises(IdentifierErr,
                       match="Convolved model\n'powlaw1d.p1'\n is set for dataset full. You should use plot_model instead."):
        ui.plot_source('full')

    with pytest.raises(IdentifierErr,
                       match="Convolved model\n'powlaw1d.p1'\n is set for dataset full. You should use get_model_plot instead."):
        ui.get_source_plot('full')

    # Test Case 2.
    #
    # The hide_logging fixture hides the logging output, so we need
    # to re-enable it for this call.
    #
    assert len (caplog.records) == 0
    with SherpaVerbosity("INFO"):
        ui.set_source('full', 'powlaw1d.p2')

    assert len (caplog.records) == 1
    logname, loglvl, logmsg = caplog.record_tuples[0]
    assert logname == "sherpa.ui.utils"
    assert loglvl == logging.WARNING
    assert logmsg == "Clearing convolved model\n'powlaw1d.p1'\nfor dataset full"

    ui.get_source('full')

    # Test Case 3
    ui.load_data('not_full', setup_ui_full.ascii)
    with pytest.raises(IdentifierErr,
                       match=r"source not_full has not been set, consider using set_source\(\) or set_model\(\)"):
        ui.get_source('not_full')


@pytest.mark.parametrize('model', ['gauss1d', 'delta1d', 'normgauss1d'])
def test_psf_model1d(model, clean_ui):
    ui.dataspace1d(1, 10)
    ui.load_psf('psf1d', model + '.mdl')
    ui.set_psf('psf1d')
    mdl = ui.get_model_component('mdl')
    assert mdl.get_center() == (4.0, )


@pytest.mark.parametrize('model', ['gauss2d', 'delta2d', 'normgauss2d'])
def test_psf_model2d(model):
    ui.dataspace2d([216, 261])
    ui.load_psf('psf2d', model + '.mdl')
    ui.set_psf('psf2d')
    mdl = ui.get_model_component('mdl')
    assert mdl.get_center() == (108.0, 130.0)


def test_does_user_model_get_cleaned(clean_ui):
    """Do user models get removed from the session by clean?"""

    mname = "test_model"
    with pytest.raises(IdentifierErr):
        ui.get_model_component(mname)

    ui.load_user_model(um_line, mname)
    mdl = ui.get_model_component(mname)
    assert mdl.name == "usermodel.{}".format(mname)
    assert isinstance(mdl, sherpa.models.basic.UserModel)

    ui.clean()

    with pytest.raises(IdentifierErr):
        ui.get_model_component(mname)


# Test simple use of load_user_model and add_user_pars
#
def test_user_model_create_pars_default(clean_ui):

    mname = "test_model"
    ui.load_user_model(um_line, mname)

    mdl = ui.get_model_component(mname)
    assert len(mdl.pars) == 1
    par = mdl.pars[0]
    assert par.name == "ampl"
    assert par.val == pytest.approx(1.0)
    assert not par.frozen
    assert par.units == ''
    assert par.min == pytest.approx(-1 * hugeval)
    assert par.max == pytest.approx(hugeval)


def test_user_model_create_pars_names(clean_ui):

    mname = "test_model"
    ui.load_user_model(um_line, mname)

    mdl = ui.get_model_component(mname)
    assert len(mdl.pars) == 1

    # add user pars doesn't change the existing instance, you have
    # to "get" the new version to see the change
    #
    ui.add_user_pars(mname, ['X1', 'x'])

    mdl = ui.get_model_component(mname)
    assert len(mdl.pars) == 2
    p0 = mdl.pars[0]
    p1 = mdl.pars[1]

    assert p0.name == 'X1'
    assert p0.val == pytest.approx(0.0)
    assert p0.units == ''
    assert not p0.frozen
    assert p0.min == pytest.approx(-1 * hugeval)
    assert p0.max == pytest.approx(hugeval)

    assert p1.name == 'x'
    assert p1.val == pytest.approx(0.0)
    assert p1.units == ''
    assert not p1.frozen
    assert p1.min == pytest.approx(-1 * hugeval)
    assert p1.max == pytest.approx(hugeval)


def test_user_model_create_pars_full(clean_ui):

    mname = "test_model"
    ui.load_user_model(um_line, mname)
    ui.add_user_pars(mname, ['pAr1', '_p'], [23.2, 3.1e2],
                     parunits=['', 'cm^2 s'],
                     parfrozen=[True, False],
                     parmins=[0, -100],
                     parmaxs=[100, 1e5])

    mdl = ui.get_model_component(mname)
    assert len(mdl.pars) == 2
    p0 = mdl.pars[0]
    p1 = mdl.pars[1]

    assert p0.name == 'pAr1'
    assert p0.val == pytest.approx(23.2)
    assert p0.units == ''
    assert p0.frozen
    assert p0.min == pytest.approx(0)
    assert p0.max == pytest.approx(100)

    assert p1.name == '_p'
    assert p1.val == pytest.approx(3.1e2)
    assert p1.units == 'cm^2 s'
    assert not p1.frozen
    assert p1.min == pytest.approx(-100)
    assert p1.max == pytest.approx(1e5)


# diagnose and test out GitHub issue #609; the inability to change
# parameter values for a user model using direct access
#
def test_user_model_change_par(clean_ui):

    mname = "test_model"
    ui.load_user_model(um_line, mname)
    ui.add_user_pars(mname, ['xXx', 'Y2'])

    mdl = ui.get_model_component(mname)
    assert len(mdl.pars) == 2
    p0 = mdl.pars[0]
    p1 = mdl.pars[1]

    assert p0.name == 'xXx'
    assert p1.name == 'Y2'
    assert p0.val == pytest.approx(0.0)
    assert p1.val == pytest.approx(0.0)

    # Use the user-supplied names:
    #
    mdl.xXx = 2.0
    assert p0.val == pytest.approx(2.0)

    mdl.Y2 = 3.0
    assert p1.val == pytest.approx(3.0)

    # Now all lower case
    #
    mdl.xxx = 4.0
    assert p0.val == pytest.approx(4.0)

    mdl.y2 = 12.0
    assert p1.val == pytest.approx(12.0)

    # Try with the set_par function
    #
    ui.set_par('test_model.xxx', 12.2)
    assert p0.val == pytest.approx(12.2)

    ui.set_par('test_model.y2', 14.0, frozen=True)
    assert p1.val == pytest.approx(14.0)
    assert p1.frozen

    ui.clean()


def test_user_model1d_eval_fail(clean_ui):
    """This is expected to fail as the number of pars does not match."""

    mname = "test_model"
    ui.load_user_model(um_line, mname)
    mdl = ui.get_model_component(mname)
    with pytest.raises(IndexError):
        mdl([2.3, 5.4, 8.7])


def test_user_model1d_eval(clean_ui):
    """Simple evaluation check for 1D case."""

    mname = "test_model"
    ui.load_user_model(um_line, mname)
    ui.add_user_pars(mname, ["slope", "intercept"])

    m = 2.1
    c = -4.8

    mdl = ui.get_model_component(mname)
    mdl.slope = m
    mdl.intercept = c

    x = np.asarray([2.3, 5.4, 8.7])
    y = mdl(x)

    yexp = x * m + c

    # This check require pytest >= 3.2.0
    #
    assert y == pytest.approx(yexp)


def test_user_model1d_fit(clean_ui):
    """Check can use in a fit."""

    mname = "test_model"
    ui.load_user_model(um_line, mname)
    ui.add_user_pars(mname, ["slope", "intercept"],
                     parvals=[1.0, 1.0])

    mdl = ui.get_model_component(mname)

    x = np.asarray([-2.4, 2.3, 5.4, 8.7, 12.3])

    # Set up the data to be scattered around y = -0.2 x + 2.8
    # Pick the deltas so that they sum to 0 (except for central
    # point)
    #
    slope = -0.2
    intercept = 2.8

    dy = np.asarray([0.1, -0.2, 0.14, -0.1, 0.2])
    ydata = x * slope + intercept + dy

    ui.load_arrays(1, x, ydata)

    ui.set_source(mname)
    ui.ignore(5.0, 6.0)  # drop the central bin

    ui.set_stat('leastsq')
    ui.set_method('simplex')
    ui.fit()

    fres = ui.get_fit_results()
    assert fres.succeeded
    assert fres.parnames == ('test_model.slope', 'test_model.intercept')
    assert fres.numpoints == 4
    assert fres.dof == 2

    # Tolerance has been adjusted to get the tests to pass on my
    # machine. It's really just to check that the values have changed
    # from their default values.
    #
    assert fres.parvals[0] == pytest.approx(slope, abs=0.01)
    assert fres.parvals[1] == pytest.approx(intercept, abs=0.05)

    # These should be the same values, so no need to use pytest.approx
    # (unless there's some internal translation between types done
    # somewhere?).
    #
    assert mdl.slope.val == fres.parvals[0]
    assert mdl.intercept.val == fres.parvals[1]


class MyCacheTestModel(ArithmeticModel):

    def calc(self, par, x):
        A = par[0]
        mylambda = par[1]
        b = par[2]
        fvec = A * np.exp(- mylambda * x) + b
        if self.counter == 0:
            assert self.cache > 0
        else:
            assert self.cache == 0
        self.counter += 1
        return fvec

    def __init__(self, name='myexp'):
        self.A = Parameter(name, 'A', 1)
        self.mylambda = Parameter(name, 'mylambda', 2)
        self.b = Parameter(name, 'b', 3)
        self.counter = 0
        ArithmeticModel.__init__(self, name, (self.A, self.mylambda, self.b))


def test_cache():
    """To make sure that the runtime fit(cache=???) works"""

    x = np.array([1.0, 2.0, 3.0])
    model = MyCacheTestModel()
    par = np.array([1.1, 2.0, 3.0])
    y = model.calc(par, x)

    data = Data1D('tmp', x, y)
    fit = Fit(data, model, LeastSq())
    fit.fit(cache=False)


def test_list_psf_ids_empty(clean_ui):
    assert ui.list_psf_ids() == []


@pytest.mark.parametrize("id", [None, 2, "2"])
def test_list_psf_ids_single(id, clean_ui):
    ui.dataspace1d(1, 10, id=id)
    ui.load_psf('psf1d', 'gauss1d.mdl')
    ui.set_psf(id, 'psf1d')

    if id is None:
        id = 1

    assert ui.list_psf_ids() == [id]


def test_list_psf_ids_multi(clean_ui):

    ui.dataspace1d(1, 10)
    ui.dataspace1d(1, 10, id="2")
    ui.load_psf('psf1d', 'gauss1d.mdl')
    ui.set_psf('psf1d')
    ui.set_psf("2", 'psf1d')

    ans = ui.list_psf_ids()
    assert ans == [1, "2"]


def test_load_conv_from_file(tmp_path, clean_ui):
    """Check we can read in a convolution model from a file.

    This does not require a valid I/O module.
    """

    tmp = tmp_path / "tmp.dat"
    tmp.write_text("! X Y\n1 10\n2 5\n 3 0\n4 1\n")

    assert ui.list_model_components() == []

    ui.load_conv("foo", str(tmp), comment="!")
    assert ui.list_model_components() == ["foo"]

    foo = ui.get_model_component("foo")
    assert isinstance(foo, ConvolutionKernel)
    assert isinstance(foo.kernel, Data1D)
    assert foo.kernel.x == pytest.approx([1, 2, 3, 4])
    assert foo.kernel.y == pytest.approx([10, 5, 0, 1])


def test_load_psf_from_file(tmp_path, clean_ui):
    """Check we can read in a PSF model from a file.

    This does not require a valid I/O module.
    """

    tmp = tmp_path / "tmp.dat"
    tmp.write_text("! X Y Z\n1 10 0.3\n2 5 0.2\n 3 0 0.1\n4 1 0.1\n")

    assert ui.list_model_components() == []

    ui.load_psf("foo", str(tmp), comment="!", colkeys=["X", "Z"])
    assert ui.list_model_components() == ["foo"]

    foo = ui.get_model_component("foo")
    assert isinstance(foo, PSFModel)
    assert isinstance(foo.kernel, Data1D)
    assert foo.kernel.x == pytest.approx([1, 2, 3, 4])
    assert foo.kernel.y == pytest.approx([0.3, 0.2, 0.1, 0.1])
