"""
Functions to parse and preprocess the datasets.

The output of each function is a dict with the following structure:
{
    topic_name: {
        articles: [{
            pub_date,
            sentences: [{
                text,
                mentioned_dates
            }]
        }],
        gold_timelines: {
            gold_timeline_name: {
                date: [sentences]
            }
        }
    }
}
"""

import itertools
import os
import pickle
from datetime import datetime
from pathlib import Path
import heideltime_util
import util
import re


def get_timeline17_dataset(path):
    """
    Returns the Timeline17 dataset as a dictionary.
    If cache/timeline17-dataset.pkl exists, it will be loaded from there,
    otherwise, it will be parsed from the given path (assuming the default folder structure).

    :param path: The path to Timeline17's 'Data' directory.
    :return: A dictionary containing the preprocessed data.
    """

    cache_filename = 'cache/ghalandari.pkl'
    if os.path.exists(cache_filename):
        return pickle.load(open(cache_filename, 'rb'))

    path = Path(path)
    data = {}

    # go through each topic directory
    for topic_dirname in util.subdirs(path):
        topic_path = path / topic_dirname
        topic_name, source = topic_dirname.split('_')
        if topic_name not in data:
            data[topic_name] = {'articles': [], 'gold_timelines': {}}

        # parse input articles
        for pub_date in util.subdirs(topic_path / 'InputDocs'):
            date_path = topic_path / 'InputDocs' / pub_date
            for article_filename in util.files(date_path, extension='txt'):
                article_file_path = date_path / article_filename
                print(article_file_path)
                article = {'pub_date': pub_date, 'sentences': []}

                # get sentence text
                with open(article_file_path) as f:
                    sentences_in_article = [{
                        'text': line.strip(),
                        'mentioned_dates': []
                    } for line in f.readlines() if line.strip()]

                # get date mentions using HeidelTime
                # and add them to the sentence data
                mentioned_dates_by_sentence = heideltime_util.mentioned_dates_by_sentence(article_file_path, pub_date)
                assert len(mentioned_dates_by_sentence) == len(sentences_in_article)
                for i in range(len(sentences_in_article)):
                    sentence = sentences_in_article[i]
                    sentence['mentioned_dates'] = mentioned_dates_by_sentence[i]

                article['sentences'] += sentences_in_article
                data[topic_name]['articles'].append(article)

        # parse gold timelines
        for gold_timeline_filename in util.files(topic_path / 'timelines', extension='txt'):
            gold_timeline_file_path = topic_path / 'timelines' / gold_timeline_filename
            gold_timeline_name = gold_timeline_filename.split('.')[0]
            data[topic_name]['gold_timelines'][gold_timeline_name] = {}
            with open(gold_timeline_file_path) as f:
                lines = [line.strip() for line in f.readlines()]
                date_groups = [list(y) for x, y in itertools.groupby(lines, lambda z: re.match('^-+$', z)) if not x]
                for date_group in date_groups:
                    date, sentences_on_date = date_group[0], date_group[1:]
                    sentences_on_date = [s.lstrip('-').strip() for s in sentences_on_date]
                    data[topic_name]['gold_timelines'][gold_timeline_name][date] = sentences_on_date

    pickle.dump(data, open(cache_filename, 'wb'))
    return data


def filter_articles_by_keywords(articles, keywords):
    result = []
    for article in articles:
        filtered_sentences = []
        for sentence in article['sentences']:
            for keyword in keywords:
                if keyword.lower() in sentence['text'].lower():
                    filtered_sentences.append(sentence)
                    break
        filtered_article = {'pub_date': article['pub_date'], 'sentences': filtered_sentences}
        result.append(filtered_article)

    return result


def filter_articles_by_date(articles, start_date_str, end_date_str):
    result = []
    for article in articles:
        pub_date = datetime.strptime(article['pub_date'], '%Y-%m-%d')
        start_date = datetime.strptime(start_date_str, '%Y-%m-%d')
        end_date = datetime.strptime(end_date_str, '%Y-%m-%d')
        if start_date <= pub_date <= end_date:
            result.append(article)
    return result
