#!/usr/bin/env python3
"""Two-stage XAUUSD headline model on CPU: gold-relevant? then direction.

A fine-tuned DeBERTa-v3-small encoder with two heads, exported to ONNX
with an int8 embedding table (transformer/train.py,
transformer/export_onnx.py).
Scoring needs only onnxruntime, tokenizers and numpy — no PyTorch.

For each headline:

- ``p_relevant``: probability the news moves the gold price (trained on
  DeepSeek |xau| >= 0.25 vs <= 0.10)
- ``p_bullish``: probability that, if it does, gold goes up
- ``relevant``: ``p_relevant >= relevance_threshold`` (from meta.json)
- ``label``: ``bullish`` / ``bearish``, or ``neutral`` when not relevant
- ``score``: ``p_relevant * (2 * p_bullish - 1)`` in [-1, 1]

Usage:
    python news_transformer.py score --headline "Fed cuts rates, signals more easing"
    python news_transformer.py score --input headlines.jsonl
"""

from __future__ import annotations

import argparse
import html
import json
import os
import sys
from pathlib import Path
from typing import Any, Iterable

import numpy as np

DEFAULT_MODEL_DIR = Path(__file__).resolve().parent / "models" / "xau_transformer"
MAX_LENGTH = 64
#: DeepSeek |xau| at or above this is "moves gold"; at or below
#: RELEVANT_MAX_IRRELEVANT it is "does not"; the band between carries no
#: relevance label in training.
RELEVANT_MIN = 0.25
RELEVANT_MAX_IRRELEVANT = 0.10


def clean_headline(headline: str) -> str:
    """The exact text normalisation used in training and scoring."""
    text = " ".join(html.unescape(str(headline or "")).split())
    if not text:
        raise ValueError("headline is empty")
    return text


def _sigmoid(x: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-x))


def load_model(model_dir: str | Path = DEFAULT_MODEL_DIR,
               threads: int | None = None) -> dict[str, Any]:
    """Load the ONNX model, tokenizer and metadata. Only load trusted files."""
    import onnxruntime as ort
    from tokenizers import Tokenizer

    model_dir = Path(model_dir)
    meta_path = model_dir / "meta.json"
    if not meta_path.is_file():
        raise FileNotFoundError(f"model not found: {meta_path}")
    meta = json.loads(meta_path.read_text(encoding="utf-8"))
    tokenizer = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
    tokenizer.enable_truncation(max_length=int(meta.get("max_length", MAX_LENGTH)))
    tokenizer.enable_padding(pad_id=int(meta["pad_token_id"]), pad_token=meta["pad_token"])
    options = ort.SessionOptions()
    if threads:
        options.intra_op_num_threads = max(1, min(int(threads), os.cpu_count() or 1))
        options.inter_op_num_threads = 1
    # ORT intra-op threads spin-wait after each run by default; in a mostly
    # idle bot process that is 4 cores at 100% between cycles.
    options.add_session_config_entry("session.intra_op.allow_spinning", "0")
    session = ort.InferenceSession(str(model_dir / meta.get("onnx_file", "model.onnx")),
                                   options, providers=["CPUExecutionProvider"])
    return {"meta": meta, "tokenizer": tokenizer, "session": session}


def score_batch(bundle: dict[str, Any], headlines: Iterable[str]) -> list[dict[str, Any]]:
    texts = [clean_headline(h) for h in headlines]
    if not texts:
        return []
    encodings = bundle["tokenizer"].encode_batch(texts)
    ids = np.array([e.ids for e in encodings], dtype=np.int64)
    mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)
    rel_logit, dir_logit = bundle["session"].run(
        ["relevance_logit", "direction_logit"],
        {"input_ids": ids, "attention_mask": mask})
    p_rel = _sigmoid(rel_logit.astype(np.float64))
    p_bull = _sigmoid(dir_logit.astype(np.float64))
    threshold = float(bundle["meta"].get("relevance_threshold", 0.5))
    out = []
    for text, pr, pb in zip(texts, p_rel, p_bull):
        relevant = bool(pr >= threshold)
        direction = "bullish" if pb >= 0.5 else "bearish"
        label = direction if relevant else "neutral"
        confidence = pr * max(pb, 1 - pb) if relevant else 1 - pr
        out.append({
            "target": "XAUUSD",
            "text": text,
            "relevant": relevant,
            "p_relevant": round(float(pr), 4),
            "direction": direction,
            "p_bullish": round(float(pb), 4),
            "label": label,
            "score": round(float(pr * (2 * pb - 1)), 4),
            "confidence": round(float(confidence), 4),
        })
    return out


def score_news(bundle: dict[str, Any], headline: str) -> dict[str, Any]:
    return score_batch(bundle, [headline])[0]


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    sub = parser.add_subparsers(dest="command", required=True)
    score = sub.add_parser("score", help="score headlines")
    score.add_argument("--model-dir", type=Path, default=DEFAULT_MODEL_DIR)
    score.add_argument("--threads", type=int, default=None)
    group = score.add_mutually_exclusive_group(required=True)
    group.add_argument("--headline")
    group.add_argument("--input", metavar="JSONL",
                       help='JSONL with a "headline" field per line, or - for stdin')
    args = parser.parse_args(argv)
    try:
        bundle = load_model(args.model_dir, args.threads)
        if args.headline:
            print(json.dumps(score_news(bundle, args.headline), indent=2))
            return 0
        handle = sys.stdin if args.input == "-" else open(args.input, encoding="utf-8")
        with handle:
            records = [json.loads(line) for line in handle if line.strip()]
        for start in range(0, len(records), 64):
            chunk = records[start:start + 64]
            for record, result in zip(chunk, score_batch(bundle, [r["headline"] for r in chunk])):
                print(json.dumps({**record, **result}))
    except (FileNotFoundError, OSError, ValueError, KeyError) as error:
        print(f"error: {error}", file=sys.stderr)
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())