src/finetune/classifier.py

"""Transaction classifier — description → SAT account + deductibility + IVA.

Two implementations behind one interface:

* ``RuleClassifier``  — deterministic keyword matching over the catalog. Works today,
  no weights, and (because it shares the catalog) it agrees with the training labels.
  It is the honest fallback the app uses until the fine-tuned model is loaded.
* ``ModelClassifier`` — wraps an LLM client (the fine-tuned MiniCPM). It prompts for
  the same JSON schema and parses it, falling back to the rule classifier on any
  parse failure so the app never breaks.

The fine-tune's job is to generalize beyond the keyword list (messy OCR text, brand
names we never enumerated, Spanish variation) — but the schema and the labels are
defined once, here and in the catalog.
"""

from __future__ import annotations

import json
import re
import unicodedata
from dataclasses import dataclass
from typing import Optional

from .catalog import ALL_CATEGORIES, DEFAULT_EXPENSE, VENDORS, Category


def _norm(s: str) -> str:
    s = "".join(c for c in unicodedata.normalize("NFD", s.lower())
                if unicodedata.category(c) != "Mn")
    return s


@dataclass
class Classification:
    sat_code: str
    cuenta: str
    kind: str
    deducible: bool
    deducible_ratio: float
    iva_tasa: str
    iva_tratamiento: str
    confidence: float
    method: str  # "rule" | "model"

    @classmethod
    def from_category(cls, cat: Category, confidence: float, method: str) -> "Classification":
        lbl = cat.label()
        return cls(lbl["sat_code"], lbl["cuenta"], lbl["kind"], lbl["deducible"],
                   lbl["deducible_ratio"], lbl["iva_tasa"], lbl["iva_tratamiento"],
                   confidence, method)

    def to_dict(self) -> dict:
        return {
            "sat_code": self.sat_code, "cuenta": self.cuenta, "kind": self.kind,
            "deducible": self.deducible, "deducible_ratio": self.deducible_ratio,
            "iva_tasa": self.iva_tasa, "iva_tratamiento": self.iva_tratamiento,
            "confidence": self.confidence, "method": self.method,
        }


class RuleClassifier:
    """Keyword scoring over the catalog. Deterministic; the training-label oracle."""

    def __init__(self, categories=ALL_CATEGORIES):
        self.categories = categories
        self._by_code = {c.code: c for c in categories}
        # Known vendor names are strong category signals (and match the dataset).
        self._vendor_hits = [(_norm(v), code)
                             for code, vendors in VENDORS.items() for v in vendors
                             if code in self._by_code]

    def classify(self, description: str) -> Classification:
        text = _norm(description)
        tokens = set(re.findall(r"[a-z0-9]+", text))
        scores: dict = {}

        for cat in self.categories:
            score = 0.0
            for kw in cat.keywords:
                k = _norm(kw)
                if " " in k:
                    # multi-word keyword: match as a phrase (more specific → higher weight)
                    if k in text:
                        score += 1.5
                elif k in tokens or (len(k) >= 4 and any(t.startswith(k) for t in tokens)):
                    # single word: whole-token match (avoids 'gas' ⊂ 'gasto'),
                    # with a prefix allowance so plurals/inflections still hit.
                    score += 1.0
            if score:
                scores[cat.code] = scores.get(cat.code, 0.0) + score

        for vendor, code in self._vendor_hits:
            if vendor in text:
                scores[code] = scores.get(code, 0.0) + 2.0

        if not scores:
            return Classification.from_category(DEFAULT_EXPENSE, 0.3, "rule")
        best_code = max(scores, key=scores.get)
        confidence = min(0.5 + 0.2 * scores[best_code], 0.95)
        return Classification.from_category(self._by_code[best_code], confidence, "rule")


_SYSTEM = (
    "Eres un clasificador contable mexicano. Dada la descripción de una transacción, "
    "responde ÚNICAMENTE con un objeto JSON con las llaves: sat_code, cuenta, kind "
    "(income|expense|investment), deducible (bool), deducible_ratio (number), "
    "iva_tasa (string), iva_tratamiento (standard|zero|exempt|none). Sin texto extra."
)


def build_prompt(description: str) -> list:
    return [
        {"role": "system", "content": _SYSTEM},
        {"role": "user", "content": f"Clasifica: \"{description}\""},
    ]


def parse_label(text: str) -> Optional[dict]:
    """Extract the JSON label from a model completion (tolerant of fences/prose)."""
    if not text:
        return None
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if not m:
        return None
    try:
        return json.loads(m.group(0))
    except json.JSONDecodeError:
        return None


class RemoteClassifier:
    """Calls the Modal-served fine-tuned classifier; falls back to rules.

    The whole point of the fine-tune in the live demo: real generalization to brand
    names / messy OCR the keyword list never saw. If the endpoint is cold, slow, or
    down, we degrade to the deterministic RuleClassifier so the app never blocks.
    """

    REQUIRED = {"sat_code", "cuenta", "kind", "deducible", "iva_tasa", "iva_tratamiento"}

    def __init__(self, endpoint: str, fallback: Optional[RuleClassifier] = None,
                 timeout: float = 60.0):
        self.endpoint = endpoint
        self.fallback = fallback or RuleClassifier()
        self.timeout = timeout

    def classify(self, description: str) -> Classification:
        try:
            import urllib.request
            payload = json.dumps({"description": description}).encode()
            req = urllib.request.Request(
                self.endpoint, data=payload,
                headers={"Content-Type": "application/json"})
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                data = json.loads(resp.read())
        except Exception:
            return self.fallback.classify(description)
        if not isinstance(data, dict) or not self.REQUIRED.issubset(data):
            return self.fallback.classify(description)
        return Classification(
            sat_code=str(data["sat_code"]), cuenta=str(data["cuenta"]),
            kind=str(data["kind"]), deducible=bool(data["deducible"]),
            deducible_ratio=float(data.get("deducible_ratio", 1.0)),
            iva_tasa=str(data["iva_tasa"]), iva_tratamiento=str(data["iva_tratamiento"]),
            confidence=float(data.get("confidence", 0.9)), method="model")


class ModelClassifier:
    """Uses the fine-tuned model; falls back to rules on any failure."""

    REQUIRED = {"sat_code", "cuenta", "kind", "deducible", "iva_tasa", "iva_tratamiento"}

    def __init__(self, generate, fallback: Optional[RuleClassifier] = None):
        # generate: Callable[[list[messages]], str] — a thin completion function.
        self.generate = generate
        self.fallback = fallback or RuleClassifier()

    def classify(self, description: str) -> Classification:
        try:
            raw = self.generate(build_prompt(description))
            data = parse_label(raw)
        except Exception:
            data = None
        if not data or not self.REQUIRED.issubset(data):
            return self.fallback.classify(description)
        return Classification(
            sat_code=str(data["sat_code"]),
            cuenta=str(data["cuenta"]),
            kind=str(data["kind"]),
            deducible=bool(data["deducible"]),
            deducible_ratio=float(data.get("deducible_ratio", 1.0)),
            iva_tasa=str(data["iva_tasa"]),
            iva_tratamiento=str(data["iva_tratamiento"]),
            confidence=float(data.get("confidence", 0.9)),
            method="model",
        )