inference.py

"""
scraper/space_a/inference.py

Работа с датами и timestamp-ами (МСК), валидация диапазонов, построение URL отчёта.
"""
from __future__ import annotations

import os
import json
import re
from datetime import datetime
from typing import List, Dict, Any, Tuple, Optional

import numpy as np
import pandas as pd

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torch.nn.functional import sigmoid

from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
from huggingface_hub import hf_hub_download

# =========================
# Пороги тональности
# =========================

# Повышенный порог для негатива снижает ложные срабатывания:
# модель должна быть уверена на 75%+, чтобы пометить текст как негативный.
# Позитив требует только 45% — он исторически предсказывается консервативнее.
_SENTIMENT_NEGATIVE_THRESHOLD = 0.75  # было ~0.50
_SENTIMENT_POSITIVE_THRESHOLD = 0.45  # было ~0.50

# =========================
# Общая очистка текста
# =========================

URL_PATTERN = re.compile(r"http\S+|www\.\S+")
MULTISPACE_PATTERN = re.compile(r"\s+")


def clean_text(text: str) -> str:
    if not isinstance(text, str):
        return ""
    text = text.replace("\n", " ").replace("\r", " ")
    text = URL_PATTERN.sub(" ", text)
    text = MULTISPACE_PATTERN.sub(" ", text)
    return text.strip()


# =========================
# Релевантность
# =========================

class RelevanceDatasetInfer(Dataset):
    def __init__(self, texts: List[str], tokenizer, max_length: int):
        self.texts = texts
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        enc = self.tokenizer(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )
        return {
            "input_ids": enc["input_ids"].squeeze(0),
            "attention_mask": enc["attention_mask"].squeeze(0),
        }


class RuBERTBinaryClassifier(nn.Module):
    def __init__(self, model_name: str, hidden_size: int = 512):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name)
        h = self.bert.config.hidden_size
        self.dropout1 = nn.Dropout(0.3)
        self.dense = nn.Linear(h, hidden_size)
        self.relu = nn.ReLU()
        self.dropout2 = nn.Dropout(0.2)
        self.classifier = nn.Linear(hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled = out.last_hidden_state[:, 0]
        x = self.dropout1(pooled)
        x = self.dense(x)
        x = self.relu(x)
        x = self.dropout2(x)
        logits = self.classifier(x).squeeze(-1)
        return logits


def load_relevance_model(relevance_repo_id: str, device: torch.device):
    config_path = hf_hub_download(relevance_repo_id, "config.json")
    with open(config_path, "r", encoding="utf-8") as f:
        conf = json.load(f)

    model_name = conf["model_name"]
    max_length = int(conf["max_length"])
    threshold = float(conf["threshold"])

    tokenizer = AutoTokenizer.from_pretrained(relevance_repo_id)
    model = RuBERTBinaryClassifier(model_name)
    weights_path = hf_hub_download(relevance_repo_id, "pytorch_model.bin")
    state = torch.load(weights_path, map_location=device)
    model.load_state_dict(state, strict=True)
    model.to(device)
    model.eval()
    return model, tokenizer, max_length, threshold


def infer_relevance(
    texts: List[str],
    model: nn.Module,
    tokenizer,
    max_length: int,
    threshold: float,
    device: torch.device,
    batch_size: int = 32,
) -> Tuple[np.ndarray, np.ndarray]:
    if not texts:
        return np.array([], dtype=int), np.array([], dtype=float)

    ds = RelevanceDatasetInfer(texts, tokenizer, max_length)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False)

    all_logits = []
    with torch.no_grad():
        for batch in loader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            logits = model(input_ids, attention_mask)
            all_logits.append(logits.cpu().numpy())

    if not all_logits:
        return np.array([], dtype=int), np.array([], dtype=float)

    all_logits = np.concatenate(all_logits, axis=0)
    probs = 1 / (1 + np.exp(-all_logits))
    preds = (probs >= threshold).astype(int)
    return preds, probs


# =========================
# Теги (multilabel)
# =========================

class RuBERTMultiLabelModel(nn.Module):
    def __init__(self, model_name: str, num_tags: int, hidden_size: int = 512):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name)
        h = self.bert.config.hidden_size
        self.dropout1 = nn.Dropout(0.3)
        self.dense = nn.Linear(h, hidden_size)
        self.relu = nn.ReLU()
        self.dropout2 = nn.Dropout(0.2)
        self.classifier = nn.Linear(hidden_size, num_tags)

    def forward(self, input_ids, attention_mask):
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled = out.last_hidden_state[:, 0]
        x = self.dropout1(pooled)
        x = self.dense(x)
        x = self.relu(x)
        x = self.dropout2(x)
        return self.classifier(x)


def load_tags_model(tags_repo_id: str, device: torch.device):
    tc_path = hf_hub_download(tags_repo_id, "training_config.json")
    with open(tc_path, "r", encoding="utf-8") as f:
        tc = json.load(f)

    model_name = tc.get("model_name", "DeepPavlov/rubert-base-cased")
    max_length = int(tc.get("max_length", 512))
    threshold = float(tc.get("threshold", 0.3))
    hidden_size = int(tc.get("hidden_size", 512))

    id2tag_path = hf_hub_download(tags_repo_id, "id2tag.json")
    with open(id2tag_path, "r", encoding="utf-8") as f:
        id2tag_raw = json.load(f)

    id2tag = {int(k): v for k, v in id2tag_raw.items()}
    num_tags = len(id2tag)
    all_tags = [id2tag[i] for i in sorted(id2tag.keys())]

    tokenizer = AutoTokenizer.from_pretrained(tags_repo_id)

    model = RuBERTMultiLabelModel(model_name, num_tags, hidden_size=hidden_size)
    weights_path = hf_hub_download(tags_repo_id, "pytorch_model.bin")
    state = torch.load(weights_path, map_location=device)
    model.load_state_dict(state, strict=True)
    model.to(device)
    model.eval()

    return model, tokenizer, max_length, all_tags, threshold


def predict_for_long_text_tags(
    model: nn.Module,
    text: str,
    tokenizer,
    device: torch.device,
    max_length: int,
) -> np.ndarray:
    enc_full = tokenizer(
        text,
        add_special_tokens=True,
        truncation=False,
        return_tensors="pt",
    )
    input_ids_full = enc_full["input_ids"][0]
    attention_full = enc_full["attention_mask"][0]
    seq_len = input_ids_full.shape[0]

    if seq_len <= max_length:
        input_ids = input_ids_full.unsqueeze(0).to(device)
        attention_mask = attention_full.unsqueeze(0).to(device)
        with torch.no_grad():
            logits = model(input_ids, attention_mask)
            probs = sigmoid(logits).squeeze(0).cpu().numpy()
        return probs

    logits_list = []
    for start in range(0, seq_len, max_length):
        end = start + max_length
        ids_chunk = input_ids_full[start:end]
        att_chunk = attention_full[start:end]
        if ids_chunk.shape[0] == 0:
            continue
        input_ids = ids_chunk.unsqueeze(0).to(device)
        attention_mask = att_chunk.unsqueeze(0).to(device)
        with torch.no_grad():
            logits = model(input_ids, attention_mask)
            logits_list.append(logits.squeeze(0).cpu().numpy())

    if not logits_list:
        enc = tokenizer(
            text,
            add_special_tokens=True,
            max_length=max_length,
            truncation=True,
            return_tensors="pt",
        )
        with torch.no_grad():
            logits = model(
                enc["input_ids"].to(device),
                enc["attention_mask"].to(device),
            )
        return sigmoid(logits).squeeze(0).cpu().numpy()

    logits_avg = np.mean(np.stack(logits_list, axis=0), axis=0)
    probs = 1 / (1 + np.exp(-logits_avg))
    return probs


def infer_tags_for_texts(
    texts_clean: List[str],
    model: nn.Module,
    tokenizer,
    max_length: int,
    all_tags: List[str],
    threshold: float,
    device: torch.device,
    return_probs: bool = False,
):
    results_tags = []
    results_probs = []

    for text in texts_clean:
        probs = predict_for_long_text_tags(model, text, tokenizer, device, max_length)
        mask = probs >= threshold
        idxs = np.where(mask)[0]
        tags = [all_tags[i] for i in idxs]

        results_tags.append(tags)
        results_probs.append(probs)

    if return_probs:
        return results_tags, results_probs
    return results_tags


# =========================
# Тональность
# =========================

# Индексы классов softmax: 0 = негатив, 1 = нейтрально, 2 = позитив
# Используется для кастомной архитектуры RuBERTSentimentClassifier.
# Для tbsa-модели (TBSAInferenceClassifier) маппинг фиксирован:
# tone_to_idx в обучении: -1→0, 0→1, 1→2, поэтому col0=neg, col1=neu, col2=pos.
_SENTIMENT_CLASS_NEGATIVE = 0
_SENTIMENT_CLASS_NEUTRAL  = 1
_SENTIMENT_CLASS_POSITIVE = 2

# TBSA-таргет — слово, которое подставляется перед текстом при токенизации.
# Должно совпадать с тем, что использовалось при обучении (_tbsa_train_engine.py).
_TBSA_TARGET = "Авито"


def index_to_tone(i: int) -> int:
    i = int(i)
    if i not in (0, 1, 2):
        raise ValueError(f"Некорректный индекс класса: {i}")
    return i - 1


def _apply_sentiment_thresholds(
    preds_tone: np.ndarray,
    probs: np.ndarray,
    negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD,
    positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD,
    neg_col: int = _SENTIMENT_CLASS_NEGATIVE,
    pos_col: int = _SENTIMENT_CLASS_POSITIVE,
) -> np.ndarray:
    """Применяет раздельные пороги к вектору тональностей.

    Правила:
    - Предсказан негатив (-1): если prob_negative < negative_threshold → нейтраль.
    - Предсказан позитив (+1): если prob_positive < positive_threshold → нейтраль.
    - Предсказана нейтраль (0): порог не применяется — остаётся нейтраль.

    Args:
        preds_tone:          массив int из {-1, 0, 1}, shape (N,)
        probs:               softmax-вероятности, shape (N, 3);
                             порядок колонок определяется neg_col / pos_col.
        negative_threshold:  минимальная уверенность для сохранения негатива
        positive_threshold:  минимальная уверенность для сохранения позитива
        neg_col:             индекс колонки с вероятностью негатива
        pos_col:             индекс колонки с вероятностью позитива

    Returns:
        Скорректированный массив тональностей, shape (N,).
    """
    result = preds_tone.copy()
    for i, tone in enumerate(preds_tone):
        if tone == -1:
            if probs[i, neg_col] < negative_threshold:
                result[i] = 0
        elif tone == 1:
            if probs[i, pos_col] < positive_threshold:
                result[i] = 0
    return result


class SentimentDatasetInfer(Dataset):
    """Датасет инференса тональности.

    Параметр tbsa_target: если задан (например "Авито"), токенизация
    выполняется как sentence-pair tokenizer(target, text) →
    [CLS] Авито [SEP] текст [SEP], что соответствует формату обучения
    TBSAClassifier. Также возвращает token_type_ids если токенайзер их
    поддерживает.

    Если tbsa_target=None — обычная одиночная токенизация.
    """
    def __init__(
        self,
        texts_clean: List[str],
        tokenizer,
        max_length: int,
        tbsa_target: Optional[str] = None,
    ):
        self.texts = texts_clean
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.tbsa_target = tbsa_target

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        if self.tbsa_target is not None:
            enc = self.tokenizer(
                self.tbsa_target,
                text,
                add_special_tokens=True,
                max_length=self.max_length,
                truncation=True,
                padding="max_length",
                return_tensors="pt",
                return_token_type_ids=True,
            )
        else:
            enc = self.tokenizer(
                text,
                add_special_tokens=True,
                max_length=self.max_length,
                truncation=True,
                padding="max_length",
                return_tensors="pt",
            )
        item = {
            "input_ids": enc["input_ids"].squeeze(0),
            "attention_mask": enc["attention_mask"].squeeze(0),
        }
        if "token_type_ids" in enc:
            item["token_type_ids"] = enc["token_type_ids"].squeeze(0)
        return item


class RuBERTSentimentClassifier(nn.Module):
    """Кастомная голова для старых чекпоинтов (не-tbsa)."""
    def __init__(self, model_name: str, num_labels: int = 3, hidden_size: int = 512):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name)
        h = self.bert.config.hidden_size
        self.dropout1 = nn.Dropout(0.3)
        self.dense = nn.Linear(h, hidden_size)
        self.relu = nn.ReLU()
        self.dropout2 = nn.Dropout(0.2)
        self.classifier = nn.Linear(hidden_size, num_labels)

    def forward(self, input_ids, attention_mask):
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled = out.last_hidden_state[:, 0]
        x = self.dropout1(pooled)
        x = self.dense(x)
        x = self.relu(x)
        x = self.dropout2(x)
        return self.classifier(x)


class TBSAInferenceClassifier(nn.Module):
    """Зеркало TBSAClassifier из _tbsa_train_engine.py для инференса.

    Архитектура: AutoModel (encoder) → Dropout → Linear(hidden, num_labels).
    Поддерживает token_type_ids, если модель их принимает (type_vocab_size > 1).
    Маппинг классов фиксирован обучением: col0=негатив, col1=нейтраль, col2=позитив.
    """
    def __init__(
        self,
        model_name: str,
        num_labels: int = 3,
        dropout: float = 0.1,
        use_mean_pooling: bool = False,
    ):
        super().__init__()
        self.use_mean_pooling = use_mean_pooling
        self.encoder = AutoModel.from_pretrained(model_name)
        enc_cfg = self.encoder.config
        self.supports_token_type_ids = getattr(enc_cfg, "type_vocab_size", 0) > 1
        print(f"[TBSAInferenceClassifier] supports_token_type_ids={self.supports_token_type_ids}")
        hidden = enc_cfg.hidden_size
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(hidden, num_labels)

    def forward(self, input_ids, attention_mask, token_type_ids=None):
        kwargs = dict(input_ids=input_ids, attention_mask=attention_mask)
        if self.supports_token_type_ids and token_type_ids is not None:
            kwargs["token_type_ids"] = token_type_ids
        out = self.encoder(**kwargs)
        if self.use_mean_pooling:
            m = attention_mask.unsqueeze(-1).float()
            pooled = (out.last_hidden_state * m).sum(1) / m.sum(1).clamp(min=1e-9)
        else:
            pooled = out.last_hidden_state[:, 0]
        return self.classifier(self.dropout(pooled))


# ── вспомогательная функция: маппинг id2label → neg/pos колонки ──────────────

def _resolve_label_cols(id2label: Dict[int, str]) -> Tuple[int, int, int]:
    """Возвращает (neg_col, neu_col, pos_col) по словарю id2label из config.

    Поддерживает любой порядок лейблов в HF-конфиге:
    например {0: 'NEGATIVE', 1: 'NEUTRAL', 2: 'POSITIVE'}
    или       {0: 'LABEL_0', 1: 'LABEL_1', 2: 'LABEL_2'} (fallback → 0,1,2).
    """
    neg_col = neu_col = pos_col = None
    for idx, label in id2label.items():
        l = label.lower()
        if any(x in l for x in ["neg", "-1", "негат"]):
            neg_col = int(idx)
        elif any(x in l for x in ["neu", "0", "нейтр"]):
            neu_col = int(idx)
        elif any(x in l for x in ["pos", "1", "позит"]):
            pos_col = int(idx)
    # fallback: если автодетект не сработал
    if neg_col is None:
        neg_col = _SENTIMENT_CLASS_NEGATIVE
    if neu_col is None:
        neu_col = _SENTIMENT_CLASS_NEUTRAL
    if pos_col is None:
        pos_col = _SENTIMENT_CLASS_POSITIVE
    return neg_col, neu_col, pos_col


def _is_tbsa_checkpoint(repo_id: str) -> bool:
    """Проверяет наличие 'tbsa' в имени репозитория (регистронезависимо)."""
    return "tbsa" in repo_id.lower()


def load_sentiment_model(sentiment_repo_id: str, device: torch.device):
    """Загружает модель тональности.

    Поддерживает два формата:
    1. tbsa / TBSAInferenceClassifier — когда репозиторий содержит 'tbsa'
       в названии или в config.json есть поле "tbsa": true.
       config.json — кастомный обучающий конфиг (_tbsa_train_engine.py):
       содержит model_name, max_length, dropout, use_mean_pooling, tbsa=True.
       Веса загружаются из best_pytorch_model.bin.
       Токенизация: sentence-pair (target, text) → [CLS] Авито [SEP] text [SEP].

    2. Кастомная архитектура RuBERTSentimentClassifier (старый формат) —
       config.json содержит ключи 'model_name' и 'max_length' (tbsa отсутствует).
       Веса загружаются из pytorch_model.bin.

    Возвращает (model, tokenizer, max_length).
    """
    # ── загружаем config.json ────────────────────────────────────────────────
    config_path = hf_hub_download(sentiment_repo_id, "config.json")
    with open(config_path, "r", encoding="utf-8") as f:
        conf = json.load(f)

    # ── определяем формат ────────────────────────────────────────────────────
    # tbsa=True: либо по имени репо, либо по явному флагу в конфиге
    is_tbsa = _is_tbsa_checkpoint(sentiment_repo_id) or bool(conf.get("tbsa", False))

    tokenizer = AutoTokenizer.from_pretrained(sentiment_repo_id)

    if is_tbsa:
        # ── tbsa: кастомный обучающий конфиг (_tbsa_train_engine.py) ─────────
        # config.json НЕ является HF PretrainedConfig — нет model_type/architectures.
        # Загружаем через TBSAInferenceClassifier (зеркало TBSAClassifier).
        model_name = conf["model_name"]           # "ai-forever/ruBERT-large"
        max_length = int(conf["max_length"])       # 256
        dropout = float(conf.get("dropout", 0.1))
        use_mean_pooling = bool(conf.get("use_mean_pooling", False))

        print(f"[INFO] TBSA-режим: model_name={model_name}, max_length={max_length}")

        model_inner = TBSAInferenceClassifier(
            model_name=model_name,
            num_labels=3,
            dropout=dropout,
            use_mean_pooling=use_mean_pooling,
        )
        weights_path = hf_hub_download(sentiment_repo_id, "best_pytorch_model.bin")
        state = torch.load(weights_path, map_location=device)
        model_inner.load_state_dict(state, strict=True)
        model_inner.to(device)
        model_inner.eval()

        # Маппинг фиксирован обучением (tone_to_idx: -1→0, 0→1, 1→2):
        # col 0 = негатив, col 1 = нейтраль, col 2 = позитив
        model = TbsaModelWrapper(
            model_inner,
            neg_col=0,
            neu_col=1,
            pos_col=2,
            tbsa_target=_TBSA_TARGET,
        )
        print(f"[INFO] tbsa neg_col=0  neu_col=1  pos_col=2  target='{_TBSA_TARGET}'")
    else:
        # ── кастомная архитектура (старый формат) ────────────────────────────
        model_name = conf["model_name"]
        max_length = int(conf["max_length"])

        model_inner = RuBERTSentimentClassifier(model_name)
        weights_path = hf_hub_download(sentiment_repo_id, "pytorch_model.bin")
        state = torch.load(weights_path, map_location=device)
        model_inner.load_state_dict(state, strict=True)
        model_inner.to(device)
        model_inner.eval()
        model = model_inner

    return model, tokenizer, max_length


class TbsaModelWrapper(nn.Module):
    """Обёртка над TBSAInferenceClassifier (или любой TBSA-моделью).

    Приводит интерфейс к общему виду:
        logits = model(input_ids, attention_mask)  → shape (B, num_labels)

    Хранит:
    - маппинг колонок (neg_col, neu_col, pos_col)
    - tbsa_target: строка-таргет для sentence-pair токенизации ("Авито")
      или None для обычной токенизации.
    """

    def __init__(
        self,
        hf_model: nn.Module,
        neg_col: int = _SENTIMENT_CLASS_NEGATIVE,
        neu_col: int = _SENTIMENT_CLASS_NEUTRAL,
        pos_col: int = _SENTIMENT_CLASS_POSITIVE,
        tbsa_target: Optional[str] = None,
    ):
        super().__init__()
        self.hf_model = hf_model
        self.neg_col = neg_col
        self.neu_col = neu_col
        self.pos_col = pos_col
        self.tbsa_target = tbsa_target

    def forward(self, input_ids, attention_mask, token_type_ids=None):
        if isinstance(self.hf_model, TBSAInferenceClassifier):
            return self.hf_model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                token_type_ids=token_type_ids,
            )
        # Для HF AutoModelForSequenceClassification (устаревший путь)
        out = self.hf_model(input_ids=input_ids, attention_mask=attention_mask)
        return out.logits


# ── единый конвертер «индекс предсказанного класса → тон {-1,0,1}» ───────────

def _preds_idx_to_tone(
    preds_idx: np.ndarray,
    model: nn.Module,
) -> np.ndarray:
    """Конвертирует argmax-индексы в тональность {-1, 0, 1}.

    Для TbsaModelWrapper использует neg_col/pos_col из обёртки.
    Для остальных моделей — index_to_tone (0→-1, 1→0, 2→+1).
    """
    if isinstance(model, TbsaModelWrapper):
        result = np.zeros(len(preds_idx), dtype=int)  # default нейтраль
        result[preds_idx == model.neg_col] = -1
        result[preds_idx == model.pos_col] = 1
        return result
    else:
        return np.array([index_to_tone(i) for i in preds_idx])


def infer_sentiment(
    texts_clean: List[str],
    model: nn.Module,
    tokenizer,
    max_length: int,
    device: torch.device,
    batch_size: int = 32,
    negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD,
    positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD,
) -> Dict[str, np.ndarray]:
    """Пакетный инференс тональности с раздельными порогами по классам."""
    _empty = {
        "logits": np.zeros((0, 3), dtype=float),
        "probs": np.zeros((0, 3), dtype=float),
        "preds_idx": np.zeros((0,), dtype=int),
        "preds_tone": np.zeros((0,), dtype=int),
    }
    if not texts_clean:
        return _empty

    # Определяем tbsa_target для sentence-pair токенизации
    tbsa_target = model.tbsa_target if isinstance(model, TbsaModelWrapper) else None

    ds = SentimentDatasetInfer(texts_clean, tokenizer, max_length, tbsa_target=tbsa_target)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False)

    all_logits = []
    with torch.no_grad():
        for batch in loader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            token_type_ids = batch.get("token_type_ids")
            if token_type_ids is not None:
                token_type_ids = token_type_ids.to(device)

            if isinstance(model, TbsaModelWrapper):
                logits = model(input_ids, attention_mask, token_type_ids=token_type_ids)
            else:
                logits = model(input_ids, attention_mask)
            all_logits.append(logits.cpu().numpy())

    if not all_logits:
        return _empty

    all_logits = np.concatenate(all_logits, axis=0)
    probs = torch.softmax(torch.tensor(all_logits), dim=-1).numpy()
    preds_idx = probs.argmax(axis=1)

    # ── конвертируем индексы → тон с учётом маппинга модели ──────────────────
    preds_tone_raw = _preds_idx_to_tone(preds_idx, model)

    # ── определяем neg/pos колонки для порогов ────────────────────────────────
    if isinstance(model, TbsaModelWrapper):
        neg_col = model.neg_col
        pos_col = model.pos_col
    else:
        neg_col = _SENTIMENT_CLASS_NEGATIVE
        pos_col = _SENTIMENT_CLASS_POSITIVE

    preds_tone = _apply_sentiment_thresholds(
        preds_tone_raw,
        probs,
        negative_threshold=negative_threshold,
        positive_threshold=positive_threshold,
        neg_col=neg_col,
        pos_col=pos_col,
    )

    return {
        "logits": all_logits,
        "probs": probs,
        "preds_idx": preds_idx,
        "preds_tone": preds_tone,
    }


def tone_to_str(tone_int: int) -> str:
    if tone_int == -1:
        return "негатив"
    elif tone_int == 0:
        return "нейтрально"
    elif tone_int == 1:
        return "позитив"
    return ""


def infer_sentiment_single(
    text_clean: str,
    model: nn.Module,
    tokenizer,
    max_length: int,
    device: torch.device,
    negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD,
    positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD,
    # Устаревший параметр для обратной совместимости:
    # если передан — используется как negative_threshold.
    tone_threshold: Optional[float] = None,
    # Новый явный параметр для управления порогом позитива из вне.
    tone_threshold_neg: Optional[float] = None,
) -> Dict[str, Any]:
    """Инференс тональности для одного текста."""
    _neg_thr = tone_threshold if tone_threshold is not None else negative_threshold
    _pos_thr = tone_threshold_neg if tone_threshold_neg is not None else positive_threshold

    res = infer_sentiment(
        [text_clean],
        model,
        tokenizer,
        max_length,
        device,
        batch_size=1,
        negative_threshold=_neg_thr,
        positive_threshold=_pos_thr,
    )
    probs = res["probs"][0]
    tone_int = int(res["preds_tone"][0])
    tone_str = tone_to_str(tone_int)

    # ── строим tone_probs в едином порядке (neg, neu, pos) ───────────────────
    if isinstance(model, TbsaModelWrapper):
        neg_col = model.neg_col
        neu_col = model.neu_col
        pos_col = model.pos_col
    else:
        neg_col = _SENTIMENT_CLASS_NEGATIVE
        neu_col = _SENTIMENT_CLASS_NEUTRAL
        pos_col = _SENTIMENT_CLASS_POSITIVE

    tone_probs = {
        "негатив":    float(probs[neg_col]),
        "нейтрально": float(probs[neu_col]),
        "позитив":    float(probs[pos_col]),
    }
    return {
        "tone_int":   tone_int,
        "tone_str":   tone_str,
        "tone_probs": tone_probs,
    }


# =========================
# Работа с JSON входом
# =========================

def load_input_json(json_path: str) -> List[Dict[str, Any]]:
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    if isinstance(data, list):
        return data
    elif isinstance(data, dict) and "items" in data:
        return data["items"]
    else:
        raise ValueError("Неизвестный формат входного JSON")


def extract_full_text_from_record(record: Dict[str, Any]) -> str:
    title = record.get("title", "") or ""
    text = record.get("text", "") or ""
    full_text_raw = (title + ". " + text).strip(". ")
    return clean_text(full_text_raw)


def format_ba_datetime(ts: Any) -> str:
    if not ts:
        return ""
    try:
        dt = datetime.fromtimestamp(ts)
        return dt.strftime("%d.%m.%Y %H:%M")
    except Exception:
        return ""


# =========================
# Пайплайн для JSON -> CSV
# =========================

def run_pipeline_json_to_new_ba_like_csv(
    input_json: str,
    relevance_model_dir: str,
    tags_model_dir: str,
    sentiment_model_dir: str,
    output_csv_path: str,
    limit_messages: Optional[int] = None,
    relevance_threshold: Optional[float] = None,
    tags_threshold: Optional[float] = None,
    # Устаревший параметр — заменён раздельными порогами ниже.
    tone_threshold: Optional[float] = None,
    sentiment_negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD,
    sentiment_positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD,
) -> str:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    rel_model, rel_tokenizer, rel_max_len, rel_thr_cfg = load_relevance_model(
        relevance_model_dir, device
    )
    tags_model, tags_tokenizer, tags_max_len, all_tags, tags_thr_cfg = load_tags_model(
        tags_model_dir, device
    )
    sent_model, sent_tokenizer, sent_max_len = load_sentiment_model(
        sentiment_model_dir, device
    )

    rel_thr = relevance_threshold if relevance_threshold is not None else rel_thr_cfg
    tags_thr = tags_threshold if tags_threshold is not None else tags_thr_cfg

    # Обратная совместимость: старый tone_threshold переопределяет neg-порог.
    neg_thr = tone_threshold if tone_threshold is not None else sentiment_negative_threshold
    pos_thr = sentiment_positive_threshold

    records = load_input_json(input_json)
    if limit_messages is not None:
        records = records[:limit_messages]

    rows_base: List[Dict[str, Any]] = []
    texts_clean: List[str] = []
    hubtype_all: List[str] = []

    for i, rec in enumerate(records):
        title = rec.get("title", "") or ""
        text = rec.get("text", "") or ""
        full_text_raw = (title + ". " + text).strip(". ")
        full_text_clean = clean_text(full_text_raw)

        time_ts = rec.get("timeCreate") or rec.get("date") or 0
        if isinstance(time_ts, (int, float)) and time_ts > 0:
            date_str = datetime.fromtimestamp(time_ts).strftime("%d.%m.%Y %H:%M")
        else:
            date_str = ""

        base_row = {
            "Дата": date_str,
            "ID сообщения": rec.get("id", i),
            "Hash сообщения": rec.get("hash", ""),
            "Заголовок": title,
            "Текст": text,
            "Источник": rec.get("hub", ""),
            "Url": rec.get("url", rec.get("href", "")),
            "Тип источника": rec.get("sourceType", ""),
            "Тип сообщения": rec.get("type", ""),
            "Сюжет": rec.get("topic", ""),
            "Автор": rec.get("authorName", ""),
            "Url автора": rec.get("authorUrl", ""),
            "Тип автора": rec.get("authorType", ""),
            "Место публикации": rec.get("placeName", ""),
            "Url места публикации": rec.get("placeUrl", ""),
            "Пол": rec.get("gender", ""),
            "Возраст": rec.get("age", ""),
            "Аудитория": rec.get("audience", ""),
            "Комментариев": rec.get("commentsCount", ""),
            "Цитируемость СМИ": "",
            "Репостов": rec.get("repostsCount", ""),
            "Лайков": rec.get("likesCount", ""),
            "Вовлеченность": "",
            "Просмотров": rec.get("viewsCount", ""),
            "Оценка": "",
            "Дублей": "",
            "Аудитория СМИ": "",
            "Тональность": "",
            "Роль объекта": "",
            "Агрессия": "",
            "Страна": rec.get("country", ""),
            "Регион": rec.get("region", ""),
            "Город": rec.get("city", ""),
            "Язык": rec.get("lang", "Русский"),
            "WOM": "",
            "Обработано": "",
            "Место": rec.get("place", ""),
            "Адрес": rec.get("address", ""),
            "Product": rec.get("product", ""),
            "General": rec.get("general", ""),
        }

        rows_base.append(base_row)
        texts_clean.append(full_text_clean)
        hubtype_all.append(rec.get("hubtype", "") or "")

    rel_preds, _ = infer_relevance(
        texts_clean, rel_model, rel_tokenizer, rel_max_len, rel_thr, device
    )

    relevant_indices = np.where(rel_preds == 1)[0]
    tags_all = [[] for _ in records]
    tone_str_all = ["" for _ in records]

    if len(relevant_indices) > 0:
        texts_rel = [texts_clean[idx] for idx in relevant_indices]

        tags_rel = infer_tags_for_texts(
            texts_rel,
            tags_model,
            tags_tokenizer,
            tags_max_len,
            all_tags,
            tags_thr,
            device,
        )
        for local_i, global_i in enumerate(relevant_indices):
            tags_for_item = tags_rel[local_i]
            if hubtype_all[global_i] == "Онлайн-СМИ" and "Новостные публикации" in all_tags:
                if "Новостные публикации" not in tags_for_item:
                    tags_for_item.append("Новостные публикации")
            tags_all[global_i] = tags_for_item

        sent_res = infer_sentiment(
            texts_rel,
            sent_model,
            sent_tokenizer,
            sent_max_len,
            device,
            negative_threshold=neg_thr,
            positive_threshold=pos_thr,
        )
        for local_i, global_i in enumerate(relevant_indices):
            tone_int = int(sent_res["preds_tone"][local_i])
            tone_str_all[global_i] = tone_to_str(tone_int)

    df = pd.DataFrame(rows_base)
    df["Релевантность"] = rel_preds.astype(int)
    df["Тональность"] = tone_str_all

    for tag in all_tags:
        df[tag] = 0

    for i, tags in enumerate(tags_all):
        for tag in tags:
            if tag in df.columns:
                df.at[i, tag] = 1

    df.to_csv(output_csv_path, sep=";", index=False, encoding="utf-8")
    return output_csv_path


# =========================
# Оценка модели (метрики по тегам)
# =========================

def evaluate_json_file_with_metrics(
    input_json: str,
    relevance_model_dir: str,
    tags_model_dir: str,
    sentiment_model_dir: str,
    limit_messages: Optional[int] = None,
    relevance_threshold: Optional[float] = None,
    tags_threshold: Optional[float] = None,
    tone_threshold: Optional[float] = None,
    sentiment_negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD,
    sentiment_positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD,
):
    """
    Считает micro-метрики и per-tag TP/FP/TN/FN/Precision/Recall/F1
    только по тегам.
    Истинные теги берём из поля 'tags' записи (список строк).
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    rel_model, rel_tokenizer, rel_max_len, rel_thr_cfg = load_relevance_model(
        relevance_model_dir, device
    )
    tags_model, tags_tokenizer, tags_max_len, all_tags, tags_thr_cfg = load_tags_model(
        tags_model_dir, device
    )
    _sent_model, _sent_tokenizer, _sent_max_len = load_sentiment_model(
        sentiment_model_dir, device
    )

    rel_thr = relevance_threshold if relevance_threshold is not None else rel_thr_cfg
    tags_thr = tags_threshold if tags_threshold is not None else tags_thr_cfg

    records = load_input_json(input_json)
    if limit_messages is not None:
        records = records[:limit_messages]

    texts_clean: List[str] = []
    true_tags_all: List[List[str]] = []
    hubtype_all: List[str] = []

    for rec in records:
        full_text = extract_full_text_from_record(rec)
        texts_clean.append(full_text)

        raw_tags = rec.get("tags", [])
        if isinstance(raw_tags, list):
            true_tags = [str(t).strip() for t in raw_tags if str(t).strip()]
        elif isinstance(raw_tags, str):
            true_tags = [t.strip() for t in raw_tags.split(";") if t.strip()]
        else:
            true_tags = []

        true_tags_all.append(true_tags)
        hubtype_all.append(rec.get("hubtype", "") or "")

    rel_preds, _ = infer_relevance(
        texts_clean, rel_model, rel_tokenizer, rel_max_len, rel_thr, device
    )

    relevant_indices = np.where(rel_preds == 1)[0]

    if len(relevant_indices) == 0:
        print("Нет релевантных записей.")
        return

    texts_rel = [texts_clean[idx] for idx in relevant_indices]

    tags_rel = infer_tags_for_texts(
        texts_rel,
        tags_model,
        tags_tokenizer,
        tags_max_len,
        all_tags,
        tags_thr,
        device,
    )

    pred_tags_all = [[] for _ in records]
    for local_i, global_i in enumerate(relevant_indices):
        tags_for_item = tags_rel[local_i]
        if hubtype_all[global_i] == "Онлайн-СМИ" and "Новостные публикации" in all_tags:
            if "Новостные публикации" not in tags_for_item:
                tags_for_item.append("Новостные публикации")
        pred_tags_all[global_i] = tags_for_item

    # ── per-tag метрики ───────────────────────────────────────────────────────
    metrics = {}
    for tag in all_tags:
        tp = fp = tn = fn = 0
        for i in relevant_indices:
            pred = tag in pred_tags_all[i]
            true = tag in true_tags_all[i]
            if pred and true:
                tp += 1
            elif pred and not true:
                fp += 1
            elif not pred and true:
                fn += 1
            else:
                tn += 1
        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        f1 = (
            2 * precision * recall / (precision + recall)
            if (precision + recall) > 0
            else 0.0
        )
        metrics[tag] = {
            "TP": tp, "FP": fp, "TN": tn, "FN": fn,
            "Precision": precision, "Recall": recall, "F1": f1,
        }

    # ── micro-метрики ─────────────────────────────────────────────────────────
    total_tp = sum(m["TP"] for m in metrics.values())
    total_fp = sum(m["FP"] for m in metrics.values())
    total_fn = sum(m["FN"] for m in metrics.values())
    micro_precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0.0
    micro_recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0.0
    micro_f1 = (
        2 * micro_precision * micro_recall / (micro_precision + micro_recall)
        if (micro_precision + micro_recall) > 0
        else 0.0
    )

    print(f"\n{'='*60}")
    print(f"  Micro Precision : {micro_precision:.4f}")
    print(f"  Micro Recall    : {micro_recall:.4f}")
    print(f"  Micro F1        : {micro_f1:.4f}")
    print(f"{'='*60}\n")
    print(f"{'Тег':<40} {'P':>6} {'R':>6} {'F1':>6} {'TP':>5} {'FP':>5} {'FN':>5}")
    print("-" * 75)
    for tag, m in sorted(metrics.items(), key=lambda x: -x[1]["F1"]):
        print(
            f"{tag:<40} {m['Precision']:>6.3f} {m['Recall']:>6.3f} "
            f"{m['F1']:>6.3f} {m['TP']:>5} {m['FP']:>5} {m['FN']:>5}"
        )

    return metrics