#!/usr/bin/env python3
# coding: utf-8

# In order to use this tool, please install comet first
# https://github.com/Unbabel/COMET

# Let's say you have a source file with N sentences in SL - eg: source.sl
# and the corresponding references (N sentences) reference.tl
# Translate your file in TL with the -n_best nbest options nbest being
# then number of hypotheses and output the target to -output target.nbest.tl
# Then you need to duplicate source and reference sentences nbest times
# for this script.
# for instance using awk '{for(i=1; i<=n; i++) print}' n=5 reference.tl \
#                          > reference.5.tl
# same for source.
# This script can be run (for instance with nbest = 5) as follows:
# python oracle_bleu.py --nbest-src source.5.sl --nbest-hyp target.5.tl \
#  --nbest-ref reference.5.tl --nbest-order 5 --output target.maxbleu.tl
# It will search in all hyp the best comet score

import codecs
from eole.bin import BaseBin, register_bin


def chunks(lgth, n):
    """Yield successive n-sized chunks from lgth."""
    for i in range(0, len(lgth), n):
        yield lgth[i : i + n]


# when choosing a reference-less model no nbest-ref is required


@register_bin(name="oracle_comet")
class OracleComet(BaseBin):
    @classmethod
    def add_args(cls, parser):
        parser.add_argument("--nbest-src", type=str, help="src repeated n times", required=True)
        parser.add_argument("--nbest-hyp", type=str, help="file with nbest to rerank", required=True)
        parser.add_argument("--nbest-ref", type=str, help="ref repeated n times", required=False)
        parser.add_argument("--nbest-order", type=int, help="nbest order", required=True)
        parser.add_argument("--output", type=str, help="output file", required=True)

    @classmethod
    def run(cls, args):
        from comet import download_model, load_from_checkpoint

        model_path = download_model("wmt21-comet-qe-mqm")
        model = load_from_checkpoint(model_path)

        with codecs.open(args.nbest_hyp, encoding="utf-8") as file:
            nbests = file.readlines()
            nbests = [item.strip() for item in nbests]
            nbests = chunks(nbests, args.nbest_order)

        if args.nbest_ref:
            with codecs.open(args.nbest_ref, encoding="utf-8") as file:
                nrefs = file.readlines()
                nrefs = [item.strip() for item in nrefs]
                nrefs = chunks(nrefs, args.nbest_order)

        with codecs.open(args.nbest_src, encoding="utf-8") as file:
            nsrcs = file.readlines()
            nsrcs = [item.strip() for item in nsrcs]
            nsrcs = chunks(nsrcs, args.nbest_order)

        with codecs.open(args.output, "w", encoding="utf-8") as output_file:
            best_indices = []
            # for nbest in nbests:
            if args.nbest_ref:
                for nsrc, nbest, nref in zip(nsrcs, nbests, nrefs):
                    data = []
                    for i in range(args.nbest_order):
                        data.append(
                            {
                                "src": nsrc[i],
                                "mt": nbest[i],
                                "ref": nref[i],
                            }
                        )
                        seg_scores, sys_score = model.predict(data, batch_size=args.nbest_order, gpus=1)
                        max_index = seg_scores.index(max(seg_scores))
                        output_file.write(data[max_index]["mt"] + "\n")
                        best_indices.append(max_index)
            else:
                for nsrc, nbest in zip(nsrcs, nbests):
                    data = []
                    for i in range(args.nbest_order):
                        data.append(
                            {
                                "src": nsrc[i],
                                "mt": nbest[i],
                            }
                        )
                        seg_scores, sys_score = model.predict(data, batch_size=args.nbest_order, gpus=1)
                        max_index = seg_scores.index(max(seg_scores))
                        output_file.write(data[max_index]["mt"] + "\n")
                        best_indices.append(max_index)

        for i in range(args.nbest_order):
            print(i, best_indices.count(i))
