"""Частеречные признаки на основе NLTK.

Используются две системы тегов: Universal POS (крупные классы) и Penn
Treebank (более дробные). Дополнительно считаются отношения и отдельные
грамматические категории.
"""

from __future__ import annotations

import pandas as pd

from ..lexicons import MODALS
from ..tokenization import ptb_tag_list, upos_tag_list
from .base import as_list, safe_div, word_tokens


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

    features: dict[str, float] = {}
    for tag in upos_tag_list():
        features[f"pos_upos_{tag}"] = safe_div(sum(1 for value in upos if value == tag), n_words)
    for tag in ptb_tag_list():
        features[f"pos_ptb_{tag}"] = safe_div(sum(1 for value in ptb if value == tag), n_words)

    noun = sum(1 for value in upos if value == "NOUN")
    verb = sum(1 for value in upos if value == "VERB")
    adjective = sum(1 for value in upos if value == "ADJ")
    adverb = sum(1 for value in upos if value == "ADV")
    pronoun = sum(1 for value in upos if value == "PRON")

    features.update(
        {
            "pos_noun_verb_ratio": safe_div(noun + 1.0, verb + 1.0),
            "pos_adj_adv_ratio": safe_div(adjective + 1.0, adverb + 1.0),
            "pos_modal_rate": safe_div(sum(1 for token in lower if token in MODALS), n_words),
            "pos_aux_rate": safe_div(sum(1 for value in upos if value == "AUX"), n_words),
            "pos_past_tense_rate": safe_div(sum(1 for value in ptb if value == "VBD"), n_words),
            "pos_present_tense_rate": safe_div(
                sum(1 for value in ptb if value in {"VBP", "VBZ"}), n_words
            ),
            "pos_gerund_rate": safe_div(sum(1 for value in ptb if value == "VBG"), n_words),
            "pos_participle_rate": safe_div(sum(1 for value in ptb if value == "VBN"), n_words),
            "pos_comparative_rate": safe_div(
                sum(1 for value in ptb if value in {"JJR", "RBR"}), n_words
            ),
            "pos_superlative_rate": safe_div(
                sum(1 for value in ptb if value in {"JJS", "RBS"}), n_words
            ),
            "pos_pronoun_rate": safe_div(pronoun, n_words),
            "pos_content_word_rate": safe_div(
                sum(1 for value in upos if value in {"NOUN", "VERB", "ADJ", "ADV", "NUM"}),
                n_words,
            ),
        }
    )
    return features