# Base Dependencies

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

import numpy

# Local Dependencies

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

from models import RelationCollection

# 3rd-Party Dependencies

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

from sklearn.base import BaseEstimator

class TokenDistanceFeature(BaseEstimator):

    """

    TokenDistanceFeature

    Computes the number of tokens between the two entities of a relation.

    Source: 

        Alimova and Tutubalina (2020) - Multiple features for clinical relation extraction: A machine learning approach

    """

    def __init__(self):

        pass

    def get_feature_names(self, input_features=None):

        return ["token_dist"]

    def create_token_distance_feature(

        self, collection: RelationCollection

    ) -> numpy.array:

        features = []

        # max = 1

        for doc in collection.middle_tokens:

            features.append([len(doc)])

            # if len(r.middle_context) > max:

            #     max = len(r.middle_context)

        return numpy.array(features)

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

        return self

    def transform(self, x: RelationCollection, y=None) -> numpy.array:

        return self.create_token_distance_feature(x)

    def fit_transform(self, x: RelationCollection, y=None) -> numpy.array:

        return self.create_token_distance_feature(x)