import nltk
from nltk.corpus import stopwords
import matplotlib.pyplot as plt

LINES = ['-', ':', '--']  # Style linii do wykresów


def main():
    # Pobiera pliki tekstowe do słownika z autorem jako kluczem.
    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')

    # Sprawdza wynik wgrywania plików.
    print(strings_by_author['doyle'][:300])

    # Przekształca tekst i przeprowadza testy stylometryczne.
    words_by_author = make_word_dict(strings_by_author)
    len_shortest_corpus = find_shortest_corpus(words_by_author)

    word_length_test(words_by_author, len_shortest_corpus)
    stopwords_test(words_by_author, len_shortest_corpus)
    parts_of_speech_test(words_by_author, len_shortest_corpus)
    vocab_test(words_by_author)
    jaccard_test(words_by_author, len_shortest_corpus)


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


def make_word_dict(strings_by_author):
    """
    Zwraca słownik zawierający listy tokenów w postaci słów
    przypisane do odpowiedniego autora.
    """
    words_by_author = dict()
    for author in strings_by_author:
        tokens = nltk.word_tokenize(strings_by_author[author])
        words_by_author[author] = ([token.lower() for token in tokens
                                    if token.isalpha()])
    return words_by_author


def find_shortest_corpus(words_by_author):
    """Zwraca długość najkrótszego korpusu."""
    word_count = []
    for author in words_by_author:
        word_count.append(len(words_by_author[author]))
        print('\nLiczba słów dla klucza {} = {}'
              .format(author, len(words_by_author[author])))
    len_shortest_corpus = min(word_count)
    print('\nDługość najkrótszego korpusu = {}\n'.format(len_shortest_corpus))
    return len_shortest_corpus


def word_length_test(words_by_author, len_shortest_corpus):
    """
    Tworzy wykres przedstawiający częstość długości słów dla autora
    z ograniczeniem danych do długości najkrótszego korpusu.
    """
    by_author_length_freq_dist = dict()
    plt.figure(1)
    plt.ion()

    for i, author in enumerate(words_by_author):
        word_lengths = [len(word) for word in words_by_author[author]
                        [:len_shortest_corpus]]
        by_author_length_freq_dist[author] = nltk.FreqDist(word_lengths)
        by_author_length_freq_dist[author].plot(
            15,
            linestyle=LINES[i],
            label=author,
            title='Częstość występowania słów o różnej długości')

    plt.legend()
    plt.ylabel("Liczba wystąpień")
    plt.xlabel("Długość słowa")
    # Odkomentuj poniższą linię, aby zobaczyć wykres na etapie kodowania.
    # plt.show()


def stopwords_test(words_by_author, len_shortest_corpus):
    """
    Tworzy wykres częstości występowania słów nieindeksowanych
    z ograniczeniem danych do długości najkrótszego korpusu.
    """
    stopwords_by_author_freq_dist = dict()
    plt.figure(2)
    # Używa zbioru, aby przyspieszyć przetwarzanie.
    stop_words = set(stopwords.words('english'))
    # print('Liczba słów nieindeksowanych = {}\n'.format(len(stop_words)))
    # print('Słowa nieindeksowane = {}\n'.format(stop_words))
    
    for i, author in enumerate(words_by_author):
        stopwords_by_author = [word for word in words_by_author[author]
                               [:len_shortest_corpus] if word in stop_words]
        stopwords_by_author_freq_dist[author] = nltk.FreqDist(
            stopwords_by_author)
        stopwords_by_author_freq_dist[author].plot(
            50,
            label=author,
            linestyle=LINES[i],
            title='50 najczęściej użytych słów nieindeksowanych')

    plt.legend()
    plt.ylabel("Liczba wystąpień")
    plt.xlabel("Słowo")
    # Odkomentuj poniższą linię, aby zobaczyć wykres na etapie kodowania.
    # plt.show()


def parts_of_speech_test(words_by_author, len_shortest_corpus):
    """
    Tworzy wykres użycia przez autora części mowy,
    takich jak rzeczowniki, czasowniki czy przysłówki.
    """
    by_author_pos_freq_dist = dict()
    plt.figure(3)
    for i, author in enumerate(words_by_author):
        pos_by_author = [pos[1] for pos in nltk.pos_tag(words_by_author[author]
                                                        [:len_shortest_corpus])]
        by_author_pos_freq_dist[author] = nltk.FreqDist(pos_by_author)
        by_author_pos_freq_dist[author].plot(35,
                                             label=author,
                                             linestyle=LINES[i],
                                             title='Części mowy')
    plt.legend()
    plt.ylabel("Liczba wystąpień")
    plt.xlabel("Część mowy")
    plt.show()
    # W przypadku użytkowników Windows PowerShell może zaistnieć
    # potrzeba użycia polecenia: plt.show(block=True).


def vocab_test(words_by_author):
    """
    Porównuje słownictwo używane przez autorów,
    korzystając z testu zgodności chi-kwadrat.
    """
    chisquared_by_author = dict()
    for author in words_by_author:
        if author != 'unknown':
            # Łączy korpus danego autora i sporny korpus,
            # a następnie znajduje 1000 najczęściej używanych słów.
            combined_corpus = (words_by_author[author] +
                               words_by_author['unknown'])
            author_proportion = (len(words_by_author[author]) /
                                 len(combined_corpus))
            combined_freq_dist = nltk.FreqDist(combined_corpus)
            most_common_words = list(combined_freq_dist.most_common(1000))
            chisquared = 0

            # Oblicza stosunek zaobserwowanej do oczekiwanej liczby wystąpień.
            for word, combined_count in most_common_words:
                observed_count_author = words_by_author[author].count(word)
                expected_count_author = combined_count * author_proportion
                chisquared += ((observed_count_author -
                                expected_count_author) ** 2 /
                               expected_count_author)
                chisquared_by_author[author] = chisquared
            print('Chi-kwadrat dla klucza {} = {:.1f}'
                  .format(author, chisquared))

    most_likely_author = min(chisquared_by_author,
                             key=chisquared_by_author.get)
    print('Biorąc pod uwagę słownictwo, autorem najprawdopodobniej '
          'jest: {}.\n'.format(most_likely_author))


def jaccard_test(words_by_author, len_shortest_corpus):
    """
    Oblicza współczynnik podobieństwa Jaccarda między spornym
    korpusem a korpusami o ustalonym autorstwie.
    """
    jaccard_by_author = dict()
    unique_words_unknown = set(words_by_author['unknown']
                               [:len_shortest_corpus])
    authors = (author for author in words_by_author if author != 'unknown')
    for author in authors:
        unique_words_author = set(words_by_author[author]
                                  [:len_shortest_corpus])
        shared_words = unique_words_author.intersection(unique_words_unknown)
        jaccard_sim = (float(len(shared_words)) / (len(unique_words_author) +
                                                   len(unique_words_unknown) -
                                                   len(shared_words)))
        jaccard_by_author[author] = jaccard_sim
        print('Indeks Jaccarda dla klucza {} = {}'.format(author, jaccard_sim))

    most_likely_author = max(jaccard_by_author, key=jaccard_by_author.get)
    print('Biorąc pod uwagę podobieństwo, autorem najprawdopodobniej '
          'jest: {}'.format(most_likely_author))


if __name__ == '__main__':
    main()
