"""Признаки местоимений и референции к говорящему/собеседнику."""

from __future__ import annotations

import math

import pandas as pd

from ..lexicons import PRONOUN_CATEGORIES, WH_WORDS
from .base import as_list, safe_div, word_tokens


def extract(row: pd.Series) -> dict[str, float]:
    tokens = [token.lower() for token in word_tokens(as_list(row, "tokens"))]
    n_words = max(len(tokens), 1)
    features: dict[str, float] = {}

    for category, vocabulary in PRONOUN_CATEGORIES.items():
        features[f"pron_{category}_rate"] = safe_div(
            sum(1 for token in tokens if token in vocabulary), n_words
        )

    first_singular = sum(
        1 for token in tokens if token in PRONOUN_CATEGORIES["first_singular"]
    )
    second = sum(1 for token in tokens if token in PRONOUN_CATEGORIES["second"])
    features["pron_i_you_log_ratio"] = math.log((first_singular + 1.0) / (second + 1.0))
    features["pron_first_person_rate"] = safe_div(
        first_singular
        + sum(1 for token in tokens if token in PRONOUN_CATEGORIES["first_plural"]),
        n_words,
    )
    features["pron_third_person_rate"] = safe_div(
        sum(
            1
            for token in tokens
            if token in PRONOUN_CATEGORIES["third_singular"]
            or token in PRONOUN_CATEGORIES["third_plural"]
        ),
        n_words,
    )
    features["pron_wh_rate"] = safe_div(
        sum(1 for token in tokens if token in WH_WORDS), n_words
    )
    return features