import argparse
import os
import time
import numpy as np
from utils import *
from regression import *
import evaluate

SRC_HYPO = read_file_to_list('files/src_hypo_prompt.txt')
REF_HYPO = read_file_to_list('files/ref_hypo_prompt.txt')


class Scorer:
    """ Support ROUGE-1,2,L, BERTScore, MoverScore, PRISM, BARTScore """

    def __init__(self, file_path, device='cuda:0'):
        """ file_path: path to the pickle file
            All the data are normal capitalized, and tokenized, including ref_summs, and sys_summ.
        """
        self.device = device
        self.data = read_pickle(file_path)
        print(f'Data loaded from {file_path}.')

    def save_data(self, path):
        save_pickle(self.data, path)

    def score(self, metrics):
        """ metrics: list of metrics """
        for metric_name in metrics:
            if metric_name == 'bert_score':
                from bert_score import BERTScorer

                # Set up BERTScore
                bert_scorer = BERTScorer(
                    lang='en',
                    idf=False,
                    rescale_with_baseline=True,
                    device=self.device
                )
                print(f'BERTScore setup finished. Begin calculating BERTScore.')

                start = time.time()
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    sys_summ = self.data[doc_id]['sys_summ']
                    P, R, F = bert_scorer.score([sys_summ] * len(ref_summs), ref_summs)
                    P = P.max().item()
                    R = R.max().item()
                    F = F.max().item()
                    self.data[doc_id]['scores']['bert_score_p'] = P
                    self.data[doc_id]['scores']['bert_score_r'] = R
                    self.data[doc_id]['scores']['bert_score_f'] = F
                print(f'Finished calculating BERTScore, time passed {time.time() - start}s.')

            elif metric_name == 'bleu':
                sacrebleu = evaluate.load("sacrebleu")
                start = time.time()
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    sys_summ = self.data[doc_id]['sys_summ']
                    seg_scores = []
                    for out_line, ref_line in zip([sys_summ] * len(ref_summs), ref_summs):
                        results = sacrebleu.compute(predictions=[out_line], references=[[ref_line]])
                        seg_scores+= [results["score"]]
                    self.data[doc_id]['scores']['bleu'] = max(seg_scores)
                print(f'Finished calculating BLEU, time passed {time.time() - start}s.')

            elif metric_name == 'chrf':
                chrf = evaluate.load("chrf")
                start = time.time()
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    sys_summ = self.data[doc_id]['sys_summ']
                    seg_scores = []
                    for out_line, ref_line in zip([sys_summ] * len(ref_summs), ref_summs):
                        results = chrf.compute(predictions=[out_line], references=[[ref_line]])
                        seg_scores+= [results["score"]]
                    self.data[doc_id]['scores']['chrf'] = max(seg_scores)
                print(f'Finished calculating chrF, time passed {time.time() - start}s.')

            elif metric_name == 'ter':
                seg_scores=[]
                ter = evaluate.load("ter")
                start = time.time()
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    sys_summ = self.data[doc_id]['sys_summ']
                    seg_scores = []
                    for out_line, ref_line in zip([sys_summ] * len(ref_summs), ref_summs):
                        results = ter.compute(predictions=[out_line], references=[[ref_line]], case_sensitive=True)
                        seg_scores+= [results["score"]]
                    self.data[doc_id]['scores']['ter'] = max(seg_scores)
                print(f'Finished calculating TER, time passed {time.time() - start}s.')

            elif metric_name == 'bleurt':
                from bleurt import score
                checkpoint = "bleurt/bleurt-large-128"
                scorer = score.BleurtScorer(checkpoint)
                start = time.time()
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    sys_summ = self.data[doc_id]['sys_summ']
                    scores_ls = scorer.score(references=ref_summs, candidates=[sys_summ] * len(ref_summs))
                    self.data[doc_id]['scores']['bleurt'] = max(scores_ls)
                print(f'Finished calculating bleurt, time passed {time.time() - start}s.')

            elif metric_name == 'sescore':
                from comet import load_from_checkpoint
                model = load_from_checkpoint('sescore_english.ckpt')
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    sys_summ = self.data[doc_id]['sys_summ']
                    data = [{"src": ref_summs_ele, "mt": sys_summ_ele} for ref_summs_ele, sys_summ_ele in \
                    zip(ref_summs, [sys_summ] * len(ref_summs))]
                    scores_ls, _ = model.predict(data, batch_size=16, gpus=1)
                    self.data[doc_id]['scores']['sescore'] = max(scores_ls)
                print(f'Finished calculating SEScore, time passed {time.time() - start}s.')

            elif metric_name == 'contra_score':
                import torch
                exp_config.device_id = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
                model = torch.load('epoch0_best_7072_zh_en.ckpt').to(exp_config.device_id)
                model.eval()
                tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-large-tok")

                start = time.time()
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    sys_summ = self.data[doc_id]['sys_summ']
                    cur_data_dict = {'pivot': ref_summs, 'mt': [sys_summ] * len(ref_summs)}
                    cur_data_loader = preprocess_data(cur_data_dict, tokenizer, exp_config.max_length, len(ref_summs), shuffle=False, sampler=False, mode='test')
                    scores_ls = []
                    for batch in cur_data_loader:
                        score = model(batch, 'last_layer').squeeze(1).tolist()
                        scores_ls.extend(score)
                    self.data[doc_id]['scores']['contra_score'] = max(scores_ls)
                print(f'Finished calculating ContraScore, time passed {time.time() - start}s.')

            elif metric_name == 'mover_score':
                from moverscore import word_mover_score, get_idf_dict

                # Set up MoverScore
                self.stop_words = []

                ref_lines = []
                for doc_id in self.data:
                    ref_lines.extend(self.data[doc_id]['ref_summs'])
                ref_lines = list(set(ref_lines))
                self.idf_refs = get_idf_dict(ref_lines)

                # IDF for all system hypos, used for MoverScore
                sys_lines = []
                for doc_id in self.data:
                    sys_summ = self.data[doc_id]['sys_summ']
                    sys_lines.append(sys_summ)
                self.idf_hyps = get_idf_dict(sys_lines)
                print(f'MoverScore setup finished. Begin calculating MoverScore.')

                start = time.time()
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    sys_summ = self.data[doc_id]['sys_summ']
                    scores = word_mover_score(ref_summs, [sys_summ] * len(ref_summs), self.idf_refs, self.idf_hyps,
                                              self.stop_words,
                                              n_gram=1, remove_subwords=True, batch_size=48, device=self.device)
                    score = max(scores)
                    self.data[doc_id]['scores']['mover_score'] = score
                print(f'Finished calculating MoverScore, time passed {time.time() - start}s.')

            elif metric_name == 'rouge':
                from gehrmann_rouge_opennmt.rouge_baselines.baseline import baseline_main

                def rouge(dic):
                    """ Get r, p, f scores """
                    r1_, r2_, rl_ = [], [], []
                    for k in dic:
                        r1_.append([dic[k]['rouge_1_recall'], dic[k]['rouge_1_precision'], dic[k]['rouge_1_f_score']])
                        r2_.append([dic[k]['rouge_2_recall'], dic[k]['rouge_2_precision'], dic[k]['rouge_2_f_score']])
                        rl_.append([dic[k]['rouge_l_recall'], dic[k]['rouge_l_precision'], dic[k]['rouge_l_f_score']])
                    return r1_, r2_, rl_

                print(f'Begin calculating ROUGE.')
                start = time.time()
                blockPrint()
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    sys_summ = self.data[doc_id]['sys_summ']
                    sys_summ = sys_summ.lower()
                    write_list_to_file([sys_summ] * len(ref_summs), 'hypo.txt')
                    write_list_to_file(ref_summs, 'ref.txt')
                    args = argparse.Namespace(check_repeats=True, delete=True, get_each_score=True, stemming=True,
                                              method='sent_no_tag', n_bootstrap=1000, run_google_rouge=False,
                                              run_rouge=True, source='./hypo.txt', target='./ref.txt',
                                              ref_sep='||NEVER||', num_ref=1, temp_dir='./temp/')

                    scores = baseline_main(args, return_pyrouge_scores=True)['individual_score_results']
                    r1, r2, rl = rouge(scores)
                    r1 = np.max(r1, axis=0)
                    r2 = np.max(r2, axis=0)
                    rl = np.max(rl, axis=0)

                    self.data[doc_id]['scores']['rouge1_r'] = r1[0]
                    self.data[doc_id]['scores']['rouge1_p'] = r1[1]
                    self.data[doc_id]['scores']['rouge1_f'] = r1[2]
                    self.data[doc_id]['scores']['rouge2_r'] = r2[0]
                    self.data[doc_id]['scores']['rouge2_p'] = r2[1]
                    self.data[doc_id]['scores']['rouge2_f'] = r2[2]
                    self.data[doc_id]['scores']['rougel_r'] = rl[0]
                    self.data[doc_id]['scores']['rougel_p'] = rl[1]
                    self.data[doc_id]['scores']['rougel_f'] = rl[2]
                enablePrint()
                os.system('rm -rf hypo.txt ref.txt saved_out.txt')
                print(f'Finished calculating ROUGE, time passed {time.time() - start}s.')

            elif metric_name == 'prism':
                from prism import Prism
                # Set up Prism
                self.prism = Prism(model_dir='./models/m39v1/', lang='en')
                print(f'PRISM setup finished. Begin calculating PRISM.')

                start = time.time()
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    ref_summs = [detokenize(line) for line in ref_summs]
                    sys_summ = detokenize(self.data[doc_id]['sys_summ'])
                    ref_hypo_scores, hypo_ref_scores, scores = self.prism.score(cand=[sys_summ] * len(ref_summs),
                                                                                ref=ref_summs,
                                                                                segment_scores=True)
                    ref_hypo, hypo_ref, score = max(ref_hypo_scores), max(hypo_ref_scores), max(scores)
                    self.data[doc_id]['scores']['prism_ref_hypo'] = ref_hypo
                    self.data[doc_id]['scores']['prism_hypo_ref'] = hypo_ref
                    self.data[doc_id]['scores']['prism_avg'] = score
                print(f'Finished calculating PRISM, time passed {time.time() - start}s.')

            elif metric_name == 'bart_score' or metric_name == 'bart_score_cnn' or metric_name == 'bart_score_para':
                """ Vanilla BARTScore, BARTScore-CNN, BARTScore-CNN-Para """
                from bart_score import BARTScorer

                # Set up BARTScore
                if 'cnn' in metric_name:
                    bart_scorer = BARTScorer(device=self.device, checkpoint='facebook/bart-large-cnn')
                elif 'para' in metric_name:
                    bart_scorer = BARTScorer(device=self.device, checkpoint='facebook/bart-large-cnn')
                    bart_scorer.load()
                else:
                    bart_scorer = BARTScorer(device=self.device, checkpoint='facebook/bart-large')
                print(f'BARTScore setup finished. Begin calculating BARTScore.')

                start = time.time()
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    ref_summs = [detokenize(line) for line in ref_summs]
                    sys_summ = detokenize(self.data[doc_id]['sys_summ'])

                    ref_hypo_scores = np.array(bart_scorer.score(ref_summs, [sys_summ] * len(ref_summs), batch_size=4))
                    hypo_ref_scores = np.array(bart_scorer.score([sys_summ] * len(ref_summs), ref_summs, batch_size=4))
                    ref_hypo = ref_hypo_scores.max()
                    hypo_ref = hypo_ref_scores.max()
                    avg_f = (0.5 * (ref_hypo_scores + hypo_ref_scores)).max()
                    harm_f = (ref_hypo_scores * hypo_ref_scores / (ref_hypo_scores + hypo_ref_scores)).max()

                    self.data[doc_id]['scores'][f'{metric_name}_ref_hypo'] = ref_hypo
                    self.data[doc_id]['scores'][f'{metric_name}_hypo_ref'] = hypo_ref
                    self.data[doc_id]['scores'][f'{metric_name}_avg_f'] = avg_f
                    self.data[doc_id]['scores'][f'{metric_name}_harm_f'] = harm_f
                print(f'Finished calculating BARTScore, time passed {time.time() - start}s.')

            elif metric_name.startswith('prompt'):
                """ BARTScore adding prompts """
                from bart_score import BARTScorer

                def prefix_prompt(l, p):
                    new_l = []
                    for x in l:
                        new_l.append(p + ', ' + x)
                    return new_l

                def suffix_prompt(l, p):
                    new_l = []
                    for x in l:
                        new_l.append(x + ' ' + p + ',')
                    return new_l

                if 'cnn' in metric_name:
                    name = 'bart_score_cnn'
                    bart_scorer = BARTScorer(device=self.device, checkpoint='facebook/bart-large-cnn')
                elif 'para' in metric_name:
                    name = 'bart_score_para'
                    bart_scorer = BARTScorer(device=self.device, checkpoint='facebook/bart-large-cnn')
                    bart_scorer.load()
                else:
                    name = 'bart_score'
                    bart_scorer = BARTScorer(device=self.device, checkpoint='facebook/bart-large')

                print(f'BARTScore-P setup finished. Begin calculating BARTScore-P.')
                start = time.time()
                for doc_id in self.data:
                    ref_summs = self.data[doc_id]['ref_summs']
                    ref_summs = [detokenize(line) for line in ref_summs]
                    sys_summ = detokenize(self.data[doc_id]['sys_summ'])
                    sys_summs = [sys_summ] * len(ref_summs)
                    for prompt in REF_HYPO:
                        ref_hypo_scores_en = np.array(
                            bart_scorer.score(suffix_prompt(ref_summs, prompt), sys_summs, batch_size=4))
                        hypo_ref_scores_en = np.array(
                            bart_scorer.score(suffix_prompt(sys_summs, prompt), ref_summs, batch_size=4))
                        ref_hypo_scores_de = np.array(
                            bart_scorer.score(ref_summs, prefix_prompt(sys_summs, prompt), batch_size=4))
                        hypo_ref_scores_de = np.array(
                            bart_scorer.score(sys_summs, prefix_prompt(ref_summs, prompt), batch_size=4))
                        ref_hypo_en = ref_hypo_scores_en.max()
                        hypo_ref_en = hypo_ref_scores_en.max()
                        avg_f_en = (0.5 * (ref_hypo_scores_en + hypo_ref_scores_en)).max()
                        harm_f_en = (ref_hypo_scores_en * hypo_ref_scores_en / (
                                ref_hypo_scores_en + hypo_ref_scores_en)).max()
                        ref_hypo_de = ref_hypo_scores_de.max()
                        hypo_ref_de = hypo_ref_scores_de.max()
                        avg_f_de = (0.5 * (ref_hypo_scores_de + hypo_ref_scores_de)).max()
                        harm_f_de = (ref_hypo_scores_de * hypo_ref_scores_de / (
                                ref_hypo_scores_de + hypo_ref_scores_de)).max()

                        self.data[doc_id]['scores'][f'{name}_ref_hypo_en_{prompt}'] = ref_hypo_en
                        self.data[doc_id]['scores'][f'{name}_hypo_ref_en_{prompt}'] = hypo_ref_en
                        self.data[doc_id]['scores'][f'{name}_avg_f_en_{prompt}'] = avg_f_en
                        self.data[doc_id]['scores'][f'{name}_harm_f_en_{prompt}'] = harm_f_en
                        self.data[doc_id]['scores'][f'{name}_ref_hypo_de_{prompt}'] = ref_hypo_de
                        self.data[doc_id]['scores'][f'{name}_hypo_ref_de_{prompt}'] = hypo_ref_de
                        self.data[doc_id]['scores'][f'{name}_avg_f_de_{prompt}'] = avg_f_de
                        self.data[doc_id]['scores'][f'{name}_harm_f_de_{prompt}'] = harm_f_de
                print(f'Finished calculating BARTScore, time passed {time.time() - start}s.')

            else:
                raise NotImplementedError


def main():
    parser = argparse.ArgumentParser(description='Scorer parameters')
    parser.add_argument('--file', type=str, required=True,
                        help='The data to load from.')
    parser.add_argument('--device', type=str, default='cuda:0',
                        help='The device to run on.')
    parser.add_argument('--output', type=str, required=True,
                        help='The output path to save the calculated scores.')
    parser.add_argument('--bert_score', action='store_true', default=False,
                        help='Whether to calculate BERTScore')
    parser.add_argument('--mover_score', action='store_true', default=False,
                        help='Whether to calculate MoverScore')
    parser.add_argument('--rouge', action='store_true', default=False,
                        help='Whether to calculate ROUGE')
    parser.add_argument('--bart_score', action='store_true', default=False,
                        help='Whether to calculate BARTScore')
    parser.add_argument('--bart_score_cnn', action='store_true', default=False,
                        help='Whether to calculate BARTScore-CNN')
    parser.add_argument('--bart_score_para', action='store_true', default=False,
                        help='Whether to calculate BARTScore-Para')
    parser.add_argument('--prism', action='store_true', default=False,
                        help='Whether to calculate PRISM')
    parser.add_argument('--contra_score', action='store_true', default=False,
                        help='Whether to calculate contra_score')
    parser.add_argument('--bleurt', action='store_true', default=False,
                        help='Whether to calculate bleurt')
    parser.add_argument('--bleu', action='store_true', default=False,
                        help='Whether to calculate bleu')
    parser.add_argument('--chrf', action='store_true', default=False,
                        help='Whether to calculate chrf')
    parser.add_argument('--ter', action='store_true', default=False,
                        help='Whether to calculate ter')
    parser.add_argument('--prompt', type=str, default=None,
                        help='Whether to calculate BARTScore-P. Can be bart_ref, '
                             'bart_cnn_ref, bart_para_ref')
    args = parser.parse_args()

    scorer = Scorer(args.file, args.device)

    METRICS = []
    if args.bert_score:
        METRICS.append('bert_score')
    if args.mover_score:
        METRICS.append('mover_score')
    if args.rouge:
        METRICS.append('rouge')
    if args.bart_score:
        METRICS.append('bart_score')
    if args.bart_score_cnn:
        METRICS.append('bart_score_cnn')
    if args.bart_score_para:
        METRICS.append('bart_score_para')
    if args.prism:
        METRICS.append('prism')
    if args.contra_score:
        METRICS.append('contra_score')
    if args.bleurt:
        METRICS.append('bleurt')
    if args.bleu:
        METRICS.append('bleu')
    if args.chrf:
        METRICS.append('chrf')
    if args.ter:
        METRICS.append('ter')
    if args.prompt is not None:
        prompt = args.prompt
        assert prompt in ['bart_ref', 'bart_cnn_ref', 'bart_para_ref']
        METRICS.append(f'prompt_{prompt}')

    scorer.score(METRICS)
    scorer.save_data(args.output)


if __name__ == '__main__':
    main()
