import copy
import re
import textwrap

import numpy as np
import pytest

import pandas.util._test_decorators as td

import pandas as pd
from pandas import DataFrame
import pandas.util.testing as tm

jinja2 = pytest.importorskip('jinja2')
from pandas.io.formats.style import Styler, _get_level_lengths  # noqa  # isort:skip


class TestStyler(object):

    def setup_method(self, method):
        np.random.seed(24)
        self.s = DataFrame({'A': np.random.permutation(range(6))})
        self.df = DataFrame({'A': [0, 1], 'B': np.random.randn(2)})
        self.f = lambda x: x
        self.g = lambda x: x

        def h(x, foo='bar'):
            return pd.Series(
                'color: {foo}'.format(foo=foo), index=x.index, name=x.name)

        self.h = h
        self.styler = Styler(self.df)
        self.attrs = pd.DataFrame({'A': ['color: red', 'color: blue']})
        self.dataframes = [
            self.df,
            pd.DataFrame({'f': [1., 2.], 'o': ['a', 'b'],
                          'c': pd.Categorical(['a', 'b'])})
        ]

    def test_init_non_pandas(self):
        with pytest.raises(TypeError):
            Styler([1, 2, 3])

    def test_init_series(self):
        result = Styler(pd.Series([1, 2]))
        assert result.data.ndim == 2

    def test_repr_html_ok(self):
        self.styler._repr_html_()

    def test_repr_html_mathjax(self):
        # gh-19824
        assert 'tex2jax_ignore' not in self.styler._repr_html_()

        with pd.option_context('display.html.use_mathjax', False):
            assert 'tex2jax_ignore' in self.styler._repr_html_()

    def test_update_ctx(self):
        self.styler._update_ctx(self.attrs)
        expected = {(0, 0): ['color: red'],
                    (1, 0): ['color: blue']}
        assert self.styler.ctx == expected

    def test_update_ctx_flatten_multi(self):
        attrs = DataFrame({"A": ['color: red; foo: bar',
                                 'color: blue; foo: baz']})
        self.styler._update_ctx(attrs)
        expected = {(0, 0): ['color: red', ' foo: bar'],
                    (1, 0): ['color: blue', ' foo: baz']}
        assert self.styler.ctx == expected

    def test_update_ctx_flatten_multi_traliing_semi(self):
        attrs = DataFrame({"A": ['color: red; foo: bar;',
                                 'color: blue; foo: baz;']})
        self.styler._update_ctx(attrs)
        expected = {(0, 0): ['color: red', ' foo: bar'],
                    (1, 0): ['color: blue', ' foo: baz']}
        assert self.styler.ctx == expected

    def test_copy(self):
        s2 = copy.copy(self.styler)
        assert self.styler is not s2
        assert self.styler.ctx is s2.ctx  # shallow
        assert self.styler._todo is s2._todo

        self.styler._update_ctx(self.attrs)
        self.styler.highlight_max()
        assert self.styler.ctx == s2.ctx
        assert self.styler._todo == s2._todo

    def test_deepcopy(self):
        s2 = copy.deepcopy(self.styler)
        assert self.styler is not s2
        assert self.styler.ctx is not s2.ctx
        assert self.styler._todo is not s2._todo

        self.styler._update_ctx(self.attrs)
        self.styler.highlight_max()
        assert self.styler.ctx != s2.ctx
        assert s2._todo == []
        assert self.styler._todo != s2._todo

    def test_clear(self):
        s = self.df.style.highlight_max()._compute()
        assert len(s.ctx) > 0
        assert len(s._todo) > 0
        s.clear()
        assert len(s.ctx) == 0
        assert len(s._todo) == 0

    def test_render(self):
        df = pd.DataFrame({"A": [0, 1]})
        style = lambda x: pd.Series(["color: red", "color: blue"], name=x.name)
        s = Styler(df, uuid='AB').apply(style)
        s.render()
        # it worked?

    def test_render_empty_dfs(self):
        empty_df = DataFrame()
        es = Styler(empty_df)
        es.render()
        # An index but no columns
        DataFrame(columns=['a']).style.render()
        # A column but no index
        DataFrame(index=['a']).style.render()
        # No IndexError raised?

    def test_render_double(self):
        df = pd.DataFrame({"A": [0, 1]})
        style = lambda x: pd.Series(["color: red; border: 1px",
                                     "color: blue; border: 2px"], name=x.name)
        s = Styler(df, uuid='AB').apply(style)
        s.render()
        # it worked?

    def test_set_properties(self):
        df = pd.DataFrame({"A": [0, 1]})
        result = df.style.set_properties(color='white',
                                         size='10px')._compute().ctx
        # order is deterministic
        v = ["color: white", "size: 10px"]
        expected = {(0, 0): v, (1, 0): v}
        assert result.keys() == expected.keys()
        for v1, v2 in zip(result.values(), expected.values()):
            assert sorted(v1) == sorted(v2)

    def test_set_properties_subset(self):
        df = pd.DataFrame({'A': [0, 1]})
        result = df.style.set_properties(subset=pd.IndexSlice[0, 'A'],
                                         color='white')._compute().ctx
        expected = {(0, 0): ['color: white']}
        assert result == expected

    def test_empty_index_name_doesnt_display(self):
        # https://github.com/pandas-dev/pandas/pull/12090#issuecomment-180695902
        df = pd.DataFrame({'A': [1, 2], 'B': [3, 4], 'C': [5, 6]})
        result = df.style._translate()

        expected = [[{'class': 'blank level0', 'type': 'th', 'value': '',
                      'is_visible': True, 'display_value': ''},
                     {'class': 'col_heading level0 col0',
                      'display_value': 'A',
                      'type': 'th',
                      'value': 'A',
                      'is_visible': True,
                      },
                     {'class': 'col_heading level0 col1',
                      'display_value': 'B',
                      'type': 'th',
                      'value': 'B',
                      'is_visible': True,
                      },
                     {'class': 'col_heading level0 col2',
                      'display_value': 'C',
                      'type': 'th',
                      'value': 'C',
                      'is_visible': True,
                      }]]

        assert result['head'] == expected

    def test_index_name(self):
        # https://github.com/pandas-dev/pandas/issues/11655
        df = pd.DataFrame({'A': [1, 2], 'B': [3, 4], 'C': [5, 6]})
        result = df.set_index('A').style._translate()

        expected = [[{'class': 'blank level0', 'type': 'th', 'value': '',
                      'display_value': '', 'is_visible': True},
                     {'class': 'col_heading level0 col0', 'type': 'th',
                      'value': 'B', 'display_value': 'B', 'is_visible': True},
                     {'class': 'col_heading level0 col1', 'type': 'th',
                      'value': 'C', 'display_value': 'C', 'is_visible': True}],
                    [{'class': 'index_name level0', 'type': 'th',
                      'value': 'A'},
                     {'class': 'blank', 'type': 'th', 'value': ''},
                     {'class': 'blank', 'type': 'th', 'value': ''}]]

        assert result['head'] == expected

    def test_multiindex_name(self):
        # https://github.com/pandas-dev/pandas/issues/11655
        df = pd.DataFrame({'A': [1, 2], 'B': [3, 4], 'C': [5, 6]})
        result = df.set_index(['A', 'B']).style._translate()

        expected = [[
            {'class': 'blank', 'type': 'th', 'value': '',
             'display_value': '', 'is_visible': True},
            {'class': 'blank level0', 'type': 'th', 'value': '',
             'display_value': '', 'is_visible': True},
            {'class': 'col_heading level0 col0', 'type': 'th',
             'value': 'C', 'display_value': 'C', 'is_visible': True}],
            [{'class': 'index_name level0', 'type': 'th',
              'value': 'A'},
             {'class': 'index_name level1', 'type': 'th',
              'value': 'B'},
             {'class': 'blank', 'type': 'th', 'value': ''}]]

        assert result['head'] == expected

    def test_numeric_columns(self):
        # https://github.com/pandas-dev/pandas/issues/12125
        # smoke test for _translate
        df = pd.DataFrame({0: [1, 2, 3]})
        df.style._translate()

    def test_apply_axis(self):
        df = pd.DataFrame({'A': [0, 0], 'B': [1, 1]})
        f = lambda x: ['val: {max}'.format(max=x.max()) for v in x]
        result = df.style.apply(f, axis=1)
        assert len(result._todo) == 1
        assert len(result.ctx) == 0
        result._compute()
        expected = {(0, 0): ['val: 1'], (0, 1): ['val: 1'],
                    (1, 0): ['val: 1'], (1, 1): ['val: 1']}
        assert result.ctx == expected

        result = df.style.apply(f, axis=0)
        expected = {(0, 0): ['val: 0'], (0, 1): ['val: 1'],
                    (1, 0): ['val: 0'], (1, 1): ['val: 1']}
        result._compute()
        assert result.ctx == expected
        result = df.style.apply(f)  # default
        result._compute()
        assert result.ctx == expected

    def test_apply_subset(self):
        axes = [0, 1]
        slices = [pd.IndexSlice[:], pd.IndexSlice[:, ['A']],
                  pd.IndexSlice[[1], :], pd.IndexSlice[[1], ['A']],
                  pd.IndexSlice[:2, ['A', 'B']]]
        for ax in axes:
            for slice_ in slices:
                result = self.df.style.apply(self.h, axis=ax, subset=slice_,
                                             foo='baz')._compute().ctx
                expected = {(r, c): ['color: baz']
                            for r, row in enumerate(self.df.index)
                            for c, col in enumerate(self.df.columns)
                            if row in self.df.loc[slice_].index and
                            col in self.df.loc[slice_].columns}
                assert result == expected

    def test_applymap_subset(self):
        def f(x):
            return 'foo: bar'

        slices = [pd.IndexSlice[:], pd.IndexSlice[:, ['A']],
                  pd.IndexSlice[[1], :], pd.IndexSlice[[1], ['A']],
                  pd.IndexSlice[:2, ['A', 'B']]]

        for slice_ in slices:
            result = self.df.style.applymap(f, subset=slice_)._compute().ctx
            expected = {(r, c): ['foo: bar']
                        for r, row in enumerate(self.df.index)
                        for c, col in enumerate(self.df.columns)
                        if row in self.df.loc[slice_].index and
                        col in self.df.loc[slice_].columns}
            assert result == expected

    def test_applymap_subset_multiindex(self):
        # GH 19861
        # Smoke test for applymap
        def color_negative_red(val):
            """
            Takes a scalar and returns a string with
            the css property `'color: red'` for negative
            strings, black otherwise.
            """
            color = 'red' if val < 0 else 'black'
            return 'color: %s' % color

        dic = {
            ('a', 'd'): [-1.12, 2.11],
            ('a', 'c'): [2.78, -2.88],
            ('b', 'c'): [-3.99, 3.77],
            ('b', 'd'): [4.21, -1.22],
        }

        idx = pd.IndexSlice
        df = pd.DataFrame(dic, index=[0, 1])

        (df.style
         .applymap(color_negative_red, subset=idx[:, idx['b', 'd']])
         .render())

    def test_where_with_one_style(self):
        # GH 17474
        def f(x):
            return x > 0.5

        style1 = 'foo: bar'

        result = self.df.style.where(f, style1)._compute().ctx
        expected = {(r, c): [style1 if f(self.df.loc[row, col]) else '']
                    for r, row in enumerate(self.df.index)
                    for c, col in enumerate(self.df.columns)}
        assert result == expected

    def test_where_subset(self):
        # GH 17474
        def f(x):
            return x > 0.5

        style1 = 'foo: bar'
        style2 = 'baz: foo'

        slices = [pd.IndexSlice[:], pd.IndexSlice[:, ['A']],
                  pd.IndexSlice[[1], :], pd.IndexSlice[[1], ['A']],
                  pd.IndexSlice[:2, ['A', 'B']]]

        for slice_ in slices:
            result = self.df.style.where(f, style1, style2,
                                         subset=slice_)._compute().ctx
            expected = {(r, c):
                        [style1 if f(self.df.loc[row, col]) else style2]
                        for r, row in enumerate(self.df.index)
                        for c, col in enumerate(self.df.columns)
                        if row in self.df.loc[slice_].index and
                        col in self.df.loc[slice_].columns}
            assert result == expected

    def test_where_subset_compare_with_applymap(self):
        # GH 17474
        def f(x):
            return x > 0.5