"""chronos-x CLI: forecast from a saved model and a cached grid (private data stays local)."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from ..data.windows import SymbolGrid, WindowConfig, make_example
from ..train.sampling import collate
from .serialization import load_model


def main() -> None:
    ap = argparse.ArgumentParser(description="CHRONOS-X forecast at one origin")
    ap.add_argument("model"); ap.add_argument("grid_dir"); ap.add_argument("--session", type=int, required=True)
    ap.add_argument("--slot", type=int, required=True); ap.add_argument("--horizons", default="5,15,30,60")
    a = ap.parse_args()
    import torch
    m = load_model(Path(a.model)); m.eval()
    g = SymbolGrid.load(Path(a.grid_dir))
    ex = make_example(g, a.session, a.slot, WindowConfig()); ex["sym"] = np.int64(0)
    b = collate([ex])
    with torch.no_grad():
        o = m(b)
    out = {"symbol": g.symbol, "session_date": g.meta["session_dates"][a.session], "slot": a.slot, "sigma": float(ex["sigma"])}
    for H in (int(h) for h in a.horizons.split(",")):
        var = float((o["scale"][0, :H] ** 2 * (o["nu"][0, :H] / (o["nu"][0, :H] - 2))).sum())
        out[f"H{H}"] = {"mu_cum": float(o["mu"][0, :H].sum()), "sd_cum_indep": var ** 0.5}
    print(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()