import numpy as np

from numpy.testing import assert_allclose

from echo import CallbackProperty, ListCallbackProperty
from glue.core import Data, DataCollection

from .test_state import clone
from ..state_objects import (State, StateAttributeLimitsHelper,
                             StateAttributeSingleValueHelper,
                             StateAttributeHistogramHelper)


class SimpleTestState(State):
    a = CallbackProperty()
    b = CallbackProperty()
    flat = ListCallbackProperty()
    nested = ListCallbackProperty()


def test_state_serialization():

    state1 = SimpleTestState()
    state1.a = 2
    state1.b = 'hello'
    state1.flat = [1, 3, 4]

    sub_state = SimpleTestState()
    sub_state.a = 3
    sub_state.b = 'blah'
    sub_state.flat = [1, 2]
    sub_state.nested = []

    state1.nested = [1, 3, sub_state]

    state2 = clone(state1)

    assert state2.a == 2
    assert state2.b == 'hello'
    assert state2.flat == [1, 3, 4]
    assert state2.nested[0:2] == [1, 3]
    assert state2.nested[2].a == 3
    assert state2.nested[2].b == 'blah'
    assert state2.nested[2].flat == [1, 2]
    assert state2.nested[2].nested == []


EXPECTED_STR = """
a: 2
b: hello
flat: <CallbackList with 3 elements>
nested: <CallbackList with 3 elements>
"""

EXPECTED_REPR = """
<SimpleTestState
  a: 2
  b: hello
  flat: <CallbackList with 3 elements>
  nested: <CallbackList with 3 elements>
>
"""


def test_state_str_repr():

    state1 = SimpleTestState()
    state1.a = 2
    state1.b = 'hello'
    state1.flat = [1, 3, 4]

    sub_state = SimpleTestState()

    state1.nested = [1, 3, sub_state]

    assert str(state1) == EXPECTED_STR.strip()
    assert repr(state1) == EXPECTED_REPR.strip()


class TestStateAttributeLimitsHelper():

    def setup_method(self, method):

        self.data = Data(x=np.linspace(-100, 100, 10000),
                         y=np.linspace(2, 3, 10000), label='test_data')

        self.data_collection = DataCollection([self.data])

        class SimpleState(State):

            layer = CallbackProperty()
            comp = CallbackProperty()
            lower = CallbackProperty()
            upper = CallbackProperty()
            log = CallbackProperty(False)
            scale = CallbackProperty(100)

        self.state = SimpleState()

        self.helper = StateAttributeLimitsHelper(self.state, attribute='comp',
                                                 lower='lower', upper='upper',
                                                 percentile='scale', log='log')
        self.state.data = self.data
        self.state.comp = self.data.id['x']

        self.x_id = self.data.main_components[0]
        self.y_id = self.data.main_components[1]

    def test_minmax(self):
        assert self.helper.lower == -100
        assert self.helper.upper == +100

    def test_change_attribute(self):
        self.helper.attribute = self.y_id
        assert self.helper.lower == 2
        assert self.helper.upper == 3
        self.helper.attribute = self.x_id
        assert self.helper.lower == -100
        assert self.helper.upper == +100

    def test_change_percentile(self):

        # Changing scale mode updates the limits
        self.helper.percentile = 99.5
        assert_allclose(self.helper.lower, -99.5)
        assert_allclose(self.helper.upper, +99.5)
        self.helper.percentile = 99
        assert_allclose(self.helper.lower, -99)
        assert_allclose(self.helper.upper, +99)
        self.helper.percentile = 90
        assert_allclose(self.helper.lower, -90)
        assert_allclose(self.helper.upper, +90)

        # When switching to custom, the last limits are retained
        self.helper.percentile = "Custom"
        assert_allclose(self.helper.lower, -90)
        assert_allclose(self.helper.upper, +90)

    def test_percentile_cached(self):
        # Make sure that if we change scale and change attribute, the scale
        # modes are cached on a per-attribute basis.
        self.helper.percentile = 99.5
        self.state.comp = self.y_id
        assert self.helper.percentile == 100
        self.helper.percentile = 99
        self.state.comp = self.x_id
        assert self.helper.percentile == 99.5
        self.state.comp = self.y_id
        assert self.helper.percentile == 99

    def test_flip_button(self):

        self.helper.flip_limits()

        assert self.helper.lower == +100
        assert self.helper.upper == -100

        # Make sure that values were re-cached when flipping
        self.state.comp = self.y_id
        assert self.helper.lower == 2
        assert self.helper.upper == 3
        self.state.comp = self.x_id
        assert self.helper.lower == +100
        assert self.helper.upper == -100

    def test_manual_edit(self):

        # Make sure that values are re-cached when edited manually
        self.helper.percentile = "Custom"
        self.state.lower = -122
        self.state.upper = 234
        self.helper.log = True
        assert self.helper.lower == -122
        assert self.helper.upper == 234
        assert self.helper.log
        self.state.comp = self.y_id
        assert self.helper.lower == 2
        assert self.helper.upper == 3
        assert not self.helper.log
        self.state.comp = self.x_id
        assert self.helper.lower == -122
        assert self.helper.upper == 234
        assert self.helper.log


class TestStateAttributeSingleValueHelper():

    def setup_method(self, method):

        self.data = Data(x=np.linspace(-100, 30, 9999),
                         y=np.linspace(2, 3, 9999), label='test_data')

        self.data_collection = DataCollection([self.data])

        class SimpleState(State):

            layer = CallbackProperty()
            comp = CallbackProperty()
            val = CallbackProperty()

        self.state = SimpleState()

        self.helper = StateAttributeSingleValueHelper(self.state, attribute='comp',
                                                      function=np.nanmedian, value='val')

        self.state.data = self.data

        self.state.comp = self.data.id['x']

        self.x_id = self.data.main_components[0]
        self.y_id = self.data.main_components[1]

    def test_value(self):
        assert self.helper.value == -35.

    def test_change_attribute(self):
        self.helper.attribute = self.y_id
        assert self.helper.value == 2.5
        self.helper.attribute = self.x_id
        assert self.helper.value == -35

    def test_manual_edit(self):
        self.state.val = 42.
        assert self.helper.value == 42
        self.state.comp = self.y_id
        assert self.helper.value == 2.5
        self.state.comp = self.x_id
        assert self.helper.value == 42


class TestStateAttributeHistogramHelper():

    def setup_method(self, method):

        self.data = Data(x=[-3.2, 4.3, 2.2, 5.4, 7.2, -1.1, 2.3],
                         y=['a', 'f', 'd', 'e', 'f', 'f', 'a'], label='test_data')

        self.data_collection = DataCollection([self.data])

        class SimpleState(State):

            layer = CallbackProperty()
            comp = CallbackProperty()
            x_min = CallbackProperty()
            x_max = CallbackProperty()
            n_bin = CallbackProperty()

        self.state = SimpleState()

        self.helper = StateAttributeHistogramHelper(self.state, attribute='comp',
                                                    lower='x_min', upper='x_max', n_bin='n_bin')

        self.state.data = self.data

    def test_default_numerical(self):
        self.state.comp = self.data.id['x']
        assert self.state.x_min == -3.2
        assert self.state.x_max == 7.2
        assert self.state.n_bin == 15

    def test_default_categorical(self):
        self.state.comp = self.data.id['y']
        assert self.state.x_min == -0.5
        assert self.state.x_max == 3.5
        assert self.state.n_bin == 4

    def test_hitting_limits(self):

        # FIXME: here we modify the internal defaults rather than making a new
        # state helper, but this could be improved
        self.helper._default_n_bin = 4
        self.helper._max_n_bin = 3

        self.state.comp = self.data.id['x']
        assert self.state.x_min == -3.2
        assert self.state.x_max == 7.2
        assert self.state.n_bin == 4

        self.state.comp = self.data.id['y']
        assert self.state.x_min == -0.5
        assert self.state.x_max == 3.5
        assert self.state.n_bin == 3

    def test_caching(self):
        self.state.comp = self.data.id['x']
        self.state.x_min = 2
        self.state.x_max = 7
        self.state.n_bin = 8
        self.state.comp = self.data.id['y']
        self.state.x_min = 1.5
        self.state.x_max = 3.5
        self.state.n_bin = 3
        self.state.comp = self.data.id['x']
        assert self.state.x_min == 2
        assert self.state.x_max == 7
        assert self.state.n_bin == 8
        self.state.comp = self.data.id['y']
        assert self.state.x_min == 1.5
        assert self.state.x_max == 3.5
        assert self.state.n_bin == 3


def test_histogram_helper_common_n_bin():

    data = Data(x=[-3.2, 4.3, 2.2],
                y=['a', 'f', 'd'],
                z=[1.1, 2.3, 1.2],
                label='test_data')

    class SimpleState(State):

        layer = CallbackProperty()
        comp = CallbackProperty()
        x_min = CallbackProperty()
        x_max = CallbackProperty()
        n_bin = CallbackProperty()
        common = CallbackProperty()

    state = SimpleState()

    helper = StateAttributeHistogramHelper(state, attribute='comp',
                                           lower='x_min', upper='x_max', n_bin='n_bin',
                                           common_n_bin='common')

    state.data = data

    state.comp = data.id['x']
    state.n_bin = 9
    state.comp = data.id['y']
    assert state.n_bin == 3
    state.comp = data.id['z']
    assert state.n_bin == 15

    state.n_bin = 12

    state.common = True

    state.comp = data.id['x']
    assert state.n_bin == 12

    state.n_bin = 11

    state.comp = data.id['y']
    assert state.n_bin == 3
    state.comp = data.id['z']
    assert state.n_bin == 11

    state.common = False
    state.n_bin = 13

    state.comp = data.id['x']
    assert state.n_bin == 11


def test_histogram_helper_common_n_bin_active():

    # Make sure that common_n_bin works as expected if True from start

    data = Data(x=[-3.2, 4.3, 2.2],
                y=['a', 'f', 'd'],
                z=[1.1, 2.3, 1.2],
                label='test_data')

    class SimpleState(State):

        layer = CallbackProperty()
        comp = CallbackProperty()
        x_min = CallbackProperty()
        x_max = CallbackProperty()
        n_bin = CallbackProperty()
        common = CallbackProperty(True)

    state = SimpleState()

    helper = StateAttributeHistogramHelper(state, attribute='comp',
                                           lower='x_min', upper='x_max', n_bin='n_bin',
                                           common_n_bin='common')

    state.data = data

    state.comp = data.id['x']
    state.n_bin = 9
    state.comp = data.id['z']
    assert state.n_bin == 9

    state.n_bin = 12

    state.common = True

    state.comp = data.id['x']
    assert state.n_bin == 12

    state.n_bin = 11

    state.comp = data.id['y']
    assert state.n_bin == 3
    state.comp = data.id['z']
    assert state.n_bin == 11

    state.common = False
    state.n_bin = 13

    state.comp = data.id['x']
    assert state.n_bin == 11


def test_limits_helper_initial_values():

    # Regression test for a bug that occurred if the limits cache was empty
    # but some attributes were set to values - in this case we don't want to
    # override the existing values.

    data = Data(x=np.linspace(-100, 100, 10000),
                y=np.linspace(2, 3, 10000), label='test_data')

    class SimpleState(State):

        layer = CallbackProperty()
        comp = CallbackProperty()
        lower = CallbackProperty()
        upper = CallbackProperty()

    state = SimpleState()
    state.lower = 1
    state.upper = 2
    state.comp = data.id['x']

    helper = StateAttributeLimitsHelper(state, attribute='comp',
                                        lower='lower', upper='upper')

    assert helper.lower == 1
    assert helper.upper == 2


class DatetimeState(State):
    a = CallbackProperty()


def test_state_serialization_datetime64():

    state1 = DatetimeState()
    state1.a = np.datetime64(100, 'D')

    state2 = clone(state1)

    assert state2.a == np.datetime64(100, 'D')


def test_nan_inf_minmax():

    data = Data(x=[3, 1, -2, np.inf, np.nan], label='test_data')

    class SimpleState(State):

        layer = CallbackProperty()
        comp = CallbackProperty()
        lower = CallbackProperty()
        upper = CallbackProperty()
        percentile = CallbackProperty()
        log = CallbackProperty()

    state = SimpleState()

    helper = StateAttributeLimitsHelper(state, attribute='comp',  # noqa
                                        lower='lower', upper='upper',
                                        percentile='percentile', log='log')

    state.data = data
    state.comp = data.id['x']

    assert state.lower == -2
    assert state.upper == +3

    state.log = True

    assert state.lower == +1
    assert state.upper == +3

    state.log = False
    state.percentile = 99

    assert_allclose(state.lower, -1.97)
    assert_allclose(state.upper, +2.98)


def test_percentile_no_log():

    # Regression test for a bug that caused a crash if the state class had a
    # percentile attribute but no log.

    data = Data(x=np.linspace(-100, 100, 10000),
                y=np.linspace(2, 3, 10000), label='test_data')

    class SimpleState(State):

        layer = CallbackProperty()
        comp = CallbackProperty()
        lower = CallbackProperty()
        upper = CallbackProperty()
        scale = CallbackProperty()

    state = SimpleState()

    state.comp = data.id['x']
    state.lower = 2
    state.upper = 4

    helper = StateAttributeLimitsHelper(state, attribute='comp',
                                        lower='lower', upper='upper',
                                        percentile='scale')

    state.scale = 90

[(15, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (16, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (17, 'echo.ListCallbackProperty', 'ListCallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (18, 'echo.ListCallbackProperty', 'ListCallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (305, 'glue.core.Data', 'Data', False, 'from glue.core import Data, DataCollection\n'), (359, 'glue.core.Data', 'Data', False, 'from glue.core import Data, DataCollection\n'), (436, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (442, 'numpy.datetime64', 'np.datetime64', True, 'import numpy as np\n'), (451, 'glue.core.Data', 'Data', False, 'from glue.core import Data, DataCollection\n'), (482, 'numpy.testing.assert_allclose', 'assert_allclose', False, 'from numpy.testing import assert_allclose\n'), (483, 'numpy.testing.assert_allclose', 'assert_allclose', False, 'from numpy.testing import assert_allclose\n'), (87, 'glue.core.DataCollection', 'DataCollection', False, 'from glue.core import Data, DataCollection\n'), (125, 'numpy.testing.assert_allclose', 'assert_allclose', False, 'from numpy.testing import assert_allclose\n'), (126, 'numpy.testing.assert_allclose', 'assert_allclose', False, 'from numpy.testing import assert_allclose\n'), (128, 'numpy.testing.assert_allclose', 'assert_allclose', False, 'from numpy.testing import assert_allclose\n'), (129, 'numpy.testing.assert_allclose', 'assert_allclose', False, 'from numpy.testing import assert_allclose\n'), (131, 'numpy.testing.assert_allclose', 'assert_allclose', False, 'from numpy.testing import assert_allclose\n'), (132, 'numpy.testing.assert_allclose', 'assert_allclose', False, 'from numpy.testing import assert_allclose\n'), (136, 'numpy.testing.assert_allclose', 'assert_allclose', False, 'from numpy.testing import assert_allclose\n'), (137, 'numpy.testing.assert_allclose', 'assert_allclose', False, 'from numpy.testing import assert_allclose\n'), (193, 'glue.core.DataCollection', 'DataCollection', False, 'from glue.core import Data, DataCollection\n'), (235, 'glue.core.Data', 'Data', False, 'from glue.core import Data, DataCollection\n'), (238, 'glue.core.DataCollection', 'DataCollection', False, 'from glue.core import Data, DataCollection\n'), (312, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (313, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (314, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (315, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (316, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (317, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (366, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (367, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (368, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (369, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (370, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (371, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (418, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (419, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (420, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (421, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (446, 'numpy.datetime64', 'np.datetime64', True, 'import numpy as np\n'), (455, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (456, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (457, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (458, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (459, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (460, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (496, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (497, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (498, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (499, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (500, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (91, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (92, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (93, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (94, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (95, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (96, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (197, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (198, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (199, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (242, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (243, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (244, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (245, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (246, 'echo.CallbackProperty', 'CallbackProperty', False, 'from echo import CallbackProperty, ListCallbackProperty\n'), (413, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (414, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (491, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (492, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (84, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (85, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (190, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (191, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n')]