import abc
import unittest
import numpy as np
import six

from kafe2.core.minimizers import AVAILABLE_MINIMIZERS

from kafe2.config import kc

from kafe2.fit import IndexedFit
from kafe2.fit.indexed.cost import IndexedCostFunction_Chi2

from kafe2.test.fit.test_fit import AbstractTestFit


def simple_chi2(data, model):
    return np.sum((data - model)**2)

def simple_chi2_explicit_model_name(data, simple_indexed_model):
    return np.sum((data - simple_indexed_model)**2)

def simple_indexed_model(a=1.1, b=2.2, c=3.3):
    x = np.arange(10)
    return a * x ** 2 + b * x + c

def line_indexed_model(a=3.0, b=0.0):
    x = np.arange(10)
    return a * x + b


def analytic_solution(des_mat, cov_mat_inv, data):
    return (
        np.squeeze(np.asarray(np.linalg.inv(
            des_mat.T.dot(cov_mat_inv).dot(des_mat)
        ).dot(des_mat.T)
         .dot(cov_mat_inv)
         .dot(data))))


class TestIndexedFitBasicInterface(AbstractTestFit, unittest.TestCase):

    MINIMIZER = 'scipy'

    def setUp(self):
        self._n_points = 10

        # "jitter" for data smearing
        self._jitter = np.array([-0.3193475 , -1.2404198 , -1.4906926 , -0.78832446,
                                   -1.7638106,   0.36664261,  0.49433821,  0.0719646,
                                    1.95670326,  0.31200215])
        assert len(self._jitter) == self._n_points

        # reference initial values
        self._ref_initial_pars = np.array([1.1, 2.2, 3.3])
        self._ref_initial_model = simple_indexed_model(*self._ref_initial_pars)

        # fit data
        self._ref_data = self._ref_initial_model + self._jitter

        self._linear_design_matrix = np.array([
            np.arange(10) ** 2,
            np.arange(10),
            np.ones(10)
        ]).T

        # pre-fit cost value
        self._ref_initial_cost = simple_chi2(
            data=self._ref_data,
            model=self._ref_initial_model,
        )

        # reference matrices/errors
        self._ref_error = np.ones(self._n_points)
        self._ref_matrix_eye = np.eye(self._n_points)

        # reference fit result values
        #self._nominal_fit_result_pars = (1.1351433, 2.13736919, 2.33346549)
        self._nominal_fit_result_pars = analytic_solution(
            self._linear_design_matrix,
            np.linalg.inv(self._ref_matrix_eye),
            self._ref_data,
        )

        self._nominal_fit_result_model = simple_indexed_model(*self._nominal_fit_result_pars)
        self._nominal_fit_result_cost = simple_chi2(
            data=self._ref_data,
            model=self._nominal_fit_result_model,
        )

        # helper dict with all reference property values
        self._ref_prop_dict = dict(
            did_fit=False,
            model_count=1,

            parameter_values=self._ref_initial_pars,
            parameter_names=('a', 'b', 'c'),
            cost_function_value=self._ref_initial_cost,

            data=self._ref_data,
            data_error=self._ref_error,
            data_cov_mat=self._ref_matrix_eye,
            data_cov_mat_inverse=self._ref_matrix_eye,
            data_cor_mat=self._ref_matrix_eye,

            model=self._ref_initial_model,
            model_error=self._ref_error * 0,
            model_cov_mat=self._ref_matrix_eye * 0,
            model_cov_mat_inverse=None,
            model_cor_mat=self._ref_matrix_eye * np.nan,

            total_error=self._ref_error,
            total_cov_mat=self._ref_matrix_eye,
            total_cov_mat_inverse=self._ref_matrix_eye,
            total_cor_mat=self._ref_matrix_eye,
        )

    def _get_fit(
            self, model_function=None, cost_function=None, errors=None,
            dynamic_error_algorithm=None):
        '''convenience'''
        model_function = model_function or simple_indexed_model
        # TODO: fix default
        cost_function = cost_function or IndexedCostFunction_Chi2(
            errors_to_use='covariance')
        errors = errors or [dict(err_val=1.0)]
        dynamic_error_algorithm = dynamic_error_algorithm or "nonlinear"

        _fit = IndexedFit(
            data=self._ref_data,
            model_function=model_function,
            cost_function=cost_function,
            minimizer=self.MINIMIZER,
            dynamic_error_algorithm=dynamic_error_algorithm
        )
        for _err in errors:
            _fit.add_error(**_err)

        return _fit

    def _get_test_fits(self):
        return {
            'default': self._get_fit(),
            'explicit': self._get_fit(cost_function=simple_chi2),
            'explicit_model': self._get_fit(cost_function=simple_chi2_explicit_model_name),
            'relative_errors_data': self._get_fit(
                errors=[dict(err_val=1.0 / self._ref_data, relative=True, reference='data')]
            ),
        }

    def test_initial_state(self):
        self.run_test_for_all_fits(
            self._ref_prop_dict
        )

    def test_fit_results(self):
        self.run_test_for_all_fits(
            dict(self._ref_prop_dict,
                parameter_values=self._nominal_fit_result_pars,
                model=self._nominal_fit_result_model,
                did_fit=True,
                cost_function_value=self._nominal_fit_result_cost,
            ),
            call_before_fit=lambda f: f.do_fit(),
            rtol=1e-2
        )

    def test_set_all_parameter_values(self):
        self.run_test_for_all_fits(
            dict(self._ref_prop_dict,
                parameter_values=self._nominal_fit_result_pars,
                data=self._ref_data,
                model=self._nominal_fit_result_model,
                parameter_names=('a', 'b', 'c'),
                did_fit=False,
                model_count=1,
                cost_function_value=self._nominal_fit_result_cost
            ),
            # set parameters to their final values by hand
            call_before_fit=lambda f: f.set_all_parameter_values(
                self._nominal_fit_result_pars),
        )

    def test_set_all_parameter_values_wrong_number_raise(self):
        with self.assertRaises(ValueError):
            self._get_fit().set_all_parameter_values((1,))
        with self.assertRaises(ValueError):
            self._get_fit().set_all_parameter_values((1,2,3,4,5))

    def test_parameter_defaults(self):
        def dummy_model(a, b, c):
            return np.arange(10)

        _fit = self._get_fit(model_function=dummy_model)

        self._assert_fit_properties(
            _fit,
            dict(
                parameter_values=np.array([1, 1, 1]),
            )
        )

    def test_parameter_partial_defaults(self):
        def dummy_model(a, b, c=3.3):
            return np.arange(10)

        _fit = self._get_fit(model_function=dummy_model)

        self._assert_fit_properties(
            _fit,
            dict(
                parameter_values=np.array([1, 1, 3.3]),
            )
        )

    def test_update_data(self):
        _fit = self._get_fit()

        _fit.data = self._ref_data * 2
        _fit.add_error(err_val=1.0)

        self._assert_fit_properties(
            _fit,
            dict(
                data=self._ref_data * 2,
            )
        )

        _fit.do_fit()
        _new_estimates = np.array(self._nominal_fit_result_pars) * 2

        self._assert_fit_properties(
            _fit,
            dict(
                data=self._ref_data * 2,
                parameter_values=_new_estimates,
            ),
            rtol=1e-2
        )

    def test_update_data_length_mismatch_raise(self):
        # TODO: update when different length limitation is fixed
        _fit = self._get_fit()
        with self.assertRaises(ValueError):
            _fit.data = np.arange(self._n_points + 1)

    def test_reserved_parameter_names_raise(self):
        def dummy_model(data):
            pass

        with self.assertRaises(ValueError) as _exc:
            IndexedFit(data=self._ref_data,
                  model_function=dummy_model,
                  minimizer=self.MINIMIZER)

        self.assertIn('reserved', _exc.exception.args[0])
        self.assertIn('data', _exc.exception.args[0])

    def dummy_model_no_pars_raise(self):
        def dummy_model():
            pass

        with self.assertRaises(ValueError) as _exc:
            IndexedFit(data=self._ref_data,
                  model_function=dummy_model,
                  minimizer=self.MINIMIZER)

        self.assertIn(
            'needs at least one parameter',
            _exc.exception.args[0])

    def dummy_model_varargs_raise(self):
        # TODO: raise even without 'par'
        def dummy_model(x, par, *varargs):
            pass

        with self.assertRaises(ValueError) as _exc:
            IndexedFit(data=self._ref_data,
                  model_function=dummy_model,
                  minimizer=self.MINIMIZER)

        self.assertIn('variable', _exc.exception.args[0])
        self.assertIn('varargs', _exc.exception.args[0])

    def dummy_model_varkwargs_raise(self):
        # TODO: raise even without 'par'
        def dummy_model(x, par, **varkwargs):
            pass

        with self.assertRaises(ValueError) as _exc:
            IndexedFit(data=self._ref_data,
                  model_function=dummy_model,
                  minimizer=self.MINIMIZER)

        self.assertIn('variable', _exc.exception.args[0])
        self.assertIn('varkwargs', _exc.exception.args[0])

    def dummy_model_varargs_varkwargs_raise(self):
        # TODO: raise even without 'par'
        def dummy_model(x, par, *varargs, **varkwargs):
            pass

        with self.assertRaises(ValueError) as _exc:
            IndexedFit(data=self._ref_data,
                  model_function=dummy_model,
                  minimizer=self.MINIMIZER)

        self.assertIn('variable', _exc.exception.args[0])
        self.assertIn('varargs', _exc.exception.args[0])
        # TODO: enable when implemented
        #self.assertIn('varkwargs', _exc.exception.args[0])

    def test_report_before_fit(self):
        # TODO: check report content
        _buffer = six.StringIO()
        _fit = self._get_fit()
        _fit.report(output_stream=_buffer)
        self.assertNotEqual(_buffer.getvalue(), "")

    def test_report_after_fit(self):
        # TODO: check report content
        _buffer = six.StringIO()
        _fit = self._get_fit()
        _fit.do_fit()
        _fit.report(output_stream=_buffer)
        self.assertNotEqual(_buffer.getvalue(), "")

    def test_relative_model_error_nonlinear(self):
        _fit_data_err = self._get_fit(errors=[
            dict(err_val=1.0, relative=False, reference="data"),
            dict(err_val=0.1, relative=True, reference="data")
        ])
        _fit_model_err = self._get_fit(errors=[
            dict(err_val=1.0, relative=False, reference="data"),
            dict(err_val=0.1, relative=True, reference="model")
        ], dynamic_error_algorithm="nonlinear")
        _fit_data_err.do_fit()
        _fit_model_err.do_fit()
        self._assert_fit_results_equal(_fit_data_err, _fit_model_err, rtol=1e-2, atol=2e-2)

    def test_relative_model_error_iterative(self):
        _fit_data_err = self._get_fit(errors=[
            dict(err_val=1.0, relative=False, reference="data"),
            dict(err_val=0.1, relative=True, reference="data")
        ])
        _fit_model_err = self._get_fit(errors=[
            dict(err_val=1.0, relative=False, reference="data"),
            dict(err_val=0.1, relative=True, reference="model")
        ], dynamic_error_algorithm="iterative")
        _fit_data_err.do_fit()
        _fit_model_err.do_fit()
        self._assert_fit_results_equal(_fit_data_err, _fit_model_err, rtol=1e-2)

    def test_fit_scalar_model_function(self):
        def _model_function(a):
            return a

        _fit = IndexedFit([1, 2], _model_function)
        _fit.add_error(1.0)
        _fit.do_fit()
        self.assertAlmostEqual(_fit.parameter_values[0], 1.5)
        _fit.report()


class TestIndexedFitWithSimpleErrors(AbstractTestFit, unittest.TestCase):

    MINIMIZER = 'scipy'

    def setUp(self):
        six.get_unbound_function(TestIndexedFitBasicInterface.setUp)(self)

    def _get_fit(self, errors=None):
        '''convenience'''

        errors = errors or [dict(err_val=1.0)]

        _fit = IndexedFit(
            data=self._ref_data,
            model_function=simple_indexed_model,
            cost_function=IndexedCostFunction_Chi2(errors_to_use='covariance'),
            minimizer=self.MINIMIZER
        )

        for _err in errors:
            _fit.add_error(**_err)

        return _fit

    def _get_test_fits(self):
        return {
            'default': self._get_fit(),
            'two_errors': self._get_fit(
                errors=[
                    dict(err_val=1.0/np.sqrt(2)),
                    dict(err_val=1.0/np.sqrt(2))
                ]
            ),
            'named_errors': self._get_fit(
                errors=[
                    dict(err_val=1.0/np.sqrt(2), name="MyDataError",
                         correlation=0, relative=False, reference='data'),
                    dict(err_val=1.0/np.sqrt(2), name="MyModelError",
                         correlation=0, relative=False, reference='model'),
                ]
            )
        }

    def test_initial_state(self):
        self.run_test_for_all_fits(
            self._ref_prop_dict,
            fit_names=['default', 'two_errors'],
        )

    def test_fit_results(self):
        self.run_test_for_all_fits(
            dict(self._ref_prop_dict,
                parameter_values=self._nominal_fit_result_pars,
                model=self._nominal_fit_result_model,
                did_fit=True,
                cost_function_value=self._nominal_fit_result_cost,
            ),
            fit_names=['default', 'two_errors'],
            call_before_fit=lambda f: f.do_fit(),
            rtol=1e-2
        )

    def test_get_matching_error_all(self):
        _fit = self._get_test_fits()['named_errors']
        for _mc in (None, dict()):
            _errs = _fit.get_matching_errors(matching_criteria=_mc)
            self.assertEqual(len(_errs), 2)
            self.assertIs(_fit.data_container._error_dicts['MyDataError']['err'], _errs['MyDataError'])
            self.assertIs(_fit._param_model._error_dicts['MyModelError']['err'], _errs['MyModelError'])

    def test_get_matching_error_name(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(name='MyDataError'))
        self.assertEqual(len(_errs), 1)
        self.assertIs(_fit.data_container._error_dicts['MyDataError']['err'], _errs['MyDataError'])

    def test_get_matching_error_type_simple(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(type='simple'))
        self.assertEqual(len(_errs), 2)
        self.assertIs(_fit.data_container._error_dicts['MyDataError']['err'], _errs['MyDataError'])
        self.assertIs(_fit._param_model._error_dicts['MyModelError']['err'], _errs['MyModelError'])

    def test_get_matching_error_type_matrix(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(type='matrix'))
        self.assertEqual(len(_errs), 0)

    def test_get_matching_error_uncorrelated(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(correlated=True))
        self.assertEqual(len(_errs), 0)

    def test_get_matching_error_correlated(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(correlated=False))
        self.assertEqual(len(_errs), 2)
        self.assertIs(_fit.data_container._error_dicts['MyDataError']['err'], _errs['MyDataError'])
        self.assertIs(_fit._param_model._error_dicts['MyModelError']['err'], _errs['MyModelError'])

    def test_get_matching_error_reference_data(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(reference='data'))
        self.assertEqual(len(_errs), 1)
        self.assertIs(_fit.data_container._error_dicts['MyDataError']['err'], _errs['MyDataError'])

    def test_get_matching_error_reference_model(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(reference='model'))
        self.assertEqual(len(_errs), 1)
        self.assertIs(_fit._param_model._error_dicts['MyModelError']['err'], _errs['MyModelError'])


class TestIndexedFitWithMatrixErrors(AbstractTestFit, unittest.TestCase):

    MINIMIZER = 'scipy'

    def setUp(self):
        six.get_unbound_function(TestIndexedFitBasicInterface.setUp)(self)

    def _get_fit(self, errors=None):
        '''convenience'''

        _fit = IndexedFit(
            data=self._ref_data,
            model_function=simple_indexed_model,
            cost_function=IndexedCostFunction_Chi2(
                errors_to_use='covariance'),
            minimizer=self.MINIMIZER
        )

        if errors is None:
            _fit.add_matrix_error(err_matrix=np.eye(self._n_points), matrix_type='cov')
        else:
            for _err in errors:
                if 'err_matrix' in _err:
                    _fit.add_matrix_error(**_err)
                else:
                    _fit.add_error(**_err)

        return _fit

    def _get_test_fits(self):
        return {
            'default': self._get_fit(),
            'cor_matrix_and_error_vector': self._get_fit(
                errors=[
                    dict(err_matrix=np.eye(self._n_points),
                         matrix_type='cor',
                         err_val=1.0
                    )
                ]
            ),
            'two_matrix_errors': self._get_fit(
                errors=[
                    dict(err_matrix=np.eye(self._n_points)/2, matrix_type='cov'),
                    dict(err_matrix=np.eye(self._n_points)/2, matrix_type='cov')
                ]
            ),
            'one_matrix_one_simple_error': self._get_fit(
                errors=[
                    dict(err_matrix=np.eye(self._n_points)/2, matrix_type='cov'),
                    dict(err_val=1/np.sqrt(2))
                ]
            ),
            'named_errors': self._get_fit(
                errors=[
                    dict(err_val=1.0/np.sqrt(2), name="MySimpleDataError",
                         correlation=0, relative=False, reference='data'),
                    dict(err_val=1.0/np.sqrt(2), name="MySimpleModelError",
                         correlation=0, relative=False, reference='model'),
                    dict(err_matrix=np.eye(self._n_points)/2, matrix_type='cov',
                         name="MyMatrixDataError",
                         relative=False, reference='data'),
                    dict(err_matrix=np.eye(self._n_points)/2, matrix_type='cov',
                         name="MyMatrixModelError",
                         relative=False, reference='model'),
                ]
            )
        }

    def test_initial_state(self):
        self.run_test_for_all_fits(
            self._ref_prop_dict,
            fit_names=['default', 'cor_matrix_and_error_vector',
                       'one_matrix_one_simple_error', 'two_matrix_errors'],
        )

    def test_fit_results(self):
        self.run_test_for_all_fits(
            dict(self._ref_prop_dict,
                parameter_values=self._nominal_fit_result_pars,
                model=self._nominal_fit_result_model,
                did_fit=True,
                cost_function_value=self._nominal_fit_result_cost,
            ),
            fit_names=['default', 'cor_matrix_and_error_vector',
                       'one_matrix_one_simple_error', 'two_matrix_errors'],
            call_before_fit=lambda f: f.do_fit(),
            rtol=1e-2
        )

    def test_get_matching_error_all(self):
        _fit = self._get_test_fits()['named_errors']
        for _mc in (None, dict()):
            _errs = _fit.get_matching_errors(matching_criteria=_mc)
            self.assertEqual(len(_errs), 4)
            self.assertIs(_fit.data_container._error_dicts['MySimpleDataError']['err'], _errs['MySimpleDataError'])
            self.assertIs(_fit._param_model._error_dicts['MySimpleModelError']['err'], _errs['MySimpleModelError'])
            self.assertIs(_fit.data_container._error_dicts['MyMatrixDataError']['err'], _errs['MyMatrixDataError'])
            self.assertIs(_fit._param_model._error_dicts['MyMatrixModelError']['err'], _errs['MyMatrixModelError'])

    def test_get_matching_error_name(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(name='MySimpleDataError'))
        self.assertEqual(len(_errs), 1)
        self.assertIs(_fit.data_container._error_dicts['MySimpleDataError']['err'], _errs['MySimpleDataError'])

    def test_get_matching_error_type_simple(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(type='simple'))
        self.assertEqual(len(_errs), 2)
        self.assertIs(_fit.data_container._error_dicts['MySimpleDataError']['err'], _errs['MySimpleDataError'])
        self.assertIs(_fit._param_model._error_dicts['MySimpleModelError']['err'], _errs['MySimpleModelError'])

    def test_get_matching_error_type_matrix(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(type='matrix'))
        self.assertEqual(len(_errs), 2)
        self.assertIs(_fit.data_container._error_dicts['MyMatrixDataError']['err'], _errs['MyMatrixDataError'])
        self.assertIs(_fit._param_model._error_dicts['MyMatrixModelError']['err'], _errs['MyMatrixModelError'])

    def test_get_matching_error_uncorrelated(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(correlated=True))
        self.assertEqual(len(_errs), 0)

    def test_get_matching_error_correlated(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(correlated=False))
        # NOTE: passing 'correlated' only matches 'matrix' errors, irrespective of 'True'/'False' value passed
        self.assertEqual(len(_errs), 2)
        self.assertIs(_fit.data_container._error_dicts['MySimpleDataError']['err'], _errs['MySimpleDataError'])
        self.assertIs(_fit._param_model._error_dicts['MySimpleModelError']['err'], _errs['MySimpleModelError'])

    def test_get_matching_error_reference_data(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(reference='data'))
        self.assertEqual(len(_errs), 2)
        self.assertIs(_fit.data_container._error_dicts['MySimpleDataError']['err'], _errs['MySimpleDataError'])
        self.assertIs(_fit.data_container._error_dicts['MyMatrixDataError']['err'], _errs['MyMatrixDataError'])

    def test_get_matching_error_reference_model(self):
        _fit = self._get_test_fits()['named_errors']
        _errs = _fit.get_matching_errors(matching_criteria=dict(reference='model'))
        self.assertEqual(len(_errs), 2)
        self.assertIs(_fit._param_model._error_dicts['MySimpleModelError']['err'], _errs['MySimpleModelError'])
        self.assertIs(_fit._param_model._error_dicts['MyMatrixModelError']['err'], _errs['MyMatrixModelError'])
