# /*##########################################################################
# Copyright (C) 2016-2021 European Synchrotron Radiation Facility
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# ############################################################################*/
"""
Nominal tests of the leastsq function.
"""

import unittest

import numpy
import sys

from silx.utils import testutils
from silx.math.fit.leastsq import _logger as fitlogger


class Test_leastsq(unittest.TestCase):
    """
    Unit tests of the leastsq function.
    """

    ndims = None

    def setUp(self):
        try:
            from silx.math.fit import leastsq

            self.instance = leastsq
        except ImportError:
            self.instance = None

        def myexp(x):
            # put a (bad) filter to avoid over/underflows
            # with no python looping
            with numpy.errstate(invalid="ignore"):
                return numpy.exp(
                    x * numpy.less(abs(x), 250)
                ) - 1.0 * numpy.greater_equal(abs(x), 250)

        self.my_exp = myexp

        def gauss(x, *params):
            params = numpy.asarray(params, dtype=numpy.float64)
            result = params[0] + params[1] * x
            for i in range(2, len(params), 3):
                p = params[i : (i + 3)]
                dummy = 2.3548200450309493 * (x - p[1]) / p[2]
                result += p[0] * self.my_exp(-0.5 * dummy * dummy)
            return result

        self.gauss = gauss

        def gauss_derivative(x, params, idx):
            if idx == 0:
                return numpy.ones(len(x), numpy.float64)
            if idx == 1:
                return x
            gaussian_peak = (idx - 2) // 3
            gaussian_parameter = (idx - 2) % 3
            actual_idx = 2 + 3 * gaussian_peak
            p = params[actual_idx : (actual_idx + 3)]
            if gaussian_parameter == 0:
                return self.gauss(x, *[0, 0, 1.0, p[1], p[2]])
            if gaussian_parameter == 1:
                tmp = self.gauss(x, *[0, 0, p[0], p[1], p[2]])
                tmp *= 2.3548200450309493 * (x - p[1]) / p[2]
                return tmp * 2.3548200450309493 / p[2]
            if gaussian_parameter == 2:
                tmp = self.gauss(x, *[0, 0, p[0], p[1], p[2]])
                tmp *= 2.3548200450309493 * (x - p[1]) / p[2]
                return tmp * 2.3548200450309493 * (x - p[1]) / (p[2] * p[2])

        self.gauss_derivative = gauss_derivative

    def tearDown(self):
        self.instance = None
        self.gauss = None
        self.gauss_derivative = None
        self.my_exp = None
        self.model_function = None
        self.model_derivative = None

    def testImport(self):
        self.assertTrue(
            self.instance is not None, "Cannot import leastsq from silx.math.fit"
        )

    def testUnconstrainedFitNoWeight(self):
        parameters_actual = [10.5, 2, 1000.0, 20.0, 15]
        x = numpy.arange(10000.0)
        y = self.gauss(x, *parameters_actual)
        parameters_estimate = [0.0, 1.0, 900.0, 25.0, 10]
        model_function = self.gauss

        fittedpar, cov = self.instance(model_function, x, y, parameters_estimate)
        test_condition = numpy.allclose(parameters_actual, fittedpar)
        if not test_condition:
            msg = "Unsuccessfull fit\n"
            for i in range(len(fittedpar)):
                msg += "Expected %g obtained %g\n" % (
                    parameters_actual[i],
                    fittedpar[i],
                )
            self.assertTrue(test_condition, msg)

    def testUnconstrainedFitWeight(self):
        parameters_actual = [10.5, 2, 1000.0, 20.0, 15]
        x = numpy.arange(10000.0)
        y = self.gauss(x, *parameters_actual)
        sigma = numpy.sqrt(y)
        parameters_estimate = [0.0, 1.0, 900.0, 25.0, 10]
        model_function = self.gauss

        fittedpar, cov = self.instance(
            model_function, x, y, parameters_estimate, sigma=sigma
        )
        test_condition = numpy.allclose(parameters_actual, fittedpar)
        if not test_condition:
            msg = "Unsuccessfull fit\n"
            for i in range(len(fittedpar)):
                msg += "Expected %g obtained %g\n" % (
                    parameters_actual[i],
                    fittedpar[i],
                )
            self.assertTrue(test_condition, msg)

    def testDerivativeFunction(self):
        parameters_actual = [10.5, 2, 10000.0, 20.0, 150, 5000, 900.0, 300]
        x = numpy.arange(10000.0)
        y = self.gauss(x, *parameters_actual)
        delta = numpy.sqrt(numpy.finfo(numpy.float64).eps)
        for i in range(len(parameters_actual)):
            p = parameters_actual * 1
            if p[i] == 0:
                delta_par = delta
            else:
                delta_par = p[i] * delta
            if i > 2:
                p[0] = 0.0
                p[1] = 0.0
            p[i] += delta_par
            yPlus = self.gauss(x, *p)
            p[i] = parameters_actual[i] - delta_par
            yMinus = self.gauss(x, *p)
            numerical_derivative = (yPlus - yMinus) / (2 * delta_par)
            # numerical_derivative = (self.gauss(x, *p) - y) / delta_par
            p[i] = parameters_actual[i]
            derivative = self.gauss_derivative(x, p, i)
            diff = numerical_derivative - derivative
            test_condition = numpy.allclose(
                numerical_derivative, derivative, atol=5.0e-6
            )
            if not test_condition:
                msg = "Error calculating derivative of parameter %d." % i
                msg += "\n diff min = %g diff max = %g" % (diff.min(), diff.max())
                self.assertTrue(test_condition, msg)

    def testConstrainedFit(self):
        CFREE = 0
        CPOSITIVE = 1
        CQUOTED = 2
        CFIXED = 3
        CFACTOR = 4
        CDELTA = 5
        CSUM = 6
        parameters_actual = [10.5, 2, 10000.0, 20.0, 150, 5000, 900.0, 300]
        x = numpy.arange(10000.0)
        y = self.gauss(x, *parameters_actual)
        parameters_estimate = [0.0, 1.0, 900.0, 25.0, 10, 400, 850, 200]
        model_function = self.gauss
        model_deriv = self.gauss_derivative
        constraints_all_free = [[0, 0, 0]] * len(parameters_actual)
        constraints_all_positive = [[1, 0, 0]] * len(parameters_actual)
        constraints_delta_position = [[0, 0, 0]] * len(parameters_actual)
        constraints_delta_position[6] = [CDELTA, 3, 880]
        constraints_sum_position = constraints_all_positive * 1
        constraints_sum_position[6] = [CSUM, 3, 920]
        constraints_factor = constraints_delta_position * 1
        constraints_factor[2] = [CFACTOR, 5, 2]
        constraints_list = [
            None,
            constraints_all_free,
            constraints_all_positive,
            constraints_delta_position,
            constraints_sum_position,
        ]

        # for better code coverage, the warning recommending to set full_output
        # to True when using constraints should be shown at least once
        full_output = True
        for index, constraints in enumerate(constraints_list):
            if index == 2:
                full_output = None
            elif index == 3:
                full_output = 0
            for model_deriv in [None, self.gauss_derivative]:
                for sigma in [None, numpy.sqrt(y)]:
                    fittedpar, cov = self.instance(
                        model_function,
                        x,
                        y,
                        parameters_estimate,
                        sigma=sigma,
                        constraints=constraints,
                        model_deriv=model_deriv,
                        full_output=full_output,
                    )[:2]
                    full_output = True

        test_condition = numpy.allclose(parameters_actual, fittedpar)
        if not test_condition:
            msg = "Unsuccessfull fit\n"
            for i in range(len(fittedpar)):
                msg += "Expected %g obtained %g\n" % (
                    parameters_actual[i],
                    fittedpar[i],
                )
            self.assertTrue(test_condition, msg)

    def testUnconstrainedFitAnalyticalDerivative(self):
        parameters_actual = [10.5, 2, 1000.0, 20.0, 15]
        x = numpy.arange(10000.0)
        y = self.gauss(x, *parameters_actual)
        sigma = numpy.sqrt(y)
        parameters_estimate = [0.0, 1.0, 900.0, 25.0, 10]
        model_function = self.gauss
        model_deriv = self.gauss_derivative

        fittedpar, cov = self.instance(
            model_function,
            x,
            y,
            parameters_estimate,
            sigma=sigma,
            model_deriv=model_deriv,
        )
        test_condition = numpy.allclose(parameters_actual, fittedpar)
        if not test_condition:
            msg = "Unsuccessfull fit\n"
            for i in range(len(fittedpar)):
                msg += "Expected %g obtained %g\n" % (
                    parameters_actual[i],
                    fittedpar[i],
                )
            self.assertTrue(test_condition, msg)

    @testutils.validate_logging(fitlogger.name, warning=2)
    def testBadlyShapedData(self):
        parameters_actual = [10.5, 2, 1000.0, 20.0, 15]
        x = numpy.arange(10000.0).reshape(1000, 10)
        y = self.gauss(x, *parameters_actual)
        sigma = numpy.sqrt(y)
        parameters_estimate = [0.0, 1.0, 900.0, 25.0, 10]
        model_function = self.gauss

        for check_finite in [True, False]:
            fittedpar, cov = self.instance(
                model_function,
                x,
                y,
                parameters_estimate,
                sigma=sigma,
                check_finite=check_finite,
            )
            test_condition = numpy.allclose(parameters_actual, fittedpar)
            if not test_condition:
                msg = "Unsuccessfull fit\n"
                for i in range(len(fittedpar)):
                    msg += "Expected %g obtained %g\n" % (
                        parameters_actual[i],
                        fittedpar[i],
                    )
                self.assertTrue(test_condition, msg)

    @testutils.validate_logging(fitlogger.name, warning=3)
    def testDataWithNaN(self):
        parameters_actual = [10.5, 2, 1000.0, 20.0, 15]
        x = numpy.arange(10000.0).reshape(1000, 10)
        y = self.gauss(x, *parameters_actual)
        sigma = numpy.sqrt(y)
        parameters_estimate = [0.0, 1.0, 900.0, 25.0, 10]
        model_function = self.gauss
        x[500] = numpy.inf
        # check default behavior
        try:
            self.instance(model_function, x, y, parameters_estimate, sigma=sigma)
        except ValueError:
            info = "%s" % sys.exc_info()[1]
            self.assertTrue("array must not contain inf" in info)

        # check requested behavior
        try:
            self.instance(
                model_function,
                x,
                y,
                parameters_estimate,
                sigma=sigma,
                check_finite=True,
            )
        except ValueError:
            info = "%s" % sys.exc_info()[1]
            self.assertTrue("array must not contain inf" in info)

        fittedpar, cov = self.instance(
            model_function, x, y, parameters_estimate, sigma=sigma, check_finite=False
        )
        test_condition = numpy.allclose(parameters_actual, fittedpar)
        if not test_condition:
            msg = "Unsuccessfull fit\n"
            for i in range(len(fittedpar)):
                msg += "Expected %g obtained %g\n" % (
                    parameters_actual[i],
                    fittedpar[i],
                )
            self.assertTrue(test_condition, msg)

        # testing now with ydata containing NaN
        x = numpy.arange(10000.0).reshape(1000, 10)
        y[500] = numpy.nan
        fittedpar, cov = self.instance(
            model_function, x, y, parameters_estimate, sigma=sigma, check_finite=False
        )

        test_condition = numpy.allclose(parameters_actual, fittedpar)
        if not test_condition:
            msg = "Unsuccessfull fit\n"
            for i in range(len(fittedpar)):
                msg += "Expected %g obtained %g\n" % (
                    parameters_actual[i],
                    fittedpar[i],
                )
            self.assertTrue(test_condition, msg)

        # testing now with sigma containing NaN
        sigma[300] = numpy.nan
        fittedpar, cov = self.instance(
            model_function, x, y, parameters_estimate, sigma=sigma, check_finite=False
        )
        test_condition = numpy.allclose(parameters_actual, fittedpar)
        if not test_condition:
            msg = "Unsuccessfull fit\n"
            for i in range(len(fittedpar)):
                msg += "Expected %g obtained %g\n" % (
                    parameters_actual[i],
                    fittedpar[i],
                )
            self.assertTrue(test_condition, msg)

    def testUncertainties(self):
        """Test for validity of uncertainties in returned full-output
        dictionary. This is a non-regression test for pull request #197"""
        parameters_actual = [10.5, 2, 1000.0, 20.0, 15, 2001.0, 30.1, 16]
        x = numpy.arange(10000.0)
        y = self.gauss(x, *parameters_actual)
        parameters_estimate = [0.0, 1.0, 900.0, 25.0, 10.0, 1500.0, 20.0, 2.0]

        # test that uncertainties are not 0.
        fittedpar, cov, infodict = self.instance(
            self.gauss, x, y, parameters_estimate, full_output=True
        )
        uncertainties = infodict["uncertainties"]
        self.assertEqual(len(uncertainties), len(parameters_actual))
        self.assertEqual(len(uncertainties), len(fittedpar))
        for uncertainty in uncertainties:
            self.assertNotAlmostEqual(uncertainty, 0.0)

        # set constraint FIXED for half the parameters.
        # This should cause leastsq to return 100% uncertainty.
        parameters_estimate = [10.6, 2.1, 1000.1, 20.1, 15.1, 2001.1, 30.2, 16.1]
        CFIXED = 3
        CFREE = 0
        constraints = []
        for i in range(len(parameters_estimate)):
            if i % 2:
                constraints.append([CFIXED, 0, 0])
            else:
                constraints.append([CFREE, 0, 0])
        fittedpar, cov, infodict = self.instance(
            self.gauss,
            x,
            y,
            parameters_estimate,
            constraints=constraints,
            full_output=True,
        )
        uncertainties = infodict["uncertainties"]
        for i in range(len(parameters_estimate)):
            if i % 2:
                # test that all FIXED parameters have 100% uncertainty
                self.assertAlmostEqual(uncertainties[i], parameters_estimate[i])
