from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import numpy as np


class MotokoFeatureExtractor:
    """Normalize and stack haptic modalities into a single model tensor."""

    def __init__(self, config: dict[str, Any]) -> None:
        self.config = config
        self.max_length = int(config.get("max_length", 2048))
        self.padding_value = float(config.get("padding_value", 0.0))
        self.eps = float(config.get("normalization", {}).get("eps", 1e-6))
        self.modalities = config.get("modalities", {})

    @classmethod
    def from_config(cls, path: str | Path) -> "MotokoFeatureExtractor":
        with Path(path).open("r", encoding="utf-8") as handle:
            return cls(json.load(handle))

    def _normalize(self, values: np.ndarray) -> np.ndarray:
        mean = values.mean(axis=0, keepdims=True)
        std = values.std(axis=0, keepdims=True)
        return (values - mean) / np.maximum(std, self.eps)

    def _pad_or_trim(self, values: np.ndarray) -> np.ndarray:
        if values.shape[0] >= self.max_length:
            return values[: self.max_length]

        pad_rows = self.max_length - values.shape[0]
        pad = np.full((pad_rows, values.shape[1]), self.padding_value, dtype=values.dtype)
        return np.concatenate([values, pad], axis=0)

    def __call__(self, sample: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
        features: list[np.ndarray] = []

        for name, spec in self.modalities.items():
            if not spec.get("enabled", False):
                continue

            channels = int(spec["channels"])
            values = np.asarray(sample.get(name, np.zeros((0, channels), dtype=np.float32)))

            if values.ndim != 2 or values.shape[1] != channels:
                raise ValueError(
                    f"Expected modality '{name}' to have shape [timesteps, {channels}], "
                    f"got {values.shape}."
                )

            normalized = self._normalize(values.astype(np.float32))
            features.append(self._pad_or_trim(normalized))

        if not features:
            raise ValueError("No enabled modalities were provided.")

        stacked = np.concatenate(features, axis=1)
        attention_mask = (np.abs(stacked).sum(axis=1) > 0).astype(np.int64)

        return {
            "input_values": stacked,
            "attention_mask": attention_mask,
        }