import pytest
import numpy as np

from matplotlib.sankey import Sankey
from matplotlib.testing.decorators import check_figures_equal


def test_sankey():
    # lets just create a sankey instance and check the code runs
    sankey = Sankey()
    sankey.add()


def test_label():
    s = Sankey(flows=[0.25], labels=['First'], orientations=[-1])
    assert s.diagrams[0].texts[0].get_text() == 'First\n0.25'


def test_format_using_callable():
    # test using callable by slightly incrementing above label example

    def show_three_decimal_places(value):
        return f'{value:.3f}'

    s = Sankey(flows=[0.25], labels=['First'], orientations=[-1],
               format=show_three_decimal_places)

    assert s.diagrams[0].texts[0].get_text() == 'First\n0.250'


@pytest.mark.parametrize('kwargs, msg', (
    ({'gap': -1}, "'gap' is negative"),
    ({'gap': 1, 'radius': 2}, "'radius' is greater than 'gap'"),
    ({'head_angle': -1}, "'head_angle' is negative"),
    ({'tolerance': -1}, "'tolerance' is negative"),
    ({'flows': [1, -1], 'orientations': [-1, 0, 1]},
     r"The shapes of 'flows' \(2,\) and 'orientations'"),
    ({'flows': [1, -1], 'labels': ['a', 'b', 'c']},
     r"The shapes of 'flows' \(2,\) and 'labels'"),
    ))
def test_sankey_errors(kwargs, msg):
    with pytest.raises(ValueError, match=msg):
        Sankey(**kwargs)


@pytest.mark.parametrize('kwargs, msg', (
    ({'trunklength': -1}, "'trunklength' is negative"),
    ({'flows': [0.2, 0.3], 'prior': 0}, "The scaled sum of the connected"),
    ({'prior': -1}, "The index of the prior diagram is negative"),
    ({'prior': 1}, "The index of the prior diagram is 1"),
    ({'connect': (-1, 1), 'prior': 0}, "At least one of the connection"),
    ({'connect': (2, 1), 'prior': 0}, "The connection index to the source"),
    ({'connect': (1, 3), 'prior': 0}, "The connection index to this dia"),
    ({'connect': (1, 1), 'prior': 0, 'flows': [-0.2, 0.2],
      'orientations': [2]}, "The value of orientations"),
    ({'connect': (1, 1), 'prior': 0, 'flows': [-0.2, 0.2],
      'pathlengths': [2]}, "The lengths of 'flows'"),
    ))
def test_sankey_add_errors(kwargs, msg):
    sankey = Sankey()
    with pytest.raises(ValueError, match=msg):
        sankey.add(flows=[0.2, -0.2])
        sankey.add(**kwargs)


def test_sankey2():
    s = Sankey(flows=[0.25, -0.25, 0.5, -0.5], labels=['Foo'],
               orientations=[-1], unit='Bar')
    sf = s.finish()
    assert np.all(np.equal(np.array((0.25, -0.25, 0.5, -0.5)), sf[0].flows))
    assert sf[0].angles == [1, 3, 1, 3]
    assert all([text.get_text()[0:3] == 'Foo' for text in sf[0].texts])
    assert all([text.get_text()[-3:] == 'Bar' for text in sf[0].texts])
    assert sf[0].text.get_text() == ''
    assert np.allclose(np.array(((-1.375, -0.52011255),
                                 (1.375, -0.75506044),
                                 (-0.75, -0.41522509),
                                 (0.75, -0.8599479))),
                       sf[0].tips)

    s = Sankey(flows=[0.25, -0.25, 0, 0.5, -0.5], labels=['Foo'],
               orientations=[-1], unit='Bar')
    sf = s.finish()
    assert np.all(np.equal(np.array((0.25, -0.25, 0, 0.5, -0.5)), sf[0].flows))
    assert sf[0].angles == [1, 3, None, 1, 3]
    assert np.allclose(np.array(((-1.375, -0.52011255),
                                 (1.375, -0.75506044),
                                 (0, 0),
                                 (-0.75, -0.41522509),
                                 (0.75, -0.8599479))),
                       sf[0].tips)


@check_figures_equal(extensions=['png'])
def test_sankey3(fig_test, fig_ref):
    ax_test = fig_test.gca()
    s_test = Sankey(ax=ax_test, flows=[0.25, -0.25, -0.25, 0.25, 0.5, -0.5],
                    orientations=[1, -1, 1, -1, 0, 0])
    s_test.finish()

    ax_ref = fig_ref.gca()
    s_ref = Sankey(ax=ax_ref)
    s_ref.add(flows=[0.25, -0.25, -0.25, 0.25, 0.5, -0.5],
              orientations=[1, -1, 1, -1, 0, 0])
    s_ref.finish()