"""Concerto (Pointcept/Concerto @ 5b411171) as a transformers ``PreTrainedModel``.

``ConcertoModel.model`` is the package's ``PointTransformerV3`` (vendored verbatim in ``concerto_model.py``), built
from the checkpoint's stored config the way ``concerto.load`` builds it; ``forward`` runs ``PointTransformerV3.forward``
and returns a ``ConcertoModelOutput``: ``last_hidden_state`` (the output Point's ``feat``) and ``point``, the whole
output Point (an ``addict.Dict``: ``coord``, ``offset``, the pooling_parent / pooling_inverse chain, ...).
Runtime packages beyond torch / transformers: ``spconv``, ``torch_scatter``, ``timm``, ``addict``, ``scipy``
(as the original); ``flash_attn`` when ``enable_flash`` is on.
"""

from dataclasses import dataclass
from typing import Any

import torch
from transformers import PreTrainedModel
from transformers.utils import ModelOutput

from .concerto_model import PointTransformerV3
from .configuration_concerto import ConcertoConfig


@dataclass
class ConcertoModelOutput(ModelOutput):
    """last_hidden_state (M, C): features of the encoder's output points (``point.feat``); point: the output Point, whose
    pooling_parent / pooling_inverse chain and ``inverse`` map features back to the input points."""

    last_hidden_state: torch.Tensor | None = None
    point: Any = None


class ConcertoPreTrainedModel(PreTrainedModel):
    config_class = ConcertoConfig
    config: ConcertoConfig
    base_model_prefix = "model"
    main_input_name = "feat"


class ConcertoModel(ConcertoPreTrainedModel):
    def __init__(self, config: ConcertoConfig):
        super().__init__(config)
        # concerto.load: custom_config entries replace the stored ones
        ptv3_config = dict(config.ptv3_config)
        if config.enable_flash is not None:
            ptv3_config["enable_flash"] = config.enable_flash
        if config.enc_patch_size is not None:
            ptv3_config["enc_patch_size"] = list(config.enc_patch_size)
        self.model = PointTransformerV3(**ptv3_config)
        self.post_init()

    def forward(
        self,
        coord: torch.Tensor | None = None,
        grid_coord: torch.Tensor | None = None,
        feat: torch.Tensor | None = None,
        offset: torch.Tensor | None = None,
        batch: torch.Tensor | None = None,
        return_dict: bool | None = None,
        **kwargs,
    ) -> tuple | ConcertoModelOutput:
        """A collected point cloud (``ConcertoProcessor``): coord / grid_coord (N, 3), feat (N, C_in), and offset or batch;
        other keys the processor returns (color, inverse, ...) are kept in the Point as the original pipeline does.
        ``return_dict=False`` returns a tuple."""
        return_dict = return_dict if return_dict is not None else self.config.return_dict
        named = {"coord": coord, "grid_coord": grid_coord, "feat": feat, "offset": offset, "batch": batch}
        point = self.model({**{k: v for k, v in named.items() if v is not None}, **kwargs})
        output = ConcertoModelOutput(last_hidden_state=point.feat, point=point)
        return output if return_dict else output.to_tuple()


__all__ = ["ConcertoModel", "ConcertoModelOutput", "ConcertoPreTrainedModel"]