import unittest

import numpy as np
from six import StringIO

from kafe2.fit._base import ModelFunctionBase
from kafe2.fit.histogram import HistModelFunction
from kafe2.fit.histogram.model import HistParametricModel
from kafe2.fit.indexed import IndexedModelFunction
from kafe2.fit.indexed.model import IndexedParametricModel
from kafe2.fit.io.handle import IOStreamHandle
from kafe2.fit.representation import (
    ParametricModelYamlReader,
    ParametricModelYamlWriter,
)
from kafe2.fit.xy.model import XYParametricModel

TEST_PARAMETRIC_MODEL_HIST = """
type: histogram
n_bins: 5
bin_range: [0, 5]
model_density_function:
    type: histogram
    python_code: |
        def linear_model(x, a, b):
            return a * x + b
model_parameters: [0.0, 0.08]
"""

TEST_PARAMETRIC_MODEL_HIST_MISSING_KEYWORD = """
type: histogram
n_bins: 5
bin_range: [0, 5]
model_parameters: [0.0, 0.08]
"""

TEST_PARAMETRIC_MODEL_HIST_EXTRA_KEYWORD = (
    TEST_PARAMETRIC_MODEL_HIST
    + """
extra_keyword: 3.14
"""
)

TEST_PARAMETRIC_MODEL_HIST_WITH_ERRORS = (
    TEST_PARAMETRIC_MODEL_HIST
    + """
errors:
  - correlation_coefficient: 0.0
    error_value: 0.1
    name: test_error
    relative: false
    type: simple
"""
)

TEST_PARAMETRIC_MODEL_HIST_LEGACY = (
    TEST_PARAMETRIC_MODEL_HIST
    + """
model_density_func_antiderivative: null
"""
)

TEST_PARAMETRIC_MODEL_HIST_ANTIDERIVATIVE = (
    TEST_PARAMETRIC_MODEL_HIST
    + """
bin_evaluation: |
    def linear_model_antiderivative(x, a, b):
        return 0.5 * a * x ** 2 + b * x
"""
)

TEST_PARAMETRIC_MODEL_HIST_NUMPY_VECTORIZE = (
    TEST_PARAMETRIC_MODEL_HIST
    + """
bin_evaluation: |
    @np.vectorize
    def linear_model_antiderivative(x, a, b):
        return 0.5 * a * x ** 2 + b * x
"""
)


class TestHistParametricModelYamlRepresenter(unittest.TestCase):
    def _assert_read_model_correct(self, read_model, ref_model, errors=False):
        self.assertTrue(isinstance(read_model, HistParametricModel))

        self.assertTrue(np.all(read_model.data == ref_model.data))
        if isinstance(ref_model.bin_evaluation, str):
            self.assertEqual(read_model.bin_evaluation, ref_model.bin_evaluation)
            self.assertEqual(read_model.bin_evaluation_string, ref_model.bin_evaluation_string)
        self.assertTrue(np.all(read_model.parameters == self._test_model_parameters))
        self.assertTrue(read_model.n_bins == self._test_n_bins)
        self.assertTrue(np.all(read_model.bin_range == self._test_bin_range))
        self.assertTrue(np.all(read_model.bin_edges == self._test_bin_edges))
        self.assertTrue(
            np.allclose(
                read_model.eval_model_function_density(self._test_x),
                ref_model.eval_model_function_density(self._test_x),
            )
        )

        if errors:
            _given_error = self._test_parametric_model_with_errors.get_error("test_error")
            _read_error = read_model.get_error("test_error")
            self.assertTrue(_given_error["enabled"] == _read_error["enabled"])
            self.assertTrue(np.allclose(_given_error["err"].error, _read_error["err"].error))
            self.assertTrue(_given_error["err"].corr_coeff == _read_error["err"].corr_coeff)

    @staticmethod
    def _get_model_from_string(input_string):
        _testfile_stringstream = IOStreamHandle(StringIO(input_string))
        return ParametricModelYamlReader(_testfile_stringstream).read()

    @staticmethod
    def _get_model_from_roundtrip(model):
        _roundtrip_stringstream = IOStreamHandle(StringIO())
        _roundtrip_streamreader = ParametricModelYamlReader(_roundtrip_stringstream)
        _roundtrip_streamwriter = ParametricModelYamlWriter(model, _roundtrip_stringstream)
        _roundtrip_streamwriter.write()
        _roundtrip_stringstream.seek(0)  # return to beginning
        return _roundtrip_streamreader.read()

    @staticmethod
    def linear_model(x, a, b):
        return a * x + b

    @staticmethod
    def linear_model_antiderivative(x, a, b):
        return 0.5 * a * x**2 + b * x

    @staticmethod
    @np.vectorize
    def linear_model_numpy_vectorize(x, a, b):
        return 0.5 * a * x**2 + b * x

    def setUp(self):
        self._test_n_bins = 5
        self._test_bin_range = (0, 5)
        self._test_bin_edges = np.arange(6)
        self._test_model_parameters = np.array([0.0, 0.08])
        self._test_x = np.linspace(start=0.0, stop=5.0, num=11, endpoint=True)
        self._test_parametric_model = HistParametricModel(
            n_bins=self._test_n_bins,
            bin_range=self._test_bin_range,
            model_density_func=HistModelFunction(TestHistParametricModelYamlRepresenter.linear_model),
            model_parameters=self._test_model_parameters,
        )

        self._test_parametric_model_with_errors = HistParametricModel(
            n_bins=self._test_n_bins,
            bin_range=self._test_bin_range,
            model_density_func=HistModelFunction(TestHistParametricModelYamlRepresenter.linear_model),
            model_parameters=self._test_model_parameters,
        )
        self._test_parametric_model_with_errors.add_error(err_val=0.1, name="test_error", correlation=0, relative=False)

        self._test_parametric_model_antiderivative = HistParametricModel(
            n_bins=self._test_n_bins,
            bin_range=self._test_bin_range,
            model_density_func=HistModelFunction(TestHistParametricModelYamlRepresenter.linear_model),
            model_parameters=self._test_model_parameters,
            bin_evaluation=TestHistParametricModelYamlRepresenter.linear_model_antiderivative,
        )

        self._test_parametric_model_numpy_vectorize = HistParametricModel(
            n_bins=self._test_n_bins,
            bin_range=self._test_bin_range,
            model_density_func=HistModelFunction(TestHistParametricModelYamlRepresenter.linear_model),
            model_parameters=self._test_model_parameters,
            bin_evaluation=TestHistParametricModelYamlRepresenter.linear_model_numpy_vectorize,
        )

    def test_read_from_testfile_stream(self):
        _read_parametric_model = TestHistParametricModelYamlRepresenter._get_model_from_string(TEST_PARAMETRIC_MODEL_HIST)
        self._assert_read_model_correct(_read_parametric_model, self._test_parametric_model)

    def test_read_from_testfile_stream_antiderivative(self):
        _read_parametric_model = TestHistParametricModelYamlRepresenter._get_model_from_string(TEST_PARAMETRIC_MODEL_HIST_ANTIDERIVATIVE)
        self._assert_read_model_correct(_read_parametric_model, self._test_parametric_model_antiderivative)

    def test_read_from_testfile_stream_numpy_vectorize(self):
        _read_parametric_model = TestHistParametricModelYamlRepresenter._get_model_from_string(TEST_PARAMETRIC_MODEL_HIST_NUMPY_VECTORIZE)
        self._assert_read_model_correct(_read_parametric_model, self._test_parametric_model_numpy_vectorize)

    def test_read_from_testfile_stream_missing_keyword(self):
        with self.assertRaises(ValueError):
            _read_parametric_model = TestHistParametricModelYamlRepresenter._get_model_from_string(TEST_PARAMETRIC_MODEL_HIST_MISSING_KEYWORD)

    def test_read_from_testfile_stream_extra_keyword(self):
        with self.assertRaises(ValueError):
            _read_parametric_model = TestHistParametricModelYamlRepresenter._get_model_from_string(TEST_PARAMETRIC_MODEL_HIST_EXTRA_KEYWORD)

    def test_read_from_testfile_stream_with_errors(self):
        _read_parametric_model = TestHistParametricModelYamlRepresenter._get_model_from_string(TEST_PARAMETRIC_MODEL_HIST_WITH_ERRORS)
        self._assert_read_model_correct(_read_parametric_model, self._test_parametric_model_with_errors, errors=True)

    def test_read_from_testfile_stream_legacy(self):
        _read_parametric_model = TestHistParametricModelYamlRepresenter._get_model_from_string(TEST_PARAMETRIC_MODEL_HIST_LEGACY)
        self._assert_read_model_correct(_read_parametric_model, self._test_parametric_model)

    def test_round_trip_with_stringstream(self):
        _read_parametric_model = self._get_model_from_roundtrip(self._test_parametric_model)
        self._assert_read_model_correct(_read_parametric_model, self._test_parametric_model)

    def test_round_trip_with_stringstream_antiderivative(self):
        _read_parametric_model = self._get_model_from_roundtrip(self._test_parametric_model_antiderivative)
        self._assert_read_model_correct(_read_parametric_model, self._test_parametric_model_antiderivative)

    def test_round_trip_with_stringstream_numpy_vectorize(self):
        _read_parametric_model = self._get_model_from_roundtrip(self._test_parametric_model_numpy_vectorize)
        self._assert_read_model_correct(_read_parametric_model, self._test_parametric_model_numpy_vectorize)

    def test_round_trip_with_stringstream_with_errors(self):
        _read_parametric_model = self._get_model_from_roundtrip(self._test_parametric_model_with_errors)
        self._assert_read_model_correct(_read_parametric_model, self._test_parametric_model_with_errors, errors=True)


TEST_PARAMETRIC_MODEL_INDEXED = """
type: indexed
model_function:
    type: indexed
    python_code: |
        def linear_model(a, b):
            return a * np.arange(6) + b
model_parameters: [1.1, -1.5]
"""

TEST_PARAMETRIC_MODEL_INDEXED_MISSING_KEYWORD = """
type: indexed
model_parameters: [1.1, -1.5]
"""

TEST_PARAMETRIC_MODEL_INDEXED_EXTRA_KEYWORD = (
    TEST_PARAMETRIC_MODEL_INDEXED
    + """
extra_keyword: 3.14
"""
)

TEST_PARAMETRIC_MODEL_INDEXED_WITH_ERRORS = (
    TEST_PARAMETRIC_MODEL_INDEXED
    + """
errors:
  - correlation_coefficient: 0.0
    error_value: 0.1
    name: test_error
    relative: false
    type: simple
"""
)


class TestIndexedParametricModelYamlRepresenter(unittest.TestCase):
    TEST_X = np.arange(6)

    @staticmethod
    def linear_model(a, b):
        return a * np.arange(6) + b

    def setUp(self):
        self._test_model_parameters = np.array([1.1, -1.5])
        self._test_parametric_model = IndexedParametricModel(
            IndexedModelFunction(TestIndexedParametricModelYamlRepresenter.linear_model),
            self._test_model_parameters,
        )

        self._test_parametric_model_with_errors = IndexedParametricModel(
            IndexedModelFunction(TestIndexedParametricModelYamlRepresenter.linear_model),
            self._test_model_parameters,
        )
        self._test_parametric_model_with_errors.add_error(err_val=0.1, name="test_x_error", correlation=0, relative=False)

        self._roundtrip_stringstream = IOStreamHandle(StringIO())
        self._roundtrip_stringstream_with_errors = IOStreamHandle(StringIO())
        self._testfile_stringstream = IOStreamHandle(StringIO(TEST_PARAMETRIC_MODEL_INDEXED))
        self._testfile_stringstream_with_errors = IOStreamHandle(StringIO(TEST_PARAMETRIC_MODEL_INDEXED_WITH_ERRORS))

        self._roundtrip_streamreader = ParametricModelYamlReader(self._roundtrip_stringstream)
        self._roundtrip_streamreader_with_errors = ParametricModelYamlReader(self._roundtrip_stringstream_with_errors)
        self._roundtrip_streamwriter = ParametricModelYamlWriter(self._test_parametric_model, self._roundtrip_stringstream)
        self._roundtrip_streamwriter_with_errors = ParametricModelYamlWriter(
            self._test_parametric_model_with_errors, self._roundtrip_stringstream_with_errors
        )
        self._testfile_streamreader = ParametricModelYamlReader(self._testfile_stringstream)
        self._testfile_streamreader_with_errors = ParametricModelYamlReader(self._testfile_stringstream_with_errors)

        self._testfile_stringstream_missing_keyword = IOStreamHandle(StringIO(TEST_PARAMETRIC_MODEL_INDEXED_MISSING_KEYWORD))
        self._testfile_stringstream_extra_keyword = IOStreamHandle(StringIO(TEST_PARAMETRIC_MODEL_INDEXED_EXTRA_KEYWORD))
        self._testfile_streamreader_missing_keyword = ParametricModelYamlReader(self._testfile_stringstream_missing_keyword)
        self._testfile_streamreader_extra_keyword = ParametricModelYamlReader(self._testfile_stringstream_extra_keyword)

    def test_write_to_roundtrip_stringstream(self):
        self._roundtrip_streamwriter.write()

    def test_write_to_roundtrip_stringstream_with_errors(self):
        self._roundtrip_streamwriter_with_errors.write()

    def test_read_from_testfile_stream(self):
        _read_parametric_model = self._testfile_streamreader.read()
        self.assertTrue(isinstance(_read_parametric_model, IndexedParametricModel))

        self.assertTrue(np.allclose(_read_parametric_model.parameters, self._test_model_parameters))
        self.assertTrue(
            np.allclose(
                _read_parametric_model.eval_model_function(),
                self._test_parametric_model.eval_model_function(),
            )
        )

    def test_read_from_testfile_stream_missing_keyword(self):
        with self.assertRaises(ValueError):
            self._testfile_streamreader_missing_keyword.read()

    def test_read_from_testfile_stream_extra_keyword(self):
        with self.assertRaises(ValueError):
            self._testfile_streamreader_extra_keyword.read()

    def test_read_from_testfile_stream_with_errors(self):
        _read_parametric_model = self._testfile_streamreader_with_errors.read()
        self.assertTrue(isinstance(_read_parametric_model, IndexedParametricModel))

        self.assertTrue(np.allclose(_read_parametric_model.parameters, self._test_model_parameters))
        self.assertTrue(
            np.allclose(
                _read_parametric_model.eval_model_function(),
                self._test_parametric_model_with_errors.eval_model_function(),
            )
        )

        _given_error = self._test_parametric_model_with_errors.get_error("test_x_error")
        _read_error = _read_parametric_model.get_error("test_error")
        self.assertTrue(_given_error["enabled"] == _read_error["enabled"])
        self.assertTrue(np.allclose(_given_error["err"].error, _read_error["err"].error))
        self.assertTrue(_given_error["err"].corr_coeff == _read_error["err"].corr_coeff)

    def test_round_trip_with_stringstream(self):
        self._roundtrip_streamwriter.write()
        self._roundtrip_stringstream.seek(0)  # return to beginning
        _read_parametric_model = self._roundtrip_streamreader.read()
        self.assertTrue(isinstance(_read_parametric_model, IndexedParametricModel))

        self.assertTrue(np.allclose(_read_parametric_model.parameters, self._test_model_parameters))
        self.assertTrue(
            np.allclose(
                _read_parametric_model.eval_model_function(),
                self._test_parametric_model.eval_model_function(),
            )
        )

    def test_round_trip_with_stringstream_with_errors(self):
        self._roundtrip_streamwriter_with_errors.write()
        self._roundtrip_stringstream_with_errors.seek(0)  # return to beginning
        _read_parametric_model = self._roundtrip_streamreader_with_errors.read()
        self.assertTrue(isinstance(_read_parametric_model, IndexedParametricModel))

        self.assertTrue(np.allclose(_read_parametric_model.parameters, self._test_model_parameters))
        self.assertTrue(
            np.allclose(
                _read_parametric_model.eval_model_function(),
                self._test_parametric_model_with_errors.eval_model_function(),
            )
        )

        _given_error = self._test_parametric_model_with_errors.get_error("test_x_error")
        _read_error = _read_parametric_model.get_error("test_x_error")
        self.assertTrue(_given_error["enabled"] == _read_error["enabled"])
        self.assertTrue(np.allclose(_given_error["err"].error, _read_error["err"].error))
        self.assertTrue(_given_error["err"].corr_coeff == _read_error["err"].corr_coeff)


TEST_PARAMETRIC_MODEL_XY = """
type: xy
x_data: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
model_function:
    type: xy
    python_code: |
        def linear_model(x, a, b):
            return a * x + b
model_parameters: [1.1, -1.5]
"""

TEST_PARAMETRIC_MODEL_XY_MISSING_KEYWORD = """
type: xy
model_function:
    type: xy
    python_code: |
        def linear_model(x, a, b):
            return a * x + b
model_parameters: [1.1, -1.5]
"""

TEST_PARAMETRIC_MODEL_XY_EXTRA_KEYWORD = (
    TEST_PARAMETRIC_MODEL_XY
    + """
extra_keyword: 3.14
"""
)

TEST_PARAMETRIC_MODEL_XY_WITH_ERRORS = (
    TEST_PARAMETRIC_MODEL_XY
    + """
x_errors:
  - correlation_coefficient: 0.0
    error_value: 0.1
    name: test_x_error
    relative: false
    type: simple
"""
)


class TestXYParametricModelYamlRepresenter(unittest.TestCase):
    @staticmethod
    def linear_model(x, a, b):
        return a * x + b

    def setUp(self):
        self._test_x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
        self._test_model_parameters = np.array([1.1, -1.5])
        self._test_parametric_model = XYParametricModel(
            self._test_x,
            ModelFunctionBase(TestXYParametricModelYamlRepresenter.linear_model),
            self._test_model_parameters,
        )

        self._test_parametric_model_with_errors = XYParametricModel(
            self._test_x,
            ModelFunctionBase(TestXYParametricModelYamlRepresenter.linear_model),
            self._test_model_parameters,
        )
        self._test_parametric_model_with_errors.add_error(axis="x", err_val=0.1, name="test_x_error", correlation=0, relative=False)

        self._roundtrip_stringstream = IOStreamHandle(StringIO())
        self._roundtrip_stringstream_with_errors = IOStreamHandle(StringIO())
        self._testfile_stringstream = IOStreamHandle(StringIO(TEST_PARAMETRIC_MODEL_XY))
        self._testfile_stringstream_with_errors = IOStreamHandle(StringIO(TEST_PARAMETRIC_MODEL_XY_WITH_ERRORS))

        self._roundtrip_streamreader = ParametricModelYamlReader(self._roundtrip_stringstream)
        self._roundtrip_streamreader_with_errors = ParametricModelYamlReader(self._roundtrip_stringstream_with_errors)
        self._roundtrip_streamwriter = ParametricModelYamlWriter(self._test_parametric_model, self._roundtrip_stringstream)
        self._roundtrip_streamwriter_with_errors = ParametricModelYamlWriter(
            self._test_parametric_model_with_errors, self._roundtrip_stringstream_with_errors
        )
        self._testfile_streamreader = ParametricModelYamlReader(self._testfile_stringstream)
        self._testfile_streamreader_with_errors = ParametricModelYamlReader(self._testfile_stringstream_with_errors)

        self._testfile_stringstream_missing_keyword = IOStreamHandle(StringIO(TEST_PARAMETRIC_MODEL_XY_MISSING_KEYWORD))
        self._testfile_stringstream_extra_keyword = IOStreamHandle(StringIO(TEST_PARAMETRIC_MODEL_XY_EXTRA_KEYWORD))
        self._testfile_streamreader_missing_keyword = ParametricModelYamlReader(self._testfile_stringstream_missing_keyword)
        self._testfile_streamreader_extra_keyword = ParametricModelYamlReader(self._testfile_stringstream_extra_keyword)

    def test_write_to_roundtrip_stringstream(self):
        self._roundtrip_streamwriter.write()

    def test_write_to_roundtrip_stringstream_with_errors(self):
        self._roundtrip_streamwriter_with_errors.write()

    def test_read_from_testfile_stream(self):
        _read_parametric_model = self._testfile_streamreader.read()
        self.assertTrue(isinstance(_read_parametric_model, XYParametricModel))

        self.assertTrue(np.allclose(_read_parametric_model.parameters, self._test_model_parameters))
        _read_parametric_model.x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
        self.assertTrue(np.allclose(_read_parametric_model.y, self._test_parametric_model.y))

    def test_read_from_testfile_stream_missing_keyword(self):
        with self.assertRaises(ValueError):
            self._testfile_streamreader_missing_keyword.read()

    def test_read_from_testfile_stream_extra_keyword(self):
        with self.assertRaises(ValueError):
            self._testfile_streamreader_extra_keyword.read()

    def test_read_from_testfile_stream_with_errors(self):
        _read_parametric_model = self._testfile_streamreader_with_errors.read()
        self.assertTrue(isinstance(_read_parametric_model, XYParametricModel))

        self.assertTrue(np.allclose(_read_parametric_model.parameters, self._test_model_parameters))
        _read_parametric_model.x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
        self.assertTrue(np.allclose(_read_parametric_model.y, self._test_parametric_model_with_errors.y))

        _given_error = self._test_parametric_model_with_errors.get_error("test_x_error")
        _read_error = _read_parametric_model.get_error("test_x_error")
        self.assertTrue(_given_error["axis"] == _read_error["axis"])
        self.assertTrue(_given_error["enabled"] == _read_error["enabled"])
        self.assertTrue(np.allclose(_given_error["err"].error, _read_error["err"].error))
        self.assertTrue(_given_error["err"].corr_coeff == _read_error["err"].corr_coeff)

    def test_round_trip_with_stringstream(self):
        self._roundtrip_streamwriter.write()
        self._roundtrip_stringstream.seek(0)  # return to beginning
        _read_parametric_model = self._roundtrip_streamreader.read()
        self.assertTrue(isinstance(_read_parametric_model, XYParametricModel))

        self.assertTrue(np.allclose(_read_parametric_model.parameters, self._test_model_parameters))
        _read_parametric_model.x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
        self.assertTrue(np.allclose(_read_parametric_model.y, self._test_parametric_model.y))

    def test_round_trip_with_stringstream_with_errors(self):
        self._roundtrip_streamwriter_with_errors.write()
        self._roundtrip_stringstream_with_errors.seek(0)  # return to beginning
        _read_parametric_model = self._roundtrip_streamreader_with_errors.read()
        self.assertTrue(isinstance(_read_parametric_model, XYParametricModel))

        self.assertTrue(np.allclose(_read_parametric_model.parameters, self._test_model_parameters))
        _read_parametric_model.x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
        self.assertTrue(np.allclose(_read_parametric_model.y, self._test_parametric_model_with_errors.y))

        _given_error = self._test_parametric_model_with_errors.get_error("test_x_error")
        _read_error = _read_parametric_model.get_error("test_x_error")
        self.assertTrue(_given_error["axis"] == _read_error["axis"])
        self.assertTrue(_given_error["enabled"] == _read_error["enabled"])
        self.assertTrue(np.allclose(_given_error["err"].error, _read_error["err"].error))
        self.assertTrue(_given_error["err"].corr_coeff == _read_error["err"].corr_coeff)
