"""Признаки словарного богатства.

На коротких репликах эти признаки нестабильны. Они включены как
дополнительная группа и должны проверяться отдельно в ablation-эксперименте.
"""

from __future__ import annotations

import math
from collections import Counter

import pandas as pd

from .base import as_list, finite, safe_div, word_tokens


def _yule_k(counter: Counter[str], n_tokens: int) -> float:
    if n_tokens <= 1:
        return 0.0
    frequency_of_frequencies = Counter(counter.values())
    second_moment = sum(
        frequency * frequency * count
        for frequency, count in frequency_of_frequencies.items()
    )
    return finite(10_000.0 * (second_moment - n_tokens) / (n_tokens * n_tokens))


def _simpson_d(counter: Counter[str], n_tokens: int) -> float:
    if n_tokens <= 1:
        return 0.0
    numerator = sum(count * (count - 1) for count in counter.values())
    return finite(numerator / (n_tokens * (n_tokens - 1)))


def extract(row: pd.Series) -> dict[str, float]:
    tokens = [token.lower() for token in word_tokens(as_list(row, "tokens"))]
    counter = Counter(tokens)
    n_tokens = len(tokens)
    n_types = len(counter)
    hapax = sum(1 for count in counter.values() if count == 1)
    dis = sum(1 for count in counter.values() if count == 2)

    ttr = safe_div(n_types, n_tokens)
    guiraud = safe_div(n_types, math.sqrt(n_tokens))
    herdan = finite(math.log(n_types) / math.log(n_tokens)) if n_tokens > 1 and n_types > 1 else 0.0
    hapax_ratio = safe_div(hapax, n_types)
    dis_ratio = safe_div(dis, n_types)
    if n_types > 0 and hapax_ratio < 1.0:
        honore = finite(100.0 * math.log(n_tokens) / (1.0 - hapax_ratio))
    else:
        honore = 0.0

    top10 = sum(count for _, count in counter.most_common(10))
    return {
        "vocab_tokens": float(n_tokens),
        "vocab_types": float(n_types),
        "vocab_ttr": ttr,
        "vocab_guiraud": finite(guiraud),
        "vocab_herdan": herdan,
        "vocab_yule_k": _yule_k(counter, n_tokens),
        "vocab_simpson_d": _simpson_d(counter, n_tokens),
        "vocab_hapax_ratio": hapax_ratio,
        "vocab_dis_ratio": dis_ratio,
        "vocab_honore_r": honore,
        "vocab_top10_coverage": safe_div(top10, n_tokens),
        "vocab_mean_word_frequency": safe_div(n_tokens, n_types),
    }