"""
Custom Inference Endpoint handler for amazon/chronos-2, shaped to match the contract
server/services/forecastService.js already calls (see HF_ENDPOINT_URL there).

Deploy target: your OWN Hugging Face copy of amazon/chronos-2 (duplicated into your
namespace), with this file and requirements.txt added at the repo root. HF Inference
Endpoints auto-detects handler.py and switches the endpoint's task to "Custom" — see
hf-endpoint/chronos-2-handler/README.md for the exact steps.

Request shape sent by forecastService.js:
  { "inputs": [<daily prices...>], "parameters": { "prediction_length": 7, "num_samples": 20 } }

Response shape forecastService.js parses (server/services/forecastService.js:337-351):
  { "quantiles": { "0.1": [...], "0.5": [...], "0.9": [...] } }
Each array has `prediction_length` entries (one per forecasted day).
"""
from typing import Any, Dict

import pandas as pd
from chronos import Chronos2Pipeline

QUANTILE_LEVELS = [0.1, 0.5, 0.9]


class EndpointHandler:
    def __init__(self, path: str = ""):
        # path is the directory HF mounts the repo's model weights into.
        self.pipeline = Chronos2Pipeline.from_pretrained(path, device_map="cpu")

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        inputs = data.get("inputs")
        if not isinstance(inputs, list) or len(inputs) == 0:
            return {"error": "inputs must be a non-empty list of numbers"}

        parameters = data.get("parameters") or {}
        prediction_length = int(parameters.get("prediction_length", 7))

        context_df = pd.DataFrame({
            "id": ["route"] * len(inputs),
            # Chronos-2 only uses these to order/space the series; the collector logs one
            # observation per day, so a daily index is the right granularity here.
            "timestamp": pd.date_range(end=pd.Timestamp.today().normalize(), periods=len(inputs), freq="D"),
            "target": [float(x) for x in inputs],
        })

        pred_df = self.pipeline.predict_df(
            context_df,
            prediction_length=prediction_length,
            quantile_levels=QUANTILE_LEVELS,
            id_column="id",
            timestamp_column="timestamp",
            target="target",
        )

        quantiles = {}
        for level in QUANTILE_LEVELS:
            key = str(level)
            if key in pred_df.columns:
                quantiles[key] = pred_df[key].tolist()
            elif level == 0.5 and "predictions" in pred_df.columns:
                # Some Chronos-2 releases name the point forecast column "predictions"
                # rather than "0.5" even when 0.5 is in quantile_levels.
                quantiles[key] = pred_df["predictions"].tolist()

        if "0.5" not in quantiles:
            return {"error": f"unexpected predict_df columns: {list(pred_df.columns)}"}

        return {"quantiles": quantiles}