"""Minimal Hugging Face inference example for MultiPathFormer.

Run from a clone of the GitHub code release after installing requirements:

    python weights/inference_example.py

Set `AUGMENTED_PROMPT_JSON` and `FIRST_STEP_BASELINE_JSON` to your own files.
The augmented prompt must follow the released corridor-concat feature recipe:
TX/RX position, first-step delay/power baseline, first-step std, and scene
corridor features.
"""

from __future__ import annotations

import json
import os
from pathlib import Path

import numpy as np
from huggingface_hub import hf_hub_download

from multipathformer.inference import MultiPathFormerPredictor, prediction_to_rows


HF_REPO_ID = os.environ.get("HF_REPO_ID", "gblessed/multipathformer-foundation-27scenarios")
CHECKPOINT_FILENAME = "first_step_residual_corridor_concat_27scenarios_44710a4a_best_model_checkpoint.pth"
AUGMENTED_PROMPT_JSON = Path(os.environ.get("AUGMENTED_PROMPT_JSON", "examples/augmented_prompt.json"))
FIRST_STEP_BASELINE_JSON = Path(os.environ.get("FIRST_STEP_BASELINE_JSON", "examples/first_step_baseline.json"))


def main() -> None:
    checkpoint_path = hf_hub_download(repo_id=HF_REPO_ID, filename=CHECKPOINT_FILENAME)
    artifacts_path = hf_hub_download(repo_id=HF_REPO_ID, filename="preprocessing_artifacts.json")

    augmented_prompt = np.asarray(json.loads(AUGMENTED_PROMPT_JSON.read_text(encoding="utf-8")), dtype=np.float32)
    first_step_baseline = np.asarray(json.loads(FIRST_STEP_BASELINE_JSON.read_text(encoding="utf-8")), dtype=np.float32)

    predictor = MultiPathFormerPredictor(checkpoint_path, artifacts_path=artifacts_path)
    prediction = predictor.predict_from_augmented_prompt(augmented_prompt, first_step_baseline, max_steps=25)

    for row in prediction_to_rows(prediction)[:5]:
        print(row)


if __name__ == "__main__":
    main()