"""VoxTell 3D grounding encoder for standalone RepConTUR inference.

Adapted from the training-time VoxTellTrainRunner: the architecture is built
from an explicit ``arch_kwargs`` dict (stored in the HF config) instead of a
``plans.json`` file, and no checkpoint is loaded at construction — weights come
from the RepConTUR safetensors.
"""

import pydoc
from typing import Tuple

import torch
from torch import nn
from einops import rearrange

from .voxtell_model import VoxTellModel


class VoxTellGroundingEncoder(nn.Module):
    """VoxTell ResidualEncoder + text-prompted transformer decoder.

    The LLM ``[SEG]`` embeddings (projected to ``query_dim`` by the parent's
    ``text_hidden_fcs``) are used directly as decoder queries, bypassing
    VoxTell's original text embedding model.
    """

    def __init__(
        self,
        arch_kwargs: dict,
        kw_requires_import=('conv_op', 'norm_op', 'dropout_op', 'nonlin'),
        decoder_layer: int = 4,
        text_embedding_dim: int = 2560,
        num_maskformer_stages: int = 5,
        num_heads: int = 32,
        query_dim: int = 2048,
    ):
        super().__init__()
        self.decoder_layer = decoder_layer
        self.hidden_dim = query_dim

        resolved = dict(arch_kwargs)
        for key in kw_requires_import:
            if resolved.get(key) is not None:
                resolved[key] = pydoc.locate(resolved[key])

        self.voxtell_model = VoxTellModel(
            input_channels=1,
            **resolved,
            decoder_layer=decoder_layer,
            text_embedding_dim=text_embedding_dim,
            num_maskformer_stages=num_maskformer_stages,
            num_heads=num_heads,
            query_dim=query_dim,
            project_to_decoder_hidden_dim=query_dim,
            deep_supervision=False,
        )

        self.voxtell_model.encoder.requires_grad_(False)
        self.voxtell_model.project_text_embed.requires_grad_(False)

    def forward_with_language(
        self,
        volume: torch.Tensor,
        language_embd: torch.Tensor,
        num_objs: int = 1,
    ) -> torch.Tensor:
        """Encode a 3D volume and decode with language queries.

        Args:
            volume: (B, 1, D, H, W) preprocessed volume.
            language_embd: (B * num_objs, 1, query_dim) projected LLM embeddings.
            num_objs: objects per sample.
        Returns:
            (B * num_objs, D, H, W) segmentation logits.
        """
        batch_size = volume.shape[0]

        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
            skips = self.voxtell_model.encoder(volume)

            selected_feature = skips[self.decoder_layer]
            bottleneck_embed = rearrange(selected_feature, 'b c d h w -> b h w d c')
            bottleneck_embed = self.voxtell_model.project_bottleneck_embed(bottleneck_embed)
            bottleneck_embed = rearrange(bottleneck_embed, 'b h w d c -> (h w d) b c')

            text_embed = language_embd.squeeze(1)
            text_embed = text_embed.reshape(batch_size, num_objs, -1)
            text_embed = rearrange(text_embed, 'b n d -> n b d')

            mask_embedding, _ = self.voxtell_model.transformer_decoder(
                tgt=text_embed,
                memory=bottleneck_embed,
                pos=self.voxtell_model.pos_embed,
                memory_key_padding_mask=None,
            )
            mask_embedding = rearrange(mask_embedding, 'n b d -> b n d')

            mask_embeddings = [
                proj(mask_embedding)
                for proj in self.voxtell_model.project_to_decoder_channels
            ]

            outs = []
            for prompt_idx in range(num_objs):
                prompt_embeds = [m[:, prompt_idx:prompt_idx + 1] for m in mask_embeddings]
                outs.append(self.voxtell_model.decoder(skips, prompt_embeds))

            outs = [torch.cat(scale_outs, dim=1) for scale_outs in zip(*outs)]
            pred_masks = outs[0]

        pred_masks = pred_masks.reshape(batch_size * num_objs, *pred_masks.shape[2:])
        return pred_masks

    @torch.inference_mode()
    def predict_sliding_window(
        self,
        volume: torch.Tensor,
        language_embd: torch.Tensor,
        num_objs: int = 1,
        tile_step_size: float = 0.5,
        patch_size: Tuple[int, int, int] = (192, 192, 192),
    ) -> torch.Tensor:
        """Sliding-window inference over a full volume with Gaussian stitching.

        Args:
            volume: (1, 1, D, H, W) full preprocessed volume.
            language_embd: (num_objs, 1, query_dim) shared language queries.
            num_objs: number of objects.
            tile_step_size: step fraction of patch size.
            patch_size: inference patch size.
        Returns:
            (num_objs, D, H, W) logits at full resolution.
        """
        from acvl_utils.cropping_and_padding.padding import pad_nd_image
        from nnunetv2.inference.sliding_window_prediction import (
            compute_gaussian,
            compute_steps_for_sliding_window,
        )

        device = volume.device
        data = volume.squeeze(0)

        data, slicer_revert = pad_nd_image(
            data, patch_size, 'constant', {'value': 0}, True, None,
        )
        spatial_shape = data.shape[1:]

        steps = compute_steps_for_sliding_window(spatial_shape, patch_size, tile_step_size)
        slicers = []
        for sx in steps[0]:
            for sy in steps[1]:
                for sz in steps[2]:
                    slicers.append((
                        slice(None),
                        slice(sx, sx + patch_size[0]),
                        slice(sy, sy + patch_size[1]),
                        slice(sz, sz + patch_size[2]),
                    ))

        predicted_logits = torch.zeros(
            (num_objs, *spatial_shape), dtype=torch.float32, device=device,
        )
        n_predictions = torch.zeros(spatial_shape, dtype=torch.float32, device=device)

        gaussian = compute_gaussian(
            tuple(patch_size), sigma_scale=1.0 / 8,
            value_scaling_factor=10, device=device,
        ).float()

        for slicer in slicers:
            patch = data[slicer].unsqueeze(0).contiguous()
            pred = self.forward_with_language(patch, language_embd, num_objs)
            pred = pred.float() * gaussian

            spatial_slicer = slicer[1:]
            out_slicer = (slice(None), *spatial_slicer)
            predicted_logits[out_slicer] += pred
            n_predictions[spatial_slicer] += gaussian

        predicted_logits /= n_predictions
        predicted_logits = predicted_logits[(slice(None), *slicer_revert[1:])]
        return predicted_logits