# Base Dependencies
# ----------------
import pytest
import numpy as np
from random import randrange

# Local Dependencies
# ------------------
from features.iob_feature import IOBFeature
from constants import DDI_IOB_TAGS, N2C2_IOB_TAGS
from models import RelationCollection, Entity, Relation


# Tests
# ------
def test_iob_init():
    iob = IOBFeature("n2c2")
    assert isinstance(iob, IOBFeature)
    assert iob.dataset == "n2c2"
    assert iob.iob_tags == N2C2_IOB_TAGS
    assert iob.padding_idx is None

    iob = IOBFeature("ddi")
    assert isinstance(iob, IOBFeature)
    assert iob.dataset == "ddi"
    assert iob.iob_tags == DDI_IOB_TAGS
    assert iob.padding_idx is None


def test_iob_init_unsupported_dataset_raises():
    with pytest.raises(ValueError):
        iob = IOBFeature("i2b2")


def test_iob_get_feature_names():
    iob = IOBFeature("n2c2")
    assert iob.get_feature_names() == ["IOB"]


def test_iob_fit(n2c2_small_collection):
    iob = IOBFeature("n2c2")
    iob = iob.fit(n2c2_small_collection)
    assert isinstance(iob, IOBFeature)


def test_iob_index():
    for padding_idx in [0, 1, 5]:
        iob = IOBFeature("n2c2", padding_idx=padding_idx)

        for tag in iob.iob_tags:
            idx = iob.iob_index(tag)
            assert idx != padding_idx


def test_iob_index_unknown_tag_raises():
    iob = IOBFeature("n2c2", padding_idx=0)
    with pytest.raises(Exception):
        idx = iob.iob_idx("B-UNKN")


def test_iob_create_iob_feature(n2c2_small_collection, relation):

    iob = IOBFeature("n2c2")
    iob_feature = iob.create_iob_feature(n2c2_small_collection)
    assert len(iob_feature) == len(n2c2_small_collection)

    for i in range(len(n2c2_small_collection)):
        e1_tokens = n2c2_small_collection.entities1_tokens[i]
        e2_tokens = n2c2_small_collection.entities2_tokens[i]
        sent_tokens = n2c2_small_collection.tokens[i]

        assert len(iob_feature[i]) == len(sent_tokens)

        if len(e1_tokens) > 1 and len(e2_tokens) > 1:
            unique = 5
        elif len(e1_tokens) > 1 or len(e2_tokens) > 1:
            unique = 4
        else:
            unique = 3

        assert len(np.unique(iob_feature[i])) == unique

    one_collection = RelationCollection(relation)
    iob_feature = iob.create_iob_feature(one_collection)
    assert len(iob_feature) == 1
    assert list(iob_feature[0]) == [0, 0, 0, 1, 0, 0, 13, 14, 0, 0, 0]