# Base Dependencies

# ----------------

import numpy as np

from typing import Optional

# Local Dependencies

# ------------------

from models import RelationCollection

from nlp_pipeline import get_pipeline, set_spacy_entities

# 3rd-Party Dependencies

# ----------------------

from sklearn.base import BaseEstimator

class NegatedEntitiesFeature(BaseEstimator):

    """

    Negated Entities Feature 

    Determines if each of the target entities of a relation is negated or not.

    """

    def __init__(self, padding_idx: Optional[int] = None):

        self.padding_idx = padding_idx

    def get_feature_names(self, input_features=None):

        return ["e1_negated", "e2_negated"]

    def create_negated_entities_feature(self, collection: RelationCollection) -> list:

        features = []

        NLP = get_pipeline()

        parser = NLP.get_pipe("parser")

        negex = NLP.get_pipe("negex")

        docs = collection.tokens 

        for i, doc in enumerate(parser.pipe(docs)):

            set_spacy_entities(

                doc,

                collection.left_tokens[i],

                collection.entities1_tokens[i],

                collection.relations[i].entity1.type,

                collection.middle_tokens[i],

                collection.entities2_tokens[i],

                collection.relations[i].entity2.type,

                collection.right_tokens[i],

            )

            assert len(doc.ents) == 2

            doc = negex(doc)

            e1_negated = int(doc.ents[0]._.negex)

            e2_negated = int(doc.ents[1]._.negex)

            features.append([e1_negated, e2_negated])

        return np.array(features)

    def fit(self, x: RelationCollection, y=None):

        return self

    def transform(self, x: RelationCollection) -> list:

        return self.create_negated_entities_feature(x)

    def fit_transform(self, x: RelationCollection, y=None) -> list:

        return self.create_negated_entities_feature(x)