"""Inference script for Multi-XY Regressor (TorchScript).
Usage:
    from inference import predict, load_model
    model, stats, config = load_model()
    predictions = predict(model, your_input_features)
"""
import json
import numpy as np
import torch
from huggingface_hub import hf_hub_download

def load_model(repo_id="sagkho/multi-xy-regressor"):
    config_path = hf_hub_download(repo_id=repo_id, filename="dataset_config.json")
    model_path = hf_hub_download(repo_id=repo_id, filename="model.pt")
    stats_path = hf_hub_download(repo_id=repo_id, filename="target_stats.npz")

    with open(config_path, "r") as f:
        config = json.load(f)

    model = torch.jit.load(model_path, map_location="cpu")
    model.eval()

    stats = np.load(stats_path)

    return model, stats, config

def predict(model, features, stats=None):
    """
    Args:
        model: loaded TorchScript model
        features: np.ndarray or list of input features [batch_size, n_features]
        stats: dict-like from np.load with 'mean' and 'std'
    Returns:
        np.ndarray of shape [batch_size, n_outputs] - denormalized predictions
    """
    if isinstance(features, list):
        features = np.array(features, dtype=np.float32)
    if features.ndim == 1:
        features = features.reshape(1, -1)

    x = torch.tensor(features, dtype=torch.float32)
    with torch.no_grad():
        pred_norm = model(x).numpy()

    if stats is not None:
        pred = pred_norm * stats["std"] + stats["mean"]
    else:
        pred = pred_norm

    return pred