Source code for cann.data.datasets

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict

import numpy as np
import torch
from torch.utils.data import Dataset

from cann.config.model_config import ModelConfig




@dataclass
class ParamSurfaceDataset(Dataset):
    """
    Synthetic dataset for surrogate learning/inference.

    Each sample is a tuple (θ_i, x_i, Q*(x_i; θ_i), w) where:
      - θ_i: Model parameters (e.g., rho, kappa) as a numpy array of shape (n,).
      - x_i: Market input features (e.g., moneyness, time to maturity) as a numpy array of shape (m,).
      - Q*(x_i; θ_i): Observed market value (e.g., implied volatility or option price) as a float.
      - w: Optional weight for the sample as a float.

        This dataclass stores the arrays directly as fields (``parameters``, ``market_inputs``,
        and optional ``observed``/``weights``) and prepares torch tensors in ``__post_init__``
        for efficient training.
    """

    parameters: np.ndarray
    market_inputs: np.ndarray
    model_cfg: ModelConfig
    observed: np.ndarray | None = None
    weights: np.ndarray | None = None

    @property
    def num_parameters(self) -> int:
        """Number of model parameters (n)."""
        return self.parameters.shape[1]

    @property
    def num_samples(self) -> int:
        """Number of samples (N)."""
        return self.parameters.shape[0]

    def __post_init__(self) -> None:
        # Validate and reshape numpy arrays
        self.parameters = np.asarray(self.parameters, dtype=np.float32)
        self.market_inputs = np.asarray(self.market_inputs, dtype=np.float32)

        num_samples = self.parameters.shape[0]

        if self.market_inputs.shape[0] != num_samples:
            raise ValueError(
                f"Number of samples in market_inputs ({self.market_inputs.shape[0]}) does not match parameters ({num_samples})"
            )

        # 2. Pre-convert to PyTorch Tensors for fast access
        self._params_t = torch.from_numpy(self.parameters)
        self._market_t = torch.from_numpy(self.market_inputs)

        self._observed_t = None
        if self.observed is not None:
            self.observed = np.asarray(self.observed, dtype=np.float32)
            if self.observed.shape[0] != num_samples:
                raise ValueError(
                    f"Number of samples in observed ({self.observed.shape[0]}) does not match parameters ({num_samples})"
                )
            self._observed_t = torch.from_numpy(self.observed.reshape(-1, 1))

        self._weights_t = None
        if self.weights is not None:
            self.weights = np.asarray(self.weights, dtype=np.float32).reshape(-1)
            if self.weights.shape[0] != num_samples:
                raise ValueError(
                    f"Number of samples in weights ({self.weights.shape[0]}) does not match parameters ({num_samples})"
                )
            self._weights_t = torch.from_numpy(self.weights.reshape(-1, 1))

    def __len__(self) -> int:
        return self.parameters.shape[0]

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        """
        Retrieve a single sample from the dataset.

        Returns a dictionary with keys:
            - "params": Model parameters tensor of shape (n,).
            - "market": Market input features tensor of shape (m,).
            - "observed": Observed target tensor of shape (1,) if available.
            - "weights": Sample weight tensor of shape (1,) if available.

        Args:
            idx (int): Index of the sample to retrieve.

        Returns:
            Dict[str, torch.Tensor]: Sample dictionary containing at least ``params`` and ``market``.
        """
        sample = {
            "params": self._params_t[idx],
            "market": self._market_t[idx],
        }
        if self._observed_t is not None:
            sample["observed"] = self._observed_t[idx]

        if self._weights_t is not None:
            sample["weights"] = self._weights_t[idx]

        return sample