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

"""IMU-only COMODO inference, preserving the original checkpoint parameter names."""

from dataclasses import dataclass
from typing import Optional

import torch
from mantis.architecture import MantisV2
from torch import nn
from torch.nn import functional as F
from transformers import PreTrainedModel
from transformers.utils import ModelOutput

from .configuration_comodo import ComodoConfig


@dataclass
class ComodoOutput(ModelOutput):
    embeddings: Optional[torch.Tensor] = None
    projection: Optional[torch.Tensor] = None


class ComodoBundle(PreTrainedModel):
    """Resolve the root bundle before loading a selected, fully trained model."""

    config_class = ComodoConfig
    main_input_name = "imu_values"

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, config=None, **kwargs):
        hub_keys = ("revision", "token", "cache_dir", "local_files_only", "force_download")
        hub_kwargs = {k: kwargs[k] for k in hub_keys if k in kwargs}
        config = config or cls.config_class.from_pretrained(pretrained_model_name_or_path, **hub_kwargs)
        if config._commit_hash is not None:
            hub_kwargs["revision"] = config._commit_hash
        load_kwargs = {k: v for k, v in kwargs.items()
                       if k not in ("variant", "_from_auto", "_from_pipeline", "adapter_kwargs")}
        load_kwargs.update(hub_kwargs)
        if "torch_dtype" not in load_kwargs and config.torch_dtype is not None:
            load_kwargs["torch_dtype"] = config.torch_dtype
        if kwargs.get("_from_pipeline", getattr(config, "_release_from_pipeline", None)) == "comodo":
            # Transformers passes pipeline-specific arguments only after its
            # initial AutoModel lookup. Defer weights until variant is available.
            selector = cls(config)
            selector.name_or_path = str(pretrained_model_name_or_path)
            selector._release_load_kwargs = load_kwargs
            return selector
        variant = kwargs.pop("variant", None) or config.variant
        if variant not in ("core50", "fine150"):
            raise ValueError("variant must be 'core50' or 'fine150'")
        return ComodoModel.from_pretrained(pretrained_model_name_or_path, *args,
                                          subfolder=variant, **load_kwargs)

    def forward(self, *args, **kwargs):
        raise RuntimeError("Select a variant through pipeline('comodo', ..., variant='core50')")


class ComodoModel(PreTrainedModel):
    config_class = ComodoConfig
    main_input_name = "imu_values"

    def __init__(self, config):
        super().__init__(config)
        # Construct the architecture only. The released checkpoint contains every
        # backbone parameter, so inference never downloads a second model.
        self.imu_student = MantisV2(**config.mantis_config, device="cpu")
        self.mlp = nn.Sequential(
            nn.Linear(config.feature_dim, config.projection_hidden_dim),
            nn.BatchNorm1d(config.projection_hidden_dim),
            nn.GELU(),
            nn.Linear(config.projection_hidden_dim, config.projection_dim),
        )
        self.post_init()

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        subfolder = kwargs.get("subfolder", "")
        hub_kwargs = {k: kwargs[k] for k in
                      ("revision", "token", "cache_dir", "local_files_only") if k in kwargs}
        model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
        target = model[0] if isinstance(model, tuple) else model
        target._release_subfolder = subfolder
        target._release_hub_kwargs = hub_kwargs
        return model

    def forward(self, imu_values, normalize_embeddings=True, return_projection=False,
                return_dict=True):
        if imu_values.ndim != 3 or imu_values.shape[1] != self.config.n_channels:
            raise ValueError(f"Expected [B, {self.config.n_channels}, T] IMU values")
        if imu_values.shape[-1] != self.config.sequence_length:
            imu_values = F.interpolate(imu_values, size=self.config.sequence_length,
                                       mode="linear", align_corners=False)
        features = torch.cat([
            self.imu_student(imu_values[:, ch:ch + 1, :])
            for ch in range(self.config.n_channels)
        ], dim=1)
        projection = self.mlp(features) if return_projection else None
        embeddings = F.normalize(features, p=2, dim=1) if normalize_embeddings else features
        if not return_dict:
            return (embeddings, projection) if return_projection else (embeddings,)
        return ComodoOutput(embeddings=embeddings, projection=projection)


ComodoModel.register_for_auto_class("AutoModel")