
import os
import re
from collections import namedtuple
from random import seed, sample

from tqdm import tqdm as log_progress

from natasha import NamesExtractor
from natasha.markup import show_markup
from natasha.grammars.name import Name


FACTRU_DIR = 'factRuEval-2016'
DEVSET_DIR = os.path.join(FACTRU_DIR, 'devset')
TESTSET_DIR = os.path.join(FACTRU_DIR, 'testset')
RESULTS_DIR = os.path.join(FACTRU_DIR, 'results')

TAGS = {
    Name: 'PER'
}


Factru = namedtuple(
    'Factru',
    ['id', 'text']
)
Task1Record = namedtuple(
    'Task1Record',
    ['tag', 'start', 'size']
)


def load_text(path):
    with open(path) as file:
        return file.read()


def load_factru(dir):
    for filename in os.listdir(dir):
        match = re.match('book_(\d+)\.txt', filename)
        if match:
            id = int(match.group(1))
            path = os.path.join(dir, filename)
            text = load_text(path)
            yield Factru(id, text)


def get_task1_path(id):
    filename = 'book_{id}.task1'.format(id=id)
    return os.path.join(
        RESULTS_DIR,
        filename
    )


def format_task1(result):
    for record in result:
        yield '{record.tag} {record.start} {record.size}'.format(
            record=record
        )


def dump_lines(lines, path):
    with open(path, 'w') as file:
        for line in lines:
            file.write(line + '\n')


def dump_task1(id, result):
    path = get_task1_path(id)
    lines = format_task1(result)
    dump_lines(lines, path)


def match_to_task1(match):
    start, stop = match.span
    size = stop - start
    fact = match.fact
    tag = TAGS[type(fact)]
    return Task1Record(tag, start, size)


def maybe_mkdir(dir):
    if not os.path.exists(dir):
        os.mkdir(dir)


if __name__ == '__main__':
    factru = list(load_factru(TESTSET_DIR))
    extractor = NamesExtractor()
    maybe_mkdir(RESULTS_DIR)
    for record in log_progress(factru):
        matches = extractor(record.text)
        guess = [match_to_task1(_) for _ in matches]
        dump_task1(record.id, guess)
