"""
hv_intent.py

Hypervector intent router - word-level, supervised weights, k-NN bank.

Load:
    from hv_intent import load_v2, predict
    model = load_v2(".")
    predict(model, "what is the weather today")

Pure NumPy. No tokenizer, no gradient descent, no GPU.
"""

import json
import numpy as np
from pathlib import Path


DROP_THRESHOLD = 0.15
K_NEIGHBOURS = 7


def _encode_phrase(phrase, K, vocab_index, weights, drop_thr=DROP_THRESHOLD):
    items = []
    for t in phrase.lower().split():
        if t not in vocab_index:
            continue
        w = weights.get(t, 0.5)
        if w < drop_thr:
            continue
        items.append(K[vocab_index[t]].astype(np.float32) * w)
    if not items:
        return np.zeros(K.shape[1], dtype=np.float32)
    v = np.stack(items).sum(axis=0)
    n = np.linalg.norm(v)
    return v / n if n > 1e-30 else v


def load_v2(path):
    """Load the word-level model with supervised weights and k-NN bank."""
    path = Path(path)
    with open(path / "config.json") as f:
        config = json.load(f)
    data = np.load(path / "codebooks.npz", allow_pickle=True)

    n_cb = config["hyperparameters"]["n_codebooks"]
    codebooks = [data[f"K_{i}"] for i in range(n_cb)]
    vocab = list(data["vocab"])
    weights_arr = data["weights"]
    phrases = list(data["phrases"])
    labels = data["phrase_labels"]
    intent_labels = config["intent_labels"]

    vocab_index = {w: i for i, w in enumerate(vocab)}
    weights = {w: float(weights_arr[i]) for i, w in enumerate(vocab)}

    # Precompute bank hypervectors per codebook
    banks = []
    for K in codebooks:
        H = np.stack([_encode_phrase(p, K, vocab_index, weights)
                      for p in phrases])
        banks.append((H, labels))

    return {
        "codebooks": codebooks,
        "vocab_index": vocab_index,
        "weights": weights,
        "banks": banks,
        "intent_labels": intent_labels,
        "config": config,
    }


def predict(model, phrase, k=K_NEIGHBOURS, return_scores=False):
    """Return the predicted intent label for a single phrase."""
    C = len(model["intent_labels"])
    vote = np.zeros(C, dtype=np.float32)
    for K, (H, Y) in zip(model["codebooks"], model["banks"]):
        q = _encode_phrase(phrase, K, model["vocab_index"], model["weights"])
        s = H @ q
        k_eff = min(k, s.shape[0])
        top = np.argpartition(-s, k_eff - 1)[:k_eff]
        for j in top:
            vote[Y[j]] += max(s[j], 0.0)
    pred = int(vote.argmax())
    if return_scores:
        return model["intent_labels"][pred], vote
    return model["intent_labels"][pred]


def predict_batch(model, phrases, k=K_NEIGHBOURS):
    return [predict(model, p, k) for p in phrases]