import json
import pytest

from stanza.models import ner_tagger
from stanza.models.common.doc import Document
from stanza.models.ner.data import DataLoader
from stanza.tests import TEST_WORKING_DIR

pytestmark = [pytest.mark.travis, pytest.mark.pipeline]


ONE_SENTENCE = """
[
 [
  {
   "text": "EU",
   "ner": "B-ORG"
  },
  {
   "text": "rejects",
   "ner": "O"
  },
  {
   "text": "German",
   "ner": "B-MISC"
  },
  {
   "text": "call",
   "ner": "O"
  },
  {
   "text": "to",
   "ner": "O"
  },
  {
   "text": "boycott",
   "ner": "O"
  },
  {
   "text": "Mox",
   "ner": "B-MISC"
  },
  {
   "text": "Opal",
   "ner": "I-MISC"
  },
  {
   "text": ".",
   "ner": "O"
  }
 ]
]
"""

@pytest.fixture(scope="module")
def pretrain_file():
    return f'{TEST_WORKING_DIR}/in/tiny_emb.pt'


@pytest.fixture(scope="module")
def one_sentence_json_path(tmpdir_factory):
    filename = tmpdir_factory.mktemp('data').join("sentence.json")
    with open(filename, 'w') as fout:
        fout.write(ONE_SENTENCE)
    return filename


def test_build_vocab(pretrain_file, one_sentence_json_path, tmp_path):
    """
    Test that when loading a data file, we get back 
    """
    args = ner_tagger.parse_args(["--wordvec_pretrain_file", pretrain_file])
    pt = ner_tagger.load_pretrain(args)

    with open(one_sentence_json_path) as fin:
        train_doc = Document(json.load(fin))

    train_batch = DataLoader(train_doc, args['batch_size'], args, pt, vocab=None, evaluation=False, scheme=args.get('train_scheme'), max_batch_words=args['max_batch_words'])

    vocab = train_batch.vocab
    pt_words = list(vocab['word'])
    assert pt_words == ['<PAD>', '<UNK>', '<EMPTY>', '<ROOT>', 'unban', 'mox', 'opal']
    delta_words = list(vocab['delta'])
    assert delta_words == ['<PAD>', '<UNK>', '<EMPTY>', '<ROOT>', 'eu', 'rejects', 'german', 'call', 'to', 'boycott', 'mox', 'opal', '.']
    tags = list(vocab['tag'])
    assert tags == [['<PAD>'], ['<UNK>'], [], ['<ROOT>'], ['S-ORG'], ['O'], ['S-MISC'], ['B-MISC'], ['E-MISC']]


def test_build_vocab_ignore_repeats(pretrain_file, one_sentence_json_path, tmp_path):
    """
    Test that when loading a data file, we get back 
    """
    args = ner_tagger.parse_args(["--wordvec_pretrain_file", pretrain_file, "--emb_finetune_known_only"])
    pt = ner_tagger.load_pretrain(args)

    with open(one_sentence_json_path) as fin:
        train_doc = Document(json.load(fin))

    train_batch = DataLoader(train_doc, args['batch_size'], args, pt, vocab=None, evaluation=False, scheme=args.get('train_scheme'), max_batch_words=args['max_batch_words'])

    vocab = train_batch.vocab
    pt_words = list(vocab['word'])
    assert pt_words == ['<PAD>', '<UNK>', '<EMPTY>', '<ROOT>', 'unban', 'mox', 'opal']
    delta_words = list(vocab['delta'])
    assert delta_words == ['<PAD>', '<UNK>', '<EMPTY>', '<ROOT>', 'mox', 'opal']
    tags = list(vocab['tag'])
    assert tags == [['<PAD>'], ['<UNK>'], [], ['<ROOT>'], ['S-ORG'], ['O'], ['S-MISC'], ['B-MISC'], ['E-MISC']]