#!/usr/bin/env python3
"""inference.py: Standalone reference inference for ngc1024/tokenledger-forecast."""

import json
from pathlib import Path
import numpy as np


class TokenForecasterPredictor:
    def __init__(self, model_dir: str = "."):
        p = Path(model_dir)
        with open(p / "features.json", "r", encoding="utf-8") as f:
            self.schema = json.load(f)
        self.means = np.array(self.schema["means"], dtype=np.float32)
        self.stds = np.array(self.schema["stds"], dtype=np.float32)

        # Try loading via onnxruntime if available, else PyTorch
        self.backend = None
        onnx_file = p / "model.onnx"
        if onnx_file.exists():
            try:
                import onnxruntime as ort
                self.session = ort.InferenceSession(str(onnx_file))
                self.backend = "onnx"
            except ImportError:
                pass

        if self.backend is None:
            import torch
            from model import TokenLedgerForecaster
            self.model = TokenLedgerForecaster.from_pretrained(str(p))
            self.model.eval()
            self.backend = "torch"

    def predict(self, turn_index: int, prompt_tokens: int, prev_output_tokens: int, gap_s: float, prev_cache_hit: float = 1.0):
        """Predicts p10, p50, and p90 output token forecasts for the next turn."""
        import math
        turn_idx_sc = min(turn_index, 200) / 50.0
        prompt_k = prompt_tokens / 10000.0
        prev_out_k = prev_output_tokens / 1000.0
        gap_log = math.log1p(max(0.0, float(gap_s)))
        window_fill = min(2.0, prompt_tokens / 200000.0)

        raw = np.array([turn_idx_sc, prompt_k, prev_out_k, gap_log, prev_cache_hit, window_fill], dtype=np.float32)
        norm = ((raw - self.means) / self.stds).reshape(1, -1)

        if self.backend == "onnx":
            out = self.session.run(["quantiles_p10_p50_p90"], {"context_features": norm})[0][0]
        else:
            import torch
            with torch.no_grad():
                out = self.model(torch.tensor(norm)).numpy()[0]

        return {
            "p10_tokens": round(float(out[0])),
            "p50_tokens": round(float(out[1])),
            "p90_tokens": round(float(out[2])),
            "tail_burst_risk": "HIGH" if out[2] > 2500 else ("MEDIUM" if out[2] > 1000 else "LOW"),
        }


if __name__ == "__main__":
    predictor = TokenForecasterPredictor("repo_model" if Path("repo_model").exists() else ".")
    res = predictor.predict(turn_index=15, prompt_tokens=140000, prev_output_tokens=320, gap_s=45.0)
    print("Forecast for turn 15 (Context=140k tokens, Prev=320 tokens, Gap=45s):")
    print(json.dumps(res, indent=2))