"""
Fitting scalar and residual based objective functions
"""
import unittest
import numpy as np
import bornagain as ba


class Rosenbrock:

    def __init__(self):
        self.m_expected_minimum = 0.0
        self.m_expected_params = [1, 1.0]

    def objective_function(self, params):
        x = params["x"].value
        y = params["y"].value
        return (1 - x)**2 + 100 * (y - x**2)**2


def decaying_sin(params, x):
    amp = params['amp'].value
    phaseshift = params['phase'].value
    freq = params['frequency'].value
    decay = params['decay'].value
    return amp * np.sin(x * freq + phaseshift) * np.exp(-x * x * decay)


class DecayingSin:

    def __init__(self):
        self.m_x = np.linspace(0, 10, 100)
        self.m_params = ba.Parameters()
        self.m_params.add(ba.Parameter('amp', 10))
        self.m_params.add(ba.Parameter('decay', 0.05))
        self.m_params.add(ba.Parameter('phase', 1))
        self.m_params.add(ba.Parameter('frequency', 4))
        self.m_eps_data = np.linspace(0, 10, 100)
        self.m_eps_data.fill(0.01)
        self.m_data = decaying_sin(self.m_params, self.m_x)

    def objective_function(self, params):
        model = decaying_sin(params, self.m_x)
        return (self.m_data - model) / self.m_eps_data


class StandaloneFitTest(unittest.TestCase):

#    Rosenbrock fit test disabled while waiting for ways to set tolerance etc
#    (https://jugit.fz-juelich.de/mlz/bornagain/-/issues/693)
#
#    def test_RosenbrockFit(self):
#        print("\nStandalone fit test: Rosenbrock function")
#
#        params = ba.Parameters()
#        params.add(ba.Parameter("x", -1.2, ba.AttLimits.limited(-5, 5), 0.01))
#        params.add(ba.Parameter("y", 1, ba.AttLimits.limited(-5, 5), 0.01))
#
#        model = Rosenbrock()
#        minimizer = ba.Minimizer()
#        options = ba.MinimizerOptions()
#        options.addOption("Strategy", 2);
#        options.addOption("Tolerance", 1e-10);
#        minimizer.setOptions(options)
#
#        result = minimizer.minimize(model.objective_function, params)
#
#        print(result.toString())
#
#        # check found parameter values
#        np.testing.assert_almost_equal(result.parameters().values(),
#                                       model.m_expected_params, 2)
#
#        # check found minimum
#        np.testing.assert_almost_equal(result.minValue(),
#                                       model.m_expected_minimum, 2)
#
#        print("Done with Rosenbrock function\n")

    def test_DecayingSinFit(self):
        print("\nStandalone fit test: Decaying sin function")

        params = ba.Parameters()
        params.add(ba.Parameter('amp', 1, ba.AttLimits.positive()))
        params.add(ba.Parameter('decay', 0.1, ba.AttLimits.positive()))
        params.add(ba.Parameter('phase', 0.1, ba.AttLimits.limited(0, 3.1)))
        params.add(ba.Parameter('frequency', 1, ba.AttLimits.positive()))

        model = DecayingSin()
        minimizer = ba.Minimizer()
        result = minimizer.minimize(model.objective_function, params)

        print(result.toString())

        # check found parameter values
        np.testing.assert_almost_equal(result.parameters().values(),
                                       model.m_params.values(), 3)

        print("Done with decaying sin function\n")

if __name__ == '__main__':
    unittest.main()
