# -*- coding: utf-8 -*-
"""
Response models for the CHARM SDK.

This module defines the data structures returned by the CHARM API.
"""

from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union


@dataclass
class EmbeddingsResponse:
    """
    Response from the embeddings API (POST /predict).

    What to expect
    --------------
    embeds : list, np.ndarray, or torch.Tensor — shape (N, D)
        One embedding vector per input time series. N = number of series you sent;
        D = embedding dimension (model-dependent). Type depends on return_tensors.
    model : str
        Model name (e.g. "charm-embedding").
    usage : dict or None
        Optional server usage stats (e.g. total_tokens).
    raw : dict or None
        Full raw JSON from the server (excluding binary embeds), for debugging.

    Input correspondence
    -------------------
    If you sent descriptions of shape (N, C) and ts_array of shape (N, T, C),
    then embeds has shape (N, D). The i-th row of embeds is the embedding for
    the i-th time series.

    Example
    -------
    >>> resp = client.embeddings.create(
    ...     descriptions=[["temp", "pressure"]],
    ...     ts_array=[[[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]],
    ...     return_tensors="np"
    ... )
    >>> resp.embeds.shape  # (1, D)
    >>> resp.model
    """
    object: str
    model: str
    embeds: Union[List[List[float]], "np.ndarray", "torch.Tensor"]
    usage: Optional[Dict[str, Any]] = None
    raw: Optional[Dict[str, Any]] = None  # full raw JSON (excluding embeds)


@dataclass
class ForecastResponse:
    """
    Response from the forecast/backcast API (POST /forecast).

    What to expect (shapes)
    ----------------------
    denormalized_predictions : list, np.ndarray, or torch.Tensor
        Shape (N, abs(target_len), C, Q). Predictions in original scale; use this
        for downstream work. N = batch size, C = channels, Q = number of quantiles (e.g. 21).
    predictions : same type as above
        Shape (N, abs(target_len), C, Q). Normalized (model-internal) predictions.
    data : same type as above
        Shape (N, T, C). Echo of your input ts_array.
    target_len : int
        The argument you passed (positive = forecast, negative = backcast).
    mode : "forecast" | "backcast"
        Set from sign of target_len.
    model : str
        Model name.
    raw : dict or None
        Full raw JSON from the server, for debugging.

    Input correspondence
    -------------------
    If you sent N series of shape (T, C) and target_len = 10, then
    denormalized_predictions has shape (N, 10, C, Q). For target_len = -8,
    shape is (N, 8, C, Q) (last 8 steps reconstructed).

    Example
    -------
    >>> resp = client.prediction.create(
    ...     descriptions=[["sensor_A", "sensor_B"]],
    ...     ts_array=[[[1.0, 2.0], [1.5, 2.5], [2.0, 3.0]]],
    ...     target_len=10,
    ...     return_tensors="np"
    ... )
    >>> resp.denormalized_predictions.shape  # (1, 10, 2, Q)
    >>> resp.mode  # "forecast"
    """
    object: str
    model: str
    predictions: Union[List, "np.ndarray", "torch.Tensor"]
    data: Union[List, "np.ndarray", "torch.Tensor"]
    denormalized_predictions: Union[List, "np.ndarray", "torch.Tensor"]
    target_len: int
    mode: str
    raw: Optional[Dict[str, Any]] = None