#!/usr/bin/env python3
#
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

"""Extracts random constraints from reference files."""

import argparse
import random
import sys

from sacrebleu import extract_ngrams


def get_phrase(words, index, length):
    assert index < len(words) - length + 1
    phr = " ".join(words[index : index + length])
    for i in range(index, index + length):
        words.pop(index)
    return phr


def main(args):

    if args.seed:
        random.seed(args.seed)

    for line in sys.stdin:
        constraints = []

        def add_constraint(constraint):
            constraints.append(constraint)

        source = line.rstrip()
        if "\t" in line:
            source, target = line.split("\t")
            if args.add_sos:
                target = f"<s> {target}"
            if args.add_eos:
                target = f"{target} </s>"

            if len(target.split()) >= args.len:
                words = [target]

                num = args.number

                choices = {}
                for i in range(num):
                    if len(words) == 0:
                        break
                    segmentno = random.choice(range(len(words)))
                    segment = words.pop(segmentno)
                    tokens = segment.split()
                    phrase_index = random.choice(range(len(tokens)))
                    choice = " ".join(
                        tokens[phrase_index : min(len(tokens), phrase_index + args.len)]
                    )
                    for j in range(
                        phrase_index, min(len(tokens), phrase_index + args.len)
                    ):
                        tokens.pop(phrase_index)
                    if phrase_index > 0:
                        words.append(" ".join(tokens[0:phrase_index]))
                    if phrase_index + 1 < len(tokens):
                        words.append(" ".join(tokens[phrase_index:]))
                    choices[target.find(choice)] = choice

                    # mask out with spaces
                    target = target.replace(choice, " " * len(choice), 1)

                for key in sorted(choices.keys()):
                    add_constraint(choices[key])

        print(source, *constraints, sep="\t")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--number", "-n", type=int, default=1, help="number of phrases")
    parser.add_argument("--len", "-l", type=int, default=1, help="phrase length")
    parser.add_argument(
        "--add-sos", default=False, action="store_true", help="add <s> token"
    )
    parser.add_argument(
        "--add-eos", default=False, action="store_true", help="add </s> token"
    )
    parser.add_argument("--seed", "-s", default=0, type=int)
    args = parser.parse_args()

    main(args)