# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Baiyu Chen and the COMODO-AnyMo-Bench contributors

"""Preprocessing that matches the original Nymeria COMODO loader and collate."""

import json
from pathlib import Path

import numpy as np
import torch
from torch.nn import functional as F
from transformers import BatchFeature, ProcessorMixin
from transformers.utils.hub import cached_file


class ComodoProcessor(ProcessorMixin):
    attributes = []
    model_input_names = ["imu_values"]

    def __init__(self, mean, std, sampling_rate=60, sequence_length=512):
        super().__init__()
        self.mean = np.asarray(mean, dtype=np.float32)
        self.std = np.asarray(std, dtype=np.float32)
        if self.mean.shape != (18,) or self.std.shape != (18,):
            raise ValueError("COMODO requires normalization statistics for 18 channels")
        self.std = np.where(self.std < 1e-6, 1.0, self.std).astype(np.float32)
        self.sampling_rate = sampling_rate
        self.sequence_length = sequence_length

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        variant = kwargs.pop("variant", None)
        hub_kwargs = {key: kwargs[key] for key in
                      ("subfolder", "revision", "token", "cache_dir", "local_files_only",
                       "force_download") if key in kwargs}
        if variant is not None:
            if variant not in ("core50", "fine150"):
                raise ValueError("variant must be 'core50' or 'fine150'")
            if hub_kwargs.get("subfolder") not in (None, "", variant):
                raise ValueError("variant conflicts with subfolder")
            hub_kwargs["subfolder"] = variant
        elif not hub_kwargs.get("subfolder"):
            metadata_path = cached_file(str(pretrained_model_name_or_path), "config.json",
                                        _raise_exceptions_for_missing_entries=False, **hub_kwargs)
            metadata = json.loads(Path(metadata_path).read_text()) if metadata_path else {}
            if metadata.get("is_variant_bundle", False):
                hub_kwargs["subfolder"] = metadata.get("variant", "core50")
        config = json.loads(Path(cached_file(str(pretrained_model_name_or_path),
                                             "processor_config.json", **hub_kwargs)).read_text())
        stats = json.loads(Path(cached_file(str(pretrained_model_name_or_path),
                                            "normalization_stats.json", **hub_kwargs)).read_text())
        return cls(stats["mean"], stats["std"], config["sampling_rate"], config["sequence_length"])

    def save_pretrained(self, save_directory, **kwargs):
        target = Path(save_directory)
        target.mkdir(parents=True, exist_ok=True)
        config = {"processor_class": "ComodoProcessor", "sampling_rate": self.sampling_rate,
                  "sequence_length": self.sequence_length,
                  "auto_map": {"AutoProcessor": "processing_comodo.ComodoProcessor"}}
        (target / "processor_config.json").write_text(json.dumps(config, indent=2) + "\n")
        (target / "normalization_stats.json").write_text(json.dumps({
            "mean": self.mean.tolist(), "std": self.std.tolist()}, indent=2) + "\n")
        source = Path(__file__)
        if source.resolve() != (target / source.name).resolve():
            (target / source.name).write_text(source.read_text())
        return [str(target / "processor_config.json")]

    def __call__(self, imu, sampling_rate=60, return_tensors="pt", **kwargs):
        if kwargs:
            raise TypeError(f"Unknown preprocessing arguments: {sorted(kwargs)}")
        if sampling_rate != self.sampling_rate:
            raise ValueError(f"Provide {self.sampling_rate} Hz IMU windows; got {sampling_rate} Hz")
        if return_tensors != "pt":
            raise ValueError("Only return_tensors='pt' is supported")
        if isinstance(imu, torch.Tensor):
            imu = imu.detach().cpu().numpy()
        try:
            values = np.asarray(imu, dtype=np.float32)
        except ValueError:
            values = None  # A batch of windows with different lengths.
        if values is not None and values.ndim == 2:
            windows = [values]
        elif values is not None and values.ndim == 3:
            windows = list(values)
        elif values is None and isinstance(imu, (list, tuple)):
            windows = [np.asarray(row, dtype=np.float32) for row in imu]
        else:
            raise ValueError("Expected [T, 18], [B, T, 18], or a list of [T, 18] windows")
        resized = []
        for window in windows:
            if window.ndim != 2 or window.shape[1] != 18 or window.shape[0] == 0:
                raise ValueError("Each IMU window must have shape [T, 18] with T > 0")
            window = np.nan_to_num(window, copy=True)
            window = (window - self.mean.reshape(1, -1)) / self.std.reshape(1, -1)
            tensor = torch.from_numpy(np.ascontiguousarray(window.T)).unsqueeze(0)
            resized.append(F.interpolate(tensor, size=self.sequence_length, mode="linear",
                                         align_corners=False).squeeze(0))
        if not resized:
            raise ValueError("The input batch must contain at least one window")
        return BatchFeature({"imu_values": torch.stack(resized)})