# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.

from __future__ import absolute_import
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from fastText import train_supervised
from fastText import train_unsupervised
from fastText import load_model
from fastText import tokenize
import random
import sys
import os
import subprocess
import multiprocessing
import numpy as np
import unittest
import tempfile
import math
from scipy import stats


def compat_splitting(line):
    return line.decode('utf8').split()


def similarity(v1, v2):
    n1 = np.linalg.norm(v1)
    n2 = np.linalg.norm(v2)
    return np.dot(v1, v2) / n1 / n2


def read_vectors(model_path):
    vectors = {}
    with open(model_path, 'rb') as fin:
        for _, line in enumerate(fin):
            try:
                tab = compat_splitting(line)
                vec = np.array(tab[1:], dtype=float)
                word = tab[0]
                if np.linalg.norm(vec) == 0:
                    continue
                if word not in vectors:
                    vectors[word] = vec
            except ValueError:
                continue
            except UnicodeDecodeError:
                continue
    return vectors


def compute_similarity(model_path, data_path, vectors=None):
    if not vectors:
        vectors = read_vectors(model_path)

    mysim = []
    gold = []
    drop = 0.0
    nwords = 0.0

    with open(data_path, 'rb') as fin:
        for line in fin:
            tline = compat_splitting(line)
            word1 = tline[0].lower()
            word2 = tline[1].lower()
            nwords = nwords + 1.0

            if (word1 in vectors) and (word2 in vectors):
                v1 = vectors[word1]
                v2 = vectors[word2]
                d = similarity(v1, v2)
                mysim.append(d)
                gold.append(float(tline[2]))
            else:
                drop = drop + 1.0

    corr = stats.spearmanr(mysim, gold)
    dataset = os.path.basename(data_path)
    correlation = corr[0] * 100
    oov = math.ceil(drop / nwords * 100.0)
    return dataset, correlation, oov


def get_random_unicode(length):
    # See: https://stackoverflow.com/questions/1477294/generate-random-utf-8-string-in-python

    try:
        get_char = unichr
    except NameError:
        get_char = chr

    # Update this to include code point ranges to be sampled
    include_ranges = [
        (0x0021, 0x0021),
        (0x0023, 0x0026),
        (0x0028, 0x007E),
        (0x00A1, 0x00AC),
        (0x00AE, 0x00FF),
        (0x0100, 0x017F),
        (0x0180, 0x024F),
        (0x2C60, 0x2C7F),
        (0x16A0, 0x16F0),
        (0x0370, 0x0377),
        (0x037A, 0x037E),
        (0x0384, 0x038A),
        (0x038C, 0x038C),
    ]

    alphabet = [
        get_char(code_point)
        for current_range in include_ranges
        for code_point in range(current_range[0], current_range[1] + 1)
    ]
    return ''.join(random.choice(alphabet) for i in range(length))


def get_random_words(N, a, b):
    words = []
    for _ in range(N):
        length = random.randint(a, b)
        words.append(get_random_unicode(length))
    return words


class TestFastTextPy(unittest.TestCase):
    @classmethod
    def eprint(cls, *args, **kwargs):
        print(*args, file=sys.stderr, **kwargs)

    @classmethod
    def num_thread(cls):
        return multiprocessing.cpu_count() - 1

    @classmethod
    def build_paths(cls, train, test, output):
        train = os.path.join(cls.data_dir, train)
        test = os.path.join(cls.data_dir, test)
        output = os.path.join(cls.result_dir, output)
        return train, test, output

    @classmethod
    def build_train_args(cls, params, mode, train, output):
        args = [cls.bin, mode, "-input", train, "-output", output]
        return args + params.split(' ')

    @classmethod
    def get_train_output(cls, train_args):
        cls.eprint("Executing: " + ' '.join(train_args))
        return subprocess.check_output(train_args).decode('utf-8')

    @classmethod
    def get_path_size(cls, path):
        path_size = subprocess.check_output(["stat", "-c", "%s",
                                             path]).decode('utf-8')
        path_size = int(path_size)
        return path_size

    @classmethod
    def default_test_args(cls, model, test, quantize=False):
        return [cls.bin, "test", model, test]

    @classmethod
    def get_test_output(cls, test_args):
        cls.eprint("Executing: " + ' '.join(test_args))
        test_output = subprocess.check_output(test_args)
        test_output = test_output.decode('utf-8')
        cls.eprint("Test output:\n" + test_output)
        return list(
            map(lambda x: x.split('\t')[1], test_output.split('\n')[:-1])
        )

    @classmethod
    def train_generic_classifier(cls, train, output):
        thread = cls.num_thread()
        cls.eprint("Using {} threads".format(thread))
        sup_params = (
            "-dim 10 -lr 0.1 -wordNgrams 2 -minCount 1 -bucket 10000000 "
            "-epoch 5 -thread {}".format(thread)
        )
        mode = 'supervised'
        cls.get_train_output(
            cls.build_train_args(sup_params, mode, train, output)
        )

    @classmethod
    def train_generic_embeddings(cls, train, output):
        thread = cls.num_thread()
        cls.eprint("Using {} threads".format(thread))
        unsup_params = (
            "-thread {} -lr 0.025 -dim 100 -ws 5 -epoch 1 -minCount 5 "
            "-neg 5 -loss ns -bucket 2000000 -minn 3 -maxn 6 -t 1e-4 "
            "-lrUpdateRate 100".format(thread)
        )
        mode = 'cbow'
        cls.get_train_output(
            cls.build_train_args(unsup_params, mode, train, output)
        )

    def get_predictions_from_list(self, output, words, k):
        args = [self.bin, "predict-prob", output + '.bin', '-', str(k)]
        self.eprint("Executing: " + ' '.join(args))
        p = subprocess.Popen(
            args, stdin=subprocess.PIPE, stdout=subprocess.PIPE
        )
        test_text = ""
        if words:
            test_text = '\n'.join(words) + '\n'
        test_text = test_text.encode('utf-8')
        stdout, stderr = p.communicate(test_text)
        stdout = stdout.decode('utf-8')
        return stdout, stderr, p.returncode

    def get_word_vectors_from_list(self, output, words):
        args = [self.bin, "print-word-vectors", output + '.bin']
        self.eprint("Executing: " + ' '.join(args))
        p = subprocess.Popen(
            args, stdin=subprocess.PIPE, stdout=subprocess.PIPE
        )
        test_text = '\n'.join(words).encode('utf-8')
        stdout, stderr = p.communicate(test_text)
        return stdout


class TestFastTextPyUnit(TestFastTextPy):
    @classmethod
    def setUpClass(cls):
        cls.bin = os.environ['FASTTEXT_BIN']
        cls.data_dir = os.environ['FASTTEXT_DATA']
        cls.result_dir = tempfile.mkdtemp()
        train, _, output = cls.build_paths("fil9", "rw/rw.txt", "fil9")
        cls.train_generic_embeddings(train, output)
        cls.output = output
        train, _, output_sup = cls.build_paths(
            "dbpedia.train", "dbpedia.test", "dbpedia"
        )
        cls.train_generic_classifier(train, output_sup)
        cls.output_sup = output_sup

    @classmethod
    def tearDownClass(cls):
        pass
        # shutil.rmtree(cls.result_dir)

    # Check if get_word_vector aligns with vectors from stdin
    def test_getvector(self):
        f = load_model(self.output + '.bin')
        words, _ = f.get_words(include_freq=True)
        words += get_random_words(100, 1, 100)
        ftbin_vectors = self.get_word_vectors_from_list(self.output, words)
        ftbin_vectors = ftbin_vectors.decode('utf-8').split('\n')[:-1]
        for v in ftbin_vectors:
            word = v.split(' ')[0]
            vector = v.split(' ')[1:-1]
            vector = np.array(list(map(float, vector)))
            pvec = f.get_word_vector(word)
            # The fasttext cli returns floats with 5 digits,
            # but we use the full 6 digits.
            self.assertTrue(np.allclose(vector, pvec, rtol=1e-04))

    def test_predict(self):
        # TODO: I went a little crazy here as an exercise for
        # a rigorous test case. This could be turned into
        # a few utility functions.
        f = load_model(self.output_sup + '.bin')

        def _test(N, min_length, max_length, k, add_vocab=0):
            words = get_random_words(N, min_length, max_length)
            if add_vocab > 0:
                vocab, _ = f.get_words(include_freq=True)
                for _ in range(add_vocab):
                    ind = random.randint(0, len(vocab))
                    words += [vocab[ind]]
            all_labels = []
            all_probs = []