import os
import sys
import codecs
import argparse
from rouge import Rouge

################# main #################

def get_rouge(hypothesis, references):
    rouge = Rouge()
    preprocess_exs = lambda exs : [ex.strip().lower() for ex in exs]
    rouge_scores =  rouge.get_scores(preprocess_exs(hypothesis), preprocess_exs(references), avg=True)
    return {k: v for k, v in rouge_scores.items()}

def rouge_results_to_str(results_dict):
    return "ROUGE-F(1/2/l): {:.2f}/{:.2f}/{:.2f}\nROUGE-R(1/2/l): {:.2f}/{:.2f}/{:.2f}\nROUGE-P(1/2/l): {:.2f}/{:.2f}/{:.2f}\n".format(
        results_dict["rouge-1"]['f'] * 100,
        results_dict["rouge-2"]['f'] * 100,
        results_dict["rouge-l"]['f'] * 100,
        results_dict["rouge-1"]['r'] * 100,
        results_dict["rouge-2"]['r'] * 100,
        results_dict["rouge-l"]['r'] * 100,
        results_dict["rouge-1"]['p'] * 100,
        results_dict["rouge-2"]['p'] * 100,
        results_dict["rouge-l"]['p'] * 100,
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', type=str, default="candidate.txt", help='candidate file')
    parser.add_argument('-r', type=str, default="reference.txt", help='reference file')
    args = parser.parse_args()
    print(args)
    candidates = codecs.open(args.c, encoding="utf-8")
    references = codecs.open(args.r, encoding="utf-8")

    references = [line.strip().lower().replace('<q>', '') for line in references]
    candidates = [line.strip().lower().replace('<q>', '') for line in candidates]

    results_dict = get_rouge(candidates, references)

    print(rouge_results_to_str(results_dict))