#!/usr/bin/env python3 -u
# 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.
import argparse
import glob
from subprocess import check_call

try:
    import faiss

    has_faiss = True
except ImportError:
    has_faiss = False
import numpy as np


GB = 1024 * 1024 * 1024


def call(cmd):
    print(cmd)
    check_call(cmd, shell=True)


def get_batches(directory, lang, prefix="all_avg_pool"):
    print(f"Finding in {directory}/{prefix}.{lang}*")
    files = glob.glob(f"{directory}/{prefix}.{lang}*")
    emb_files = []
    txt_files = []
    for emb_fi in files:
        emb_files.append(emb_fi)
        txt_fi = emb_fi.replace(prefix, "sentences")
        txt_files.append(txt_fi)
    return emb_files, txt_files


def load_batch(emb_file, dim):
    embeddings = np.fromfile(emb_file, dtype=np.float32)
    num_rows = int(embeddings.shape[0] / dim)
    embeddings = embeddings.reshape((num_rows, dim))
    faiss.normalize_L2(embeddings)
    return embeddings


def knnGPU_sharded(x_batches_f, y_batches_f, dim, k, direction="x2y"):
    if not has_faiss:
        raise ImportError("Please install Faiss")
    sims = []
    inds = []
    xfrom = 0
    xto = 0
    for x_batch_f in x_batches_f:
        yfrom = 0
        yto = 0
        x_batch = load_batch(x_batch_f, dim)
        xto = xfrom + x_batch.shape[0]
        bsims, binds = [], []
        for y_batch_f in y_batches_f:
            y_batch = load_batch(y_batch_f, dim)
            neighbor_size = min(k, y_batch.shape[0])
            yto = yfrom + y_batch.shape[0]
            print("{}-{}  ->  {}-{}".format(xfrom, xto, yfrom, yto))
            idx = faiss.IndexFlatIP(dim)
            idx = faiss.index_cpu_to_all_gpus(idx)
            idx.add(y_batch)
            bsim, bind = idx.search(x_batch, neighbor_size)

            bsims.append(bsim)
            binds.append(bind + yfrom)
            yfrom += y_batch.shape[0]
            del idx
            del y_batch
        bsims = np.concatenate(bsims, axis=1)
        binds = np.concatenate(binds, axis=1)
        aux = np.argsort(-bsims, axis=1)
        sim_batch = np.zeros((x_batch.shape[0], k), dtype=np.float32)
        ind_batch = np.zeros((x_batch.shape[0], k), dtype=np.int64)
        for i in range(x_batch.shape[0]):
            for j in range(k):
                sim_batch[i, j] = bsims[i, aux[i, j]]
                ind_batch[i, j] = binds[i, aux[i, j]]
        sims.append(sim_batch)
        inds.append(ind_batch)
        xfrom += x_batch.shape[0]
        del x_batch
    sim = np.concatenate(sims, axis=0)
    ind = np.concatenate(inds, axis=0)
    return sim, ind


def score(sim, fwd_mean, bwd_mean, margin):
    return margin(sim, (fwd_mean + bwd_mean) / 2)


def score_candidates(
    sim_mat, candidate_inds, fwd_mean, bwd_mean, margin, verbose=False
):
    print(" - scoring {:d} candidates".format(sim_mat.shape[0]))
    scores = np.zeros(candidate_inds.shape)
    for i in range(scores.shape[0]):
        for j in range(scores.shape[1]):
            k = int(candidate_inds[i, j])
            scores[i, j] = score(sim_mat[i, j], fwd_mean[i], bwd_mean[k], margin)
    return scores


def load_text(files):
    all_sentences = []
    for fi in files:
        with open(fi) as sentence_fi:
            for line in sentence_fi:
                all_sentences.append(line.strip())
    print(f"Read {len(all_sentences)} sentences")
    return all_sentences


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Mine bitext")
    parser.add_argument("--src-lang", help="Source language")
    parser.add_argument("--tgt-lang", help="Target language")
    parser.add_argument(
        "--dict-path", help="Path to dictionary file", default="dict.txt"
    )
    parser.add_argument(
        "--spm-path", help="Path to SPM model file", default="sentence.bpe.model"
    )
    parser.add_argument("--dim", type=int, default=1024, help="Embedding dimension")
    parser.add_argument("--mem", type=int, default=5, help="Memory in GB")
    parser.add_argument("--src-dir", help="Source directory")
    parser.add_argument("--tgt-dir", help="Target directory")
    parser.add_argument("--output", help="Output path")
    parser.add_argument(
        "--neighborhood", type=int, default=4, help="Embedding dimension"
    )
    parser.add_argument(
        "--threshold", type=float, default=1.06, help="Threshold on mined bitext"
    )
    parser.add_argument(
        "--valid-size",
        type=int,
        default=2000,
        help="Number of sentences used for validation set",
    )
    parser.add_argument(
        "--min-count",
        type=int,
        default=50000,
        help="Min num sentences used for each language",
    )
    args = parser.parse_args()

    x_batches_f, x_sents_f = get_batches(args.src_dir, args.src_lang)
    y_batches_f, y_sents_f = get_batches(args.tgt_dir, args.tgt_lang)
    margin = lambda a, b: a / b
    y2x_sim, y2x_ind = knnGPU_sharded(
        y_batches_f, x_batches_f, args.dim, args.neighborhood, direction="y2x"
    )
    x2y_sim, x2y_ind = knnGPU_sharded(
        x_batches_f, y_batches_f, args.dim, args.neighborhood, direction="x2y"
    )

    x2y_mean = x2y_sim.mean(axis=1)
    y2x_mean = y2x_sim.mean(axis=1)
    fwd_scores = score_candidates(x2y_sim, x2y_ind, x2y_mean, y2x_mean, margin)
    bwd_scores = score_candidates(y2x_sim, y2x_ind, y2x_mean, x2y_mean, margin)
    fwd_best = x2y_ind[np.arange(x2y_sim.shape[0]), fwd_scores.argmax(axis=1)]
    bwd_best = y2x_ind[np.arange(y2x_sim.shape[0]), bwd_scores.argmax(axis=1)]
    indices = np.stack(
        (
            np.concatenate((np.arange(x2y_ind.shape[0]), bwd_best)),
            np.concatenate((fwd_best, np.arange(y2x_ind.shape[0]))),
        ),
        axis=1,
    )
    scores = np.concatenate((fwd_scores.max(axis=1), bwd_scores.max(axis=1)))

    x_sentences = load_text(x_sents_f)
    y_sentences = load_text(y_sents_f)

    threshold = args.threshold
    min_count = args.min_count
    seen_src, seen_trg = set(), set()
    directory = args.output
    call(f"mkdir -p {directory}")
    src_out = open(
        f"{directory}/all.{args.src_lang}",
        mode="w",
        encoding="utf-8",
        errors="surrogateescape",
    )
    tgt_out = open(
        f"{directory}/all.{args.tgt_lang}",
        mode="w",
        encoding="utf-8",
        errors="surrogateescape",
    )
    scores_out = open(
        f"{directory}/all.scores", mode="w", encoding="utf-8", errors="surrogateescape"
    )
    count = 0
    for i in np.argsort(-scores):
        src_ind, trg_ind = indices[i]
        if src_ind not in seen_src and trg_ind not in seen_trg:
            seen_src.add(src_ind)
            seen_trg.add(trg_ind)
            if scores[i] > threshold or count < min_count:
                if x_sentences[src_ind]:
                    print(scores[i], file=scores_out)
                    print(x_sentences[src_ind], file=src_out)
                    print(y_sentences[trg_ind], file=tgt_out)
                    count += 1
                else:
                    print(f"Ignoring sentence: {x_sentences[src_ind]}")
    src_out.close()
    tgt_out.close()
    scores_out.close()

    print(f"Found {count} pairs for threshold={threshold}")
    with open(f"{directory}/all.{args.src_lang}") as all_s, open(
        f"{directory}/all.{args.tgt_lang}"
    ) as all_t, open(f"{directory}/valid.{args.src_lang}", "w") as valid_s, open(
        f"{directory}/valid.{args.tgt_lang}", "w"
    ) as valid_t, open(
        f"{directory}/train.{args.src_lang}", "w"
    ) as train_s, open(
        f"{directory}/train.{args.tgt_lang}", "w"
    ) as train_t:
        count = 0
        for s_line, t_line in zip(all_s, all_t):
            s_line = s_line.split("\t")[1]
            t_line = t_line.split("\t")[1]
            if count >= args.valid_size:
                train_s.write(s_line)
                train_t.write(t_line)
            else:
                valid_s.write(s_line)
                valid_t.write(t_line)
                count += 1