"""Run this extractor on a document.

    pip install "adaptive-classifier>=0.3.1" torch transformers spacy dateparser numpy
    python -m spacy download en_core_web_sm

    from predict import Extractor
    x = Extractor(".")
    x.predict("Invoice INV-2207 ... Total due: $1,240.00 by 14 March 2026.",
              {"total": {"type": "number", "instructions": "The total amount due."},
               "due_date": {"type": "date", "instructions": "When payment is due."}})
    # {"total": {"value": 1240.0, "p_null": 0.07, "confidence": 0.93, "span": [..]}, "due_date": {...}}

``python predict.py request.json`` does the same for a request in the benchmark's
wire format (``state``, ``fields``, optional ``reference_date``).

How it works: candidate spans are generated from the document (regex, spaCy
entities and noun chunks, label-value lines, n-grams), the 40 most relevant per
field are kept, and one AdaptiveClassifier answers "is this the answer?" for
each, rendered with the field in context. The best candidate is the value;
below the tuned threshold the answer is null. It can only return text that is in
the document, so computed values (totals, counts) are out of its reach.
"""

from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Any, Dict, Mapping, Optional

sys.path.insert(0, str(Path(__file__).resolve().parent))

from typed_extraction import candidates as K  # noqa: E402
from typed_extraction.baselines.adaptive import AdaptiveExtractor  # noqa: E402
from typed_extraction.schema import Case, Field, Gold  # noqa: E402


class Extractor:
    def __init__(self, path: str = ".", device: str = "cpu") -> None:
        self.x = AdaptiveExtractor.load(path, device=device)

    def predict(self, state: str, fields: Mapping[str, Mapping[str, Any]],
                reference_date: Optional[str] = None) -> Dict[str, Dict[str, Any]]:
        fs = {n: Field(n, f["type"], f["instructions"],
                       {k: v for k, v in f.items() if k not in ("type", "instructions")}) for n, f in fields.items()}
        case = Case("request", "request", "request", "test", state, fs, {n: Gold(None) for n in fs}, reference_date)
        cands, sents = K.generate(state), K.sentences(state)
        short = [{n: K.shortlist(f, cands, state, reference_date, k=40, sent_spans=sents) for n, f in fs.items()}]
        return self.x.predict([case], short)[0]


if __name__ == "__main__":
    req = json.loads(Path(sys.argv[1]).read_text())
    print(json.dumps(Extractor(str(Path(__file__).resolve().parent)).predict(
        req["state"], req["fields"], req.get("reference_date")), indent=1, default=str))