#!/usr/bin/env python3
"""Joint-модель планшета: один энкодер ruBERT-tiny2, две головы — интент и слоты.

Адаптация из smart-chat/mlx-chat/ft под домен smart-home-tablet. Машинерия
(энкодер + две головы, сегментация, сбор слота срезом исходной строки) — та же,
что в оригинале; изменены только INTENTS/SLOTS под инструменты планшета.

Головы читают один и тот же last_hidden_state, форвард один — скорость на
уровне обычного классификатора (~1-2 мс на CPU):
  intent — по [CLS], классификация всей фразы;
  slots  — по каждому токену, BIO-теги (извлечение куска текста).

Значение слота достаём по символьным offsets ИСХОДНОЙ строки, а не склейкой
subword-токенов: «Биг сити лайф» иначе пришлось бы собирать из кусков.
"""
import json
import os
import re

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer

BASE = "cointegrated/rubert-tiny2"

# --- Интенты: закрытые enum-действия закодированы в САМ интент -----------------
# Принцип: BIO-слоты хороши для свободного текста (город, трек, длительность),
# но плохи для выбора из фиксированного набора (направление транспорта,
# действие с музыкой, режим). Такие enum надёжнее сделать отдельными интентами,
# чем тащить отдельной головой. Поэтому 14 «широких» tools планшета развёрнуты
# в узкие интенты. Карта интент → tool: см. tablet_map.py.
# none — ПОСЛЕДНИМ: predict() возвращает cfg["intents"][-1] на пустой вход.
INTENTS = [
    "weather_get",             # get_weather            (+ city)
    "transport_to_center",     # get_transport{to_center}
    "transport_from_center",   # get_transport{from_center}
    "transport_all",           # get_transport{all}
    "calendar_today",          # get_today_events{today}
    "calendar_week",           # get_today_events{week}
    "calendar_create",         # create_event          (+ title,date,time,owner)
    "timer_set",               # set_timer             (+ duration)
    "timer_add",               # adjust_timer{+}       (+ duration)
    "timer_subtract",          # adjust_timer{-}       (+ duration)
    "timer_cancel",            # cancel_timer
    "notes_get",               # get_notes
    "home_state",              # get_smart_home_state
    "purifier_on",             # control_air_purifier{turn_on}
    "purifier_off",            # control_air_purifier{turn_off}
    "purifier_mode",           # control_air_purifier{set_mode}  (+ mode)
    "light_on",                # control_light{on}     (+ device) — НОВЫЙ tool
    "light_off",               # control_light{off}    (+ device) — НОВЫЙ tool
    "music_play",              # control_spotify{play + query}   (+ query)
    "music_pause",             # control_spotify{pause}
    "music_resume",            # control_spotify{play}
    "music_next",              # control_spotify{next}
    "music_prev",              # control_spotify{previous}
    "music_volume",            # control_spotify{volume}         (+ level)
    "now_playing",             # get_now_playing
    "none",                    # посторонняя речь → fallback на облачный LLM
]

# --- Слоты: только СВОБОДНЫЙ текст (закрытые значения ушли в интенты) ----------
SLOTS = [
    "city",      # город для погоды:     «в Питере», «Чаща»
    "title",     # название события:      «встреча с врачом»
    "date",      # дата события:          «25 июля», «завтра»
    "time",      # время события:         «13 часов», «в 15:30»
    "owner",     # чей календарь:         «Дани», «Свете»
    "duration",  # длительность таймера:  «15 минут», «полтора часа»
    "mode",      # режим очистителя:      «ночной», «авто»
    "device",    # какая лампа:           «свет у кровати», «лампу для чтения»
    "query",     # что включить:          «Ваню Дмитриенко», «Биг сити лайф»
    "level",     # громкость:             «50», «30 процентов»
]
TAGS = ["O"] + [f"{p}-{s}" for s in SLOTS for p in ("B", "I")]

# Режем текст на слова И цифры И знаки препинания по отдельности — ровно так же,
# как при обучении (gen_slots.seg). Разбор целиком делает модель: «13:30.» →
# токены «13», «:», «30», «.», которым она сама ставит теги. Никакой
# постобработки регуляркой. Граница буква/цифра тоже режет («15минут» слитно).
TOKEN_RE = re.compile(r"[^\W\d_]+|\d+|[^\w\s]")


def segment(text):
    """Строка → [(токен, начало, конец)] с позициями в исходной строке.

    Позиции нужны, чтобы собрать значение слота срезом ИСХОДНОГО текста:
    «свет у кровати» — три токена, склеивать вручную = гадать про пробелы.
    """
    return [(m.group(), m.start(), m.end()) for m in TOKEN_RE.finditer(text)]


class JointClassifier(nn.Module):
    def __init__(self, base=BASE, n_intents=len(INTENTS), n_tags=len(TAGS)):
        super().__init__()
        self.bert = AutoModel.from_pretrained(base)
        h = self.bert.config.hidden_size
        self.dropout = nn.Dropout(0.1)
        self.intent_head = nn.Linear(h, n_intents)
        self.slot_head = nn.Linear(h, n_tags)

    def forward(self, input_ids, attention_mask):
        h = self.bert(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        h = self.dropout(h)
        return self.intent_head(h[:, 0]), self.slot_head(h)


def save(model, tok, path):
    os.makedirs(path, exist_ok=True)
    torch.save(model.state_dict(), os.path.join(path, "joint.pt"))
    tok.save_pretrained(path)
    with open(os.path.join(path, "joint_config.json"), "w", encoding="utf-8") as f:
        json.dump({"base": BASE, "intents": INTENTS, "tags": TAGS}, f,
                  ensure_ascii=False, indent=2)


def load(path, device="cpu"):
    with open(os.path.join(path, "joint_config.json"), encoding="utf-8") as f:
        cfg = json.load(f)
    tok = AutoTokenizer.from_pretrained(path)
    model = JointClassifier(cfg["base"], len(cfg["intents"]), len(cfg["tags"]))
    model.load_state_dict(torch.load(os.path.join(path, "joint.pt"), map_location=device))
    return model.to(device).eval(), tok, cfg


@torch.no_grad()
def predict(model, tok, cfg, text, device="cpu"):
    """Строка → (intent, {slot: значение}, score).

    score — вероятность интента (softmax по голове интента). Потребитель
    (route.ts) сравнивает с порогом: score < 0.7 → трактовать как none и падать
    на облачный LLM.

    Сегментация на инференсе — ровно как при обучении (segment + split_words),
    иначе «13:30» рвётся иначе и теги съезжают.
    """
    segs = segment(text)
    if not segs:
        return cfg["intents"][-1], {}, 1.0
    words = [s[0] for s in segs]
    enc = tok([words], is_split_into_words=True, truncation=True,
              max_length=48, return_tensors="pt")
    word_ids = enc.word_ids(0)
    feed = {k: v.to(device) for k, v in enc.items()}
    intent_logits, slot_logits = model(feed["input_ids"], feed["attention_mask"])

    probs = torch.softmax(intent_logits[0], -1)
    best = int(probs.argmax())
    intent, score = cfg["intents"][best], float(probs[best])
    tag_ids = slot_logits[0].argmax(-1).tolist()

    # тег на слово — с первого subword, хвост игнорируем (у него при обучении -100)
    word_tags, prev = {}, None
    for tid, wid in zip(tag_ids, word_ids):
        if wid is not None and wid != prev:
            word_tags[wid] = cfg["tags"][tid]
        prev = wid

    # BIO-склейка: значение слота — срез исходного текста от начала первого
    # токена спана до конца последнего.
    slots, cur, start, end = {}, None, 0, 0
    for i, (_, s, e) in enumerate(segs):
        tag = word_tags.get(i, "O")
        if tag.startswith("B-"):
            if cur:
                slots.setdefault(cur, text[start:end])
            cur, start, end = tag[2:], s, e
        elif tag.startswith("I-") and cur == tag[2:]:
            end = e
        else:
            if cur:
                slots.setdefault(cur, text[start:end])
            cur = None
    if cur:
        slots.setdefault(cur, text[start:end])

    slots = {k: v for k, v in slots.items() if v}
    return intent, slots, score