"""Tworzy interpunkcyjną mapę cieplną."""
import math
from string import punctuation
import nltk
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
# Zainstaluj pakiet seaborn następującym poleceniem: pip install seaborn.
import seaborn as sns

PUNCT_SET = set(punctuation)


def main():
    # Pobiera pliki tekstowe do słownika (autor jako klucz).
    strings_by_author = dict()
    strings_by_author['doyle'] = text_to_string('hound.txt')
    strings_by_author['wells'] = text_to_string('war.txt')
    strings_by_author['unknown'] = text_to_string('lost.txt')

    # Przekształca teksty w tokeny, zostawiając jedynie
    # znaki interpunkcyjne.
    punct_by_author = make_punct_dict(strings_by_author)

    # Przekształca znaki interpunkcyjne w wartości numeryczne
    # i tworzy mapy cieplne.
    plt.ion()
    for author in punct_by_author:
        heat = convert_punct_to_number(punct_by_author, author)
        # Przycina do największego wymiaru dla tablicy dwuwymiarowej.
        arr = np.array((heat[:6561]))
        arr_reshaped = arr.reshape(int(math.sqrt(len(arr))),
                                   int(math.sqrt(len(arr))))
        fig, ax = plt.subplots(figsize=(7, 7))
        sns.heatmap(arr_reshaped,
                    cmap=ListedColormap(['blue', 'yellow']),
                    square=True,
                    ax=ax)
        ax.set_title('Mapa średników dla klucza {}'.format(author))
    plt.show()


def text_to_string(filename):
    """Odczytuje plik tekstowy i zwraca ciąg znaków."""
    with open(filename) as infile:
        return infile.read()


def make_punct_dict(strings_by_author):
    """
    Zwraca słownik autorów i tokenów ze znakami interpunkcyjnymi
    z korpusów tych autorów.
    """
    punct_by_author = dict()
    for author in strings_by_author:
        tokens = nltk.word_tokenize(strings_by_author[author])
        punct_by_author[author] = ([token for token in tokens
                                    if token in PUNCT_SET])
        print("Liczba znaków interpunkcyjnych dla klucza {} = {}"
              .format(author, len(punct_by_author[author])))
    return punct_by_author


def convert_punct_to_number(punct_by_author, author):
    """
    Zwraca listę znaków interpunkcyjnych przekształconych
    na wartości liczbowe.
    """
    heat_vals = []
    for char in punct_by_author[author]:
        if char == ';':
            value = 1
        else:
            value = 2
        heat_vals.append(value)
    return heat_vals


if __name__ == '__main__':
    main()
