import argparse

import gold_timeline
import dataset
import date_selection
import sentence_selection
import summarization
import evaluation
import util

topics = {
    'bpoil': {
        'keywords': ['bp', 'oil', 'spill'],
        'gold_timelines': [
            'data/in/timeline17/Data/bpoil_bbc/timelines/bbc.txt',
            'data/in/timeline17/Data/bpoil_foxnews/timelines/foxnews.txt',
            'data/in/timeline17/Data/bpoil_guardian/timelines/guardian.txt',
            'data/in/timeline17/Data/bpoil_reuters/timelines/rbpoil.txt',
            'data/in/timeline17/Data/bpoil_washingtonpost/timelines/washingtonpost.txt',
        ],
    },
    'EgyptianProtest': {
        'keywords': ['egypt', 'egyptian'],
        'gold_timelines': [
            'data/in/timeline17/Data/EgyptianProtest_cnn/timelines/tl8.txt'
        ],
    },
    'Finan': {
        'keywords': ['financial', 'economic', 'crisis'],
        'gold_timelines': [
            'data/in/timeline17/Data/Finan_washingtonpost/timelines/financial_crisis_wallstreet_washingtonpost.txt'
        ],
    },
    'H1N1': {
        'keywords': ['h1n1', 'swine', 'flu'],
        'gold_timelines': [
            'data/in/timeline17/Data/H1N1_bbc/timelines/influenza_A_h1n1_scotland_bbc.txt',
            'data/in/timeline17/Data/H1N1_guardian/timelines/swine_flu_guardian.txt',
            'data/in/timeline17/Data/H1N1_reuters/timelines/h1n1_reuters.txt',
        ],
    },
    'haiti': {
        'keywords': ['haiti', 'quake', 'earthquake'],
        'gold_timelines': ['data/in/timeline17/Data/haiti_bbc/timelines/bbc.txt'],
    },
    'IraqWar': {
        'keywords': ['iraq', 'iraqi'],
        'gold_timelines': [
            'data/in/timeline17/Data/IraqWar_guardian/timelines/tl7.txt'
        ],
    },
    'LibyaWar': {
        'keywords': ['libya', 'libyan'],
        'gold_timelines': [
            'data/in/timeline17/Data/LibyaWar_cnn/timelines/tl6.txt',
            'data/in/timeline17/Data/LibyaWar_reuters/timelines/tl5.txt',
        ],
    },
    'MJ': {
        'keywords': ['michael', 'jackson'],
        'gold_timelines': ['data/in/timeline17/Data/MJ_bbc/timelines/bbc.co.uk.txt'],
    },
    'SyrianCrisis': {
        'keywords': ['syria', 'syrian'],
        'gold_timelines': [
            'data/in/timeline17/Data/SyrianCrisis_bbc/timelines/tl1.txt',
            'data/in/timeline17/Data/SyrianCrisis_reuters/timelines/tl2.1.txt',
            'data/in/timeline17/Data/SyrianCrisis_reuters/timelines/tl2.2.txt',
            'data/in/timeline17/Data/SyrianCrisis_reuters/timelines/tl2.3.txt',
        ],
    },
}


def main(arguments):
    def print_if_not_csv(fargs, *args, **kwargs):
        if not arguments.output_csv:
            print(fargs, *args, **kwargs)

    print_if_not_csv('Loading dataset')
    timeline17_dataset = dataset.get_timeline17_dataset('data/in/timeline17/Data')

    ar1_f_scores = []
    ar2_f_scores = []
    date_f1_scores = []

    if arguments.action == 'evaluate' and arguments.output_csv:
        print('topic;gold_timeline;AR1-F;AR2-F;Date-F1')

    for topic in ['bpoil']:  # topics.keys():

        keywords = topics[topic]['keywords']

        print_if_not_csv('Loading articles')
        articles = timeline17_dataset[topic]['articles']

        print_if_not_csv('Loading gold timelines')
        gold_tls = timeline17_dataset[topic]['gold_timelines'].values()

        for gold_tl_filename in timeline17_dataset[topic]['gold_timelines'].keys():

            print_if_not_csv(f'Topic: {topic}; gold timeline: {gold_tl_filename}')

            gold_tl = timeline17_dataset[topic]['gold_timelines'][gold_tl_filename]

            # print_if_not_csv('    Filtering articles by date')
            dates_in_gold_tl = sorted(gold_tl.keys())
            start_date = dates_in_gold_tl[0]  # util.first_day_of_month(dates_in_gold_tl[0])
            end_date = dates_in_gold_tl[-1]  # util.last_day_of_month(dates_in_gold_tl[-1])
            # articles = dataset.filter_articles_by_date(articles, start_date, end_date)

            print_if_not_csv('    Selecting dates')
            num_dates = len(gold_tl)
            avg_num_sentences = sum([len(gold_tl[date]) for date in gold_tl]) / len(gold_tl)
            selected_dates = date_selection.select_dates_by_mention_count(articles, num_dates, start_date, end_date)

            print_if_not_csv('    Selecting candidate sentences and building summary')
            summary = {}
            num_sentences = int(avg_num_sentences)
            for date in selected_dates:
                candidate_sentences, date_vector = sentence_selection.candidate_sentences(articles, date, keywords)
                if len(candidate_sentences) == 0:
                    continue
                summary_for_date = summarization.summarize(candidate_sentences, num_sentences)
                summary[date] = summary_for_date

            if arguments.action == 'run':
                for date, sentences in summary.items():
                    print()
                    print('        ' + date)
                    for s in sentences:
                        print('            ' + s)
                print()
                print()
                print()

            if arguments.action == 'evaluate':
                print_if_not_csv('    Evaluating')
                eval_result = evaluation.evaluate(summary, gold_tls)
                ar1_f = eval_result["rouge_1"]["f_score"]
                ar2_f = eval_result["rouge_2"]["f_score"]
                date_f1 = evaluation.date_f1(summary, gold_tl)

                ar1_f_scores.append(ar1_f)
                ar2_f_scores.append(ar2_f)
                date_f1_scores.append(date_f1)

                if arguments.output_csv:
                    print(f'{topic};{gold_tl_filename};{ar1_f};{ar2_f};{date_f1}')
                else:
                    print(f'    AR1-F: {ar1_f}; AR2-F: {ar2_f}; Date-F1: {date_f1}')

    if arguments.action == 'evaluate' and arguments.output_csv:
        avg_ar1_f = sum(ar1_f_scores) / len(ar1_f_scores)
        avg_ar2_f = sum(ar2_f_scores) / len(ar2_f_scores)
        avg_date_f1 = sum(date_f1_scores) / len(date_f1_scores)
        print(f'average;;{avg_ar1_f};{avg_ar2_f};{avg_date_f1}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--action', required=True, choices=['run', 'evaluate'], default='run')
    parser.add_argument('--output_csv', help='output results in a CSV format (only applies if --action is "evaluate")',
                        action='store_true')
    main(parser.parse_args())
