from wordcloud import WordCloud, get_single_color_func, ImageColorGenerator

import numpy as np
import pytest

from random import Random
from numpy.testing import assert_array_equal
from PIL import Image
import xml.etree.ElementTree as ET

import matplotlib
matplotlib.use('Agg')

THIS = """The Zen of Python, by Tim Peters

Beautiful is better than ugly.
Explicit is better than implicit.
Simple is better than complex.
Complex is better than complicated.
Flat is better than nested.
Sparse is better than dense.
Readability counts.
Special cases aren't special enough to break the rules.
Although practicality beats purity.
Errors should never pass silently.
Unless explicitly silenced.
In the face of ambiguity, refuse the temptation to guess.
There should be one-- and preferably only one --obvious way to do it.
Although that way may not be obvious at first unless you're Dutch.
Now is better than never.
Although never is often better than *right* now.
If the implementation is hard to explain, it's a bad idea.
If the implementation is easy to explain, it may be a good idea.
Namespaces are one honking great idea -- let's do more of those!

3 . 14 15 92 65 35   89 79 32 38 46   26 433
    83 27 95 02 88   41 97 16 93 99   37 510
    58 20 97 49 44   59 23 07 81 64   06 286
    20 89 98 62 80   34 82 53 42 11   70 679
    82 14 80 86 51   32 82 30 66 47   09 384
    46 09 55 05 82   23 17 25 35 94   08 128
"""

STOPWORDED_COLLOCATIONS = """
thank you very much
thank you very much
thank you very much
thanks
"""

STOPWORDED_COLLOCATIONS_UPPERCASE = """
Thank you very much
Thank you very much
Thank you very much
thank you very much
hi There
Hi there
Hi There
thanks
"""

SMALL_CANVAS = """
better late than never someone will say
"""


def test_collocations():
    wc = WordCloud(collocations=False, stopwords=set())
    wc.generate(THIS)

    wc2 = WordCloud(collocations=True, stopwords=set())
    wc2.generate(THIS)

    assert "is better" in wc2.words_
    assert "is better" not in wc.words_
    assert "way may" not in wc2.words_


def test_collocation_stopwords():
    wc = WordCloud(collocations=True, stopwords={"you", "very"}, collocation_threshold=9)
    wc.generate(STOPWORDED_COLLOCATIONS)

    assert "thank you" not in wc.words_
    assert "very much" not in wc.words_
    assert "thank" in wc.words_
    # a bigram of all stopwords will be removed
    assert "you very" not in wc.words_


def test_collocation_stopwords_uppercase():
    wc = WordCloud(collocations=True, stopwords={"thank", "hi", "there"}, collocation_threshold=9)
    wc.generate(STOPWORDED_COLLOCATIONS_UPPERCASE)

    assert "Thank you" not in wc.words_
    assert "thank you" not in wc.words_
    assert "Thank" not in wc.words_
    # a bigram of all stopwords will be removed
    assert "hi There" not in wc.words_
    assert "Hi there" not in wc.words_
    assert "Hi There" not in wc.words_


def test_plurals_numbers():
    text = THIS + "\n" + "1 idea 2 ideas three ideas although many Ideas"
    wc = WordCloud(stopwords=[]).generate(text)
    # not capitalized usually
    assert "Ideas" not in wc.words_
    # plural removed
    assert "ideas" not in wc.words_
    # usually capitalized
    assert "although" not in wc.words_
    assert "idea" in wc.words_
    assert "Although" in wc.words_
    assert "better than" in wc.words_


def test_multiple_s():
    text = 'flo flos floss flosss'
    wc = WordCloud(stopwords=[]).generate(text)
    assert "flo" in wc.words_
    assert "flos" not in wc.words_
    assert "floss" in wc.words_
    assert "flosss" in wc.words_
    # not normalizing means that the one with just one s is kept
    wc = WordCloud(stopwords=[], normalize_plurals=False).generate(text)
    assert "flo" in wc.words_
    assert "flos" in wc.words_
    assert "floss" in wc.words_
    assert "flosss" in wc.words_


def test_empty_text():
    # test originally empty text raises an exception
    wc = WordCloud(stopwords=[])
    with pytest.raises(ValueError):
        wc.generate('')

    # test empty-after-filtering text raises an exception
    wc = WordCloud(stopwords=['a', 'b'])
    with pytest.raises(ValueError):
        wc.generate('a b a')


def test_default():
    # test that default word cloud creation and conversions work
    wc = WordCloud(max_words=50)
    wc.generate(THIS)

    # check for proper word extraction
    assert len(wc.words_) == wc.max_words

    # check that we got enough words
    assert len(wc.layout_) == wc.max_words

    # check image export
    wc_image = wc.to_image()
    assert wc_image.size == (wc.width, wc.height)

    # check that numpy conversion works
    wc_array = np.array(wc)
    assert_array_equal(wc_array, wc.to_array())

    # check size
    assert wc_array.shape == (wc.height, wc.width, 3)


def test_stopwords_lowercasing():
    # test that capitalized stopwords work.
    wc = WordCloud(stopwords=["Beautiful"])
    processed = wc.process_text(THIS)
    words = [count[0] for count in processed]
    assert "Beautiful" not in words


def test_writing_to_file(tmpdir):
    wc = WordCloud()
    wc.generate(THIS)

    # check writing to file
    filename = str(tmpdir.join("word_cloud.png"))
    wc.to_file(filename)
    loaded_image = Image.open(filename)
    assert loaded_image.size == (wc.width, wc.height)


def test_check_errors():
    wc = WordCloud()

    try:
        np.array(wc)
        raise AssertionError("np.array(wc) didn't raise")
    except ValueError as e:
        assert "call generate" in str(e)

    try:
        wc.recolor()
        raise AssertionError("wc.recolor didn't raise")
    except ValueError as e:
        assert "call generate" in str(e)


def test_svg_syntax():
    wc = WordCloud()
    wc.generate(THIS)
    svg = wc.to_svg()
    ET.fromstring(svg)


def test_recolor():
    wc = WordCloud(max_words=50, colormap="jet")
    wc.generate(THIS)
    array_before = wc.to_array()
    wc.recolor()
    array_after = wc.to_array()
    # check that the same places are filled
    assert_array_equal(array_before.sum(axis=-1) != 0,
                       array_after.sum(axis=-1) != 0)
    # check that they are not the same
    assert np.abs(array_before - array_after).sum() > 10000

    # check that recoloring is deterministic
    wc.recolor(random_state=10)
    wc_again = wc.to_array()
    assert_array_equal(wc_again, wc.recolor(random_state=10))


def test_random_state():
    # check that random state makes everything deterministic
    wc = WordCloud(random_state=0)
    wc2 = WordCloud(random_state=0)
    wc.generate(THIS)
    wc2.generate(THIS)
    assert_array_equal(wc, wc2)


def test_mask():
    # test masks

    # check that using an empty mask is equivalent to not using a mask
    wc = WordCloud(random_state=42)
    wc.generate(THIS)
    mask = np.zeros(np.array(wc).shape[:2], dtype=int)
    wc_mask = WordCloud(mask=mask, random_state=42)
    wc_mask.generate(THIS)
    assert_array_equal(wc, wc_mask)

    # use actual nonzero mask
    mask = np.zeros((234, 456), dtype=int)
    mask[100:150, 300:400] = 255

    wc = WordCloud(mask=mask)
    wc.generate(THIS)
    wc_array = np.array(wc)
    assert mask.shape == wc_array.shape[:2]
    assert_array_equal(wc_array[mask != 0], 0)
    assert wc_array[mask == 0].sum() > 10000


def test_mask_contour():
    # test mask contour is created, learn more at:
    # https://github.com/amueller/word_cloud/pull/348#issuecomment-370883873
    mask = np.zeros((234, 456), dtype=int)
    mask[100:150, 300:400] = 255

    sm = WordCloud(mask=mask, contour_width=1, contour_color='blue')
    sm.generate(THIS)
    sm_array = np.array(sm)
    sm_total = sm_array[100:150, 300:400].sum()

    lg = WordCloud(mask=mask, contour_width=20, contour_color='blue')
    lg.generate(THIS)
    lg_array = np.array(lg)
    lg_total = lg_array[100:150, 300:400].sum()

    sc = WordCloud(mask=mask, contour_width=1, scale=2, contour_color='blue')
    sc.generate(THIS)
    sc_array = np.array(sc)
    sc_total = sc_array[100:150, 300:400].sum()

    # test `contour_width`
    assert lg_total > sm_total

    # test contour varies with `scale`
    assert sc_total > sm_total

    # test `contour_color`
    assert all(sm_array[100, 300] == [0, 0, 255])


def test_single_color_func():
    # test single color function for different color formats
    random = Random(42)

    red_function = get_single_color_func('red')
    assert red_function(random_state=random) == 'rgb(181, 0, 0)'

    hex_function = get_single_color_func('#00b4d2')
    assert hex_function(random_state=random) == 'rgb(0, 48, 56)'

    rgb_function = get_single_color_func('rgb(0,255,0)')
    assert rgb_function(random_state=random) == 'rgb(0, 107, 0)'

    rgb_perc_fun = get_single_color_func('rgb(80%,60%,40%)')
    assert rgb_perc_fun(random_state=random) == 'rgb(97, 72, 48)'

    hsl_function = get_single_color_func('hsl(0,100%,50%)')
    assert hsl_function(random_state=random) == 'rgb(201, 0, 0)'


def test_single_color_func_grey():
    # grey is special as it's a corner case
    random = Random(42)

    red_function = get_single_color_func('darkgrey')
    assert red_function(random_state=random) == 'rgb(181, 181, 181)'
    assert red_function(random_state=random) == 'rgb(56, 56, 56)'


def test_process_text():
    # test that process function returns a dict
    wc = WordCloud(max_words=50)
    result = wc.process_text(THIS)

    # check for proper return type
    assert isinstance(result, dict)


def test_process_text_default_patterns():
    wc = WordCloud(stopwords=set(), include_numbers=True, min_word_length=2)
    words = wc.process_text(THIS)

    wc2 = WordCloud(stopwords=set(), include_numbers=True, min_word_length=1)
    words2 = wc2.process_text(THIS)

    assert "a" not in words
    assert "3" not in words

    assert "a" in words2
    assert "3" in words2


def test_process_text_regexp_parameter():
    # test that word processing is influenced by `regexp`
    wc = WordCloud(max_words=50, regexp=r'\w{5}')
    words = wc.process_text(THIS)

    assert 'than' not in words


def test_generate_from_frequencies():
    # test that generate_from_frequencies() takes input argument dicts
    wc = WordCloud(max_words=50)
    words = wc.process_text(THIS)
    result = wc.generate_from_frequencies(words)

    assert isinstance(result, WordCloud)


def test_relative_scaling_zero():
    # non-regression test for non-integer font size
    wc = WordCloud(relative_scaling=0)
    wc.generate(THIS)


def test_unicode_stopwords():
    wc_unicode = WordCloud(stopwords=[u'Beautiful'])
    try:
        words_unicode = wc_unicode.process_text(unicode(THIS))
    except NameError:  # PY3
        words_unicode = wc_unicode.process_text(THIS)

    wc_str = WordCloud(stopwords=['Beautiful'])
    words_str = wc_str.process_text(str(THIS))

    assert words_unicode == words_str


def test_include_numbers():
    wc_numbers = WordCloud(include_numbers=True)
    wc = wc_numbers.process_text(THIS)

    assert '14' in wc.keys()


def test_min_word_length():
    wc_numbers = WordCloud(min_word_length=5)
    wc = wc_numbers.process_text(THIS)
    word_lengths = [len(word) for word in wc.keys()]

    assert min(word_lengths) == 5


def test_recolor_too_small():
    # check exception is raised when image is too small
    colouring = np.array(Image.new('RGB', size=(20, 20)))
    wc = WordCloud(width=30, height=30, random_state=0, min_font_size=1).generate(THIS)
    image_colors = ImageColorGenerator(colouring)
    with pytest.raises(ValueError, match='ImageColorGenerator is smaller than the canvas'):
        wc.recolor(color_func=image_colors)


def test_recolor_too_small_set_default():
    # check no exception is raised when default colour is used
    colouring = np.array(Image.new('RGB', size=(20, 20)))
    wc = WordCloud(max_words=50, width=30, height=30, min_font_size=1).generate(THIS)
    image_colors = ImageColorGenerator(colouring, default_color=(0, 0, 0))
    wc.recolor(color_func=image_colors)


def test_small_canvas():
    # check font size fallback works on small canvas
    wc = WordCloud(max_words=50, width=21, height=21)
    wc.generate(SMALL_CANVAS)
    assert len(wc.layout_) > 0


def test_tiny_canvas():
    # check exception if canvas too small for fallback
    w = WordCloud(max_words=50, width=1, height=1)
    with pytest.raises(ValueError, match="Couldn't find space to draw"):
        w.generate(THIS)
    assert len(w.layout_) == 0


def test_coloring_black_works():
    # check that using black colors works.
    mask = np.zeros((50, 50, 3))
    image_colors = ImageColorGenerator(mask)
    wc = WordCloud(width=50, height=50, random_state=42,
                   color_func=image_colors, min_font_size=1)
    wc.generate(THIS)


def test_repeat():
    short_text = "Some short text"
    wc = WordCloud(stopwords=[]).generate(short_text)
    assert len(wc.layout_) == 3
    wc = WordCloud(max_words=50, stopwords=[], repeat=True).generate(short_text)
    # multiple of word count larger than max_words
    assert len(wc.layout_) == 51
    # relative scaling doesn't work well with repeat
    assert wc.relative_scaling == 0
    # all frequencies are 1
    assert len(wc.words_) == 3
    assert_array_equal(list(wc.words_.values()), 1)
    frequencies = [w[0][1] for w in wc.layout_]
    assert_array_equal(frequencies, 1)
    repetition_text = "Some short text with text"
    wc = WordCloud(max_words=52, stopwords=[], repeat=True)
    wc.generate(repetition_text)
    assert len(wc.words_) == 4
    # normalized frequencies
    assert wc.words_['text'] == 1
    assert wc.words_['with'] == .5
    assert len(wc.layout_), wc.max_words
    frequencies = [w[0][1] for w in wc.layout_]
    # check that frequencies are sorted
    assert np.all(np.diff(frequencies) <= 0)


def test_zero_frequencies():

    word_cloud = WordCloud()

    word_cloud.generate_from_frequencies({'test': 1, 'test1': 0, 'test2': 0})
    assert len(word_cloud.layout_) == 1
    assert word_cloud.layout_[0][0][0] == 'test'


def test_plural_stopwords():
    x = '''was was was was was was was was was was was was was was was
    wa
    hello hello hello hello hello hello hello hello
    goodbye good bye maybe yes no'''
    w = WordCloud().generate(x)
    assert w.words_['wa'] < 1

    w = WordCloud(collocations=False).generate(x)
    assert w.words_['wa'] < 1


def test_max_font_size_as_mask_height():
    # test if max font size will respect the mask height
    x = '''hello hello hello
    bye'''

    # Get default wordcloud size
    wcd = WordCloud()
    default_size = (wcd.height, wcd.width)
    # Make sure the size we are using is larger than the default size
    size = (default_size[0] * 2, default_size[1] * 2)

    # using mask, all drawable
    mask = np.zeros(size, dtype=int)
    mask[:, :] = 0
    wc = WordCloud(mask=mask, random_state=42)
    wc.generate(x)

    # no mask
    wc2 = WordCloud(width=size[1], height=size[0], random_state=42)
    wc2.generate(x)

    # Check if the biggest element has the same font size
    assert wc.layout_[0][1] == wc2.layout_[0][1]