"""Trace Anything (ByteDance-Seed/TraceAnything @ 54677b5e) as a transformers model.

``TraceAnythingModel`` holds the submodules of the official ``TraceAnything`` under the same names (``encoder``: CroCo
ViT, ``decoder``: the trace decoder, ``ds_head_time`` / ``ds_head_track``: DPT heads; the official ``time_head`` /
``track_head`` are aliases of the latter two), so the released ``trace_anything.pt`` tensors load by name, and runs the
official ``TraceAnything.forward`` on the views built from ``pixel_values`` / ``time_steps``. The per-view result dicts
of the original are returned as ``predictions`` and stacked over the views in ``TraceAnythingOutput``.
"""

from dataclasses import dataclass

import numpy as np
import torch
from transformers import PreTrainedModel
from transformers.utils import ModelOutput

from .configuration_trace_anything import TraceAnythingConfig
from .ta_layers import get_1d_sincos_pos_embed_from_grid
from .ta_model import TraceAnything, TraceDecoder


@dataclass
class TraceAnythingOutput(ModelOutput):
    """Per view i of N (stacked from the original's per-view dicts):
    time (N, 1): the predicted time of each view; ctrl_pts3d (N, K, H, W, 3) / ctrl_conf (N, K, H, W): the K B-spline
    control points of each pixel's 3D trajectory and their confidence; track_pts3d (N, N, H, W, 3) / track_conf
    (N, N, H, W): the trajectories evaluated at the N predicted times; predictions: the original list of dicts."""

    time: torch.Tensor | None = None
    ctrl_pts3d: torch.Tensor | None = None
    ctrl_conf: torch.Tensor | None = None
    track_pts3d: torch.Tensor | None = None
    track_conf: torch.Tensor | None = None
    predictions: list | None = None


class TraceAnythingPreTrainedModel(PreTrainedModel):
    config_class = TraceAnythingConfig
    config: TraceAnythingConfig
    base_model_prefix = "trace_anything"
    main_input_name = "pixel_values"

    def _init_weights(self, module):
        super()._init_weights(module)
        # TraceDecoder.image_idx_emb is a non-persistent buffer: not in checkpoints, left uninitialised by meta-device
        # loading -> rebuild it as TraceDecoder.__init__ does
        if isinstance(module, TraceDecoder):
            embed_dim = module.image_idx_emb.shape[1]
            with torch.no_grad():
                module.image_idx_emb.copy_(torch.from_numpy(get_1d_sincos_pos_embed_from_grid(embed_dim, np.arange(1000))).float())


class TraceAnythingModel(TraceAnythingPreTrainedModel):
    _encode_images = TraceAnything._encode_images

    def __init__(self, config: TraceAnythingConfig):
        super().__init__(config)
        net = TraceAnything(encoder_args=config.encoder_args, decoder_args=config.decoder_args, head_args=config.head_args,
                            targeting_mechanism=config.targeting_mechanism, poly_degree=config.poly_degree, whether_local=False)
        if net.whether_local:
            raise ValueError("head_args.with_local_head must be False: no released checkpoint has local-head weights")
        # the official module's attributes, so its methods run on this object unchanged
        self.targeting_mechanism = net.targeting_mechanism
        self.poly_degree = net.poly_degree
        self.whether_local = net.whether_local
        self.max_parallel_views_for_head = net.max_parallel_views_for_head
        self.encoder = net.encoder
        self.decoder = net.decoder
        self.ds_head_time = net.ds_head_time
        self.ds_head_track = net.ds_head_track
        # the DPT heads register their scratch convolutions twice (layer1_rn ... and the layer_rn list): declare those
        # names as tied (target -> first name), so save_pretrained stores each tensor once and from_pretrained keeps the sharing
        first, tied = {}, {}
        for name, parameter in self.named_parameters(remove_duplicate=False):
            if id(parameter) in first:
                tied[name] = first[id(parameter)]
            else:
                first[id(parameter)] = name
        self._tied_weights_keys = tied or None
        self.post_init()

    # the official aliases (self.time_head = self.ds_head_time, ...), not registered twice
    @property
    def time_head(self):
        return self.ds_head_time

    @property
    def track_head(self):
        return self.ds_head_track

    def forward(self, pixel_values: torch.Tensor, time_steps=None, return_dict: bool | None = None, **kwargs) -> tuple | TraceAnythingOutput:
        """pixel_values (1, N, 3, H, W) in [-1, 1], H and W multiples of 16; time_steps (1, N): the time of each frame
        (default i / (N - 1)); both from ``TraceAnythingProcessor``. One video per call, as the original."""
        return_dict = return_dict if return_dict is not None else self.config.return_dict
        if pixel_values.shape[0] != 1:
            raise ValueError("one video per call: pixel_values must be (1, N, 3, H, W)")
        num_views = pixel_values.shape[1]
        if time_steps is None:
            times = [i / (num_views - 1) if num_views > 1 else 0.0 for i in range(num_views)]
        else:
            times = [float(t) for t in (time_steps[0].tolist() if isinstance(time_steps, torch.Tensor) else time_steps[0])]
        views = [{"img": pixel_values[:, i], "time_step": times[i]} for i in range(num_views)]
        predictions = TraceAnything.forward(self, views)
        output = TraceAnythingOutput(
            time=torch.stack([p["time"] for p in predictions]),
            ctrl_pts3d=torch.stack([p["ctrl_pts3d"] for p in predictions]),
            ctrl_conf=torch.stack([p["ctrl_conf"] for p in predictions]),
            track_pts3d=torch.stack([torch.cat(p["track_pts3d"]) for p in predictions]),
            track_conf=torch.stack([torch.cat(p["track_conf"]) for p in predictions]),
            predictions=predictions,
        )
        return output if return_dict else output.to_tuple()


__all__ = ["TraceAnythingModel", "TraceAnythingOutput", "TraceAnythingPreTrainedModel"]