import bs4 import time import codecs from nltk import sent_tokenize wo

import bs4
import time
import codecs
from nltk import sent_tokenize, wordpunct_tokenize, pos_tag
from nltk.corpus.reader.api import CorpusReader, CategorizedCorpusReader
from readability.readability import Unparseable, Document

CAT_PATTERN = r'([a-z_\s]+).*'
DOC_PATTERN = r'(?!\.)[a-z_\s]+/[a-z0-9]+\.json'
TAGS = ['h1', 'h2', 'h3', 'h4', 'h5', 'h6' ,'h7' ,'p' , 'li', 'div']

class HTMLCorpusReader(CategorizedCorpusReader, CorpusReader):
    """Объект чтения корпуса с HTML-документами."""
    def __init__(self, root, fileids = DOC_PATTERN, encoding = 'utf8', tags = TAGS, **kwargs):
        """Инициализирует объект чтения корпуса. Аргументы, управляющие классификацией (`cat_pattern`, `cat_map` и `cat_file`),
        передаются в конструктор `CategorizedCorpusReader`. Остальные аргументы передаются в конструктор `CorpusReader`."""
        # добавить шаблон категорий, если он не был передан в класс явно
        if not any(key.startswith('cat_') for key in kwargs.keys()):
            kwargs['cat_pattern'] = CAT_PATTERN

        # инициализировать объекты чтения корпуса nltk
        CategorizedCorpusReader.__init__(self, kwargs)
        CorpusReader.__init__(self, root, fileids, encoding)

        # сохранить теги, подлежащие извлечению
        self.tags = tags

    def resolve(self, fileids = None, categories = None):
        """Возвращает список идендификаторов файлов или названий категорий, которые передаются каждой внутренней функции
        объекта чтения корпуса. Реализована по аналогии с `CategorizedPlaintextCorpusReader` в NLTK."""
        # вызвать ошибку, если переданы оба параметра
        if fileids is not None and categories is not None:
            raise ValueError("Specify fileids or categories, not both")

        # вернуть идентификаторы файлов, ассоциированные с переданными категориями
        if categories is not None:
            return self.fileids(categories)

        # вернуть переданные идентификаторы файлов
        return fileids

    def sizes(self, fileids = None, categories = None):
        """Возращает список кортежей, полное имя и размер файла."""
        # получить список файлов
        fileids = self.resolve(fileids, categories)

        # создать генератор, полные имена и размеры файлов
        for path in self.abspaths(fileids):
            yield path, os.path.getsize(path)

    def docs(self, fileids = None, categories = None):
        """Возвращает полный текст HTML-документа."""
        # получить список файлов для чтения
        fileids = self.resolve(fileids, categories)

        # создать генератор, загружающий документы в память по одному
        for path, encoding in self.abspaths(fileids, include_encoding = True):
            with codecs.open(path, 'r', encoding = encoding) as f:
                yield f.read()

    def html(self, fileids = None, categories = None):
        """Возвращает содержимое HTML каждого документа, очищая его с помощью библиотеки readability-lxml."""
        for doc in self.docs(fileids, categories):
            try:
                yield Document(doc).summary()
            except Unparseable as e:
                print("Could not parse HTML: {}".format(e))
                continue

    def paragraphs(self, fileids = None, categories = None):
        """Использует BeautifulSoup для выделения абзацев из HTML."""
        for html in self.html(fileids, categories):
            soup = bs4.BeautifulSoup(html, 'lxml')
            for element in soup.find_all(self.tags):
                yield element.text
            soup.decompose()

    def sentences(self, fileids = None, categories = None):
        """Использует встроенный механизм NLTK для выделения предложений из абзацев."""
        for paragraph in self.paragraphs(fileids, categories):
            for sentence in sent_tokenize(paragraph):
                yield sentence

    def words(self, fileids = None, categories = None):
        """Использует встроенный механизм NLTK для выделения слов из предложений."""
        for sentence in self.sentences(fileids, categories):
            for word in wordpunct_tokenize(sentence):
                yield word

    def tokenize(self, fileids = None, categories = None):
        """Сегментирует, лексимизирует и маркирует документ в корпусе."""
        for paragraph in self.paragraphs(fileids, categories):
            yield [
                pos_tag(wordpunct_tokenize(sentence))
                for sentence in sent_tokenize(paragraph)
            ]

    def describe(self, fileids = None, categories = None):
        """Выполняет обход копруса и возвращает словарь с его характеристиками."""
        # время начала
        start = time.time()

        # структуры для подсчета
        counts = nltk.FreqDist()
        words = nltk.FreqDist()

        # выполнить обход корпуса, выделить лексемы и подсчитать их
        for paragraph in self.paragraphs(fileids, categories):
            counts['paragraphs'] += 1
            for sentence in paragraph:
                counts['sentences'] += 1
                for word, tag in sentence:
                    counts['words'] += 1
                    words[word] += 1

        # определить число файлов и категорий в корпусе
        n_fileids = len(self.resolve(fileids, categories) or self.fileids)
        n_categories = len(self.categories(self.resolve(fileids, categories)))

        return {
            'files':           n_fileids,
            'categories':      n_categories,
            'paragraphs':      counts['paragraphs'],
            'sentences':       counts['sentences'],
            'words':           counts['words'],
            'vocabulary':      len(words),
            'avg_par_in_doc':  float(counts['paragraphs']) / float(n_fileids),
            'avg_sent_in_par': float(counts['sentences']) / float(counts['paragraphs']),
            'lex_div':         float(counts['words']) / float(len(tokens)),
            'secs':            time.time() - start
        }