# Base Dependencies
# ----------------
from typing import Optional

# Local Dependencies
# ------------------
from models import RelationCollection

# 3rd-Party Dependencies
# ----------------------
from sklearn.base import BaseEstimator

# Constants
# ---------
from constants import U_POS_TAGS


class POSFeature(BaseEstimator):
    """
    PoS Tagging

    Obtains the universal POS tag of each token in the relation's sentence.
    """

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

    def get_feature_names(self, input_features=None):
        return ["POS"]

    def create_pos_feature(self, collection: RelationCollection) -> list:
        all_pos = []

        for doc in collection.tokens:
            r_pos = []
            for t in doc:
                r_pos.append(self.pos_index(t.pos_))

            all_pos.append(r_pos)

        return all_pos

    def pos_index(self, pos_tag: str):
        """
        Computes the index of the corresponding POS tag
        """
        idx = U_POS_TAGS.index(pos_tag)

        if self.padding_idx is not None and idx >= self.padding_idx:
            idx += 1
        return idx

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

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

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