"""Inferência do classificador: encoder multilíngue + regressão logística em JSON."""
import json
import re

import numpy as np
from huggingface_hub import hf_hub_download
from sentence_transformers import SentenceTransformer


def chunks(text, size=900):
    cur, out = "", []
    for p in text.split("\n"):
        if len(cur) + len(p) > size and cur:
            out.append(cur.strip())
            cur = ""
        cur += p + "\n"
    if cur.strip():
        out.append(cur.strip())
    return [c[:1400] for c in out]


class DocTypeClassifier:
    def __init__(self, head, encoder=None):
        self.classes = head["classes"]
        self.coef = np.array(head["coef"])
        self.intercept = np.array(head["intercept"])
        self.mask = re.compile(head["label_mask"], re.I)
        self.encoder = encoder or SentenceTransformer(head["encoder"], device="cpu")

    def predict(self, text):
        """Devolve [(classe, probabilidade)] da mais provável para a menos provável."""
        parts = chunks(self.mask.sub("[DOC]", text))
        v = self.encoder.encode(parts, normalize_embeddings=True, show_progress_bar=False).mean(axis=0)
        v = v / np.linalg.norm(v)
        scores = self.coef @ v + self.intercept
        p = np.exp(scores - scores.max())
        p = p / p.sum()
        return sorted(zip(self.classes, map(float, p)), key=lambda x: -x[1])


def load(repo="fernandofa42/architecture-doc-type-classifier", encoder=None):
    with open(hf_hub_download(repo, "head.json")) as fh:
        return DocTypeClassifier(json.load(fh), encoder)