"""
CT Volume Encoder Wrapper for RepConTUR.

Wraps CTViT or NNUNet3DEncoder to encode full 3D CT volumes into
a flattened token sequence of shape (B, N, dim).

The encoder is loaded from a CT-CLIP checkpoint and fully frozen.
"""

from pathlib import Path
from typing import Optional, List

import torch
import torch.nn as nn
from ._compat import print_log


class CTVolumeEncoder(nn.Module):
    """
    Wraps CTViT or NNUNet3DEncoder for 3D CT volume encoding.

    Input:  (B, 1, D, H, W) — full 3D CT volume
    Output: (B, T*H'*W', dim) — flattened spatial/temporal tokens

    The encoder is always frozen. Only downstream modules (attentional
    pooler, MLP projection) are trainable.
    """

    def __init__(
        self,
        encoder_type: str = 'ctvit',
        checkpoint_path: Optional[str] = None,
        ct_clip_source_path: Optional[str] = None,
        # CTViT params
        dim: int = 512,
        codebook_size: int = 8192,
        image_size: int = 360,
        patch_size: int = 20,
        temporal_patch_size: int = 10,
        spatial_depth: int = 4,
        temporal_depth: int = 4,
        dim_head: int = 32,
        heads: int = 8,
        # NNUNet params
        nnunet_depth: int = 5,
        nnunet_channels: Optional[List[int]] = None,
        nnunet_stride_list: Optional[List[List[int]]] = None,
    ):
        super().__init__()
        self.encoder_type = encoder_type
        self.dim = dim

        if nnunet_channels is None:
            nnunet_channels = [32, 64, 128, 256, 320, 512]
        if nnunet_stride_list is None:
            nnunet_stride_list = [[1, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]]

        # transformer_maskgit is vendored alongside this module (ct_clip_source_path ignored).
        if encoder_type == 'ctvit':
            from .ctvit import CTViT
            self.encoder = CTViT(
                dim=dim,
                codebook_size=codebook_size,
                image_size=image_size,
                patch_size=patch_size,
                temporal_patch_size=temporal_patch_size,
                spatial_depth=spatial_depth,
                temporal_depth=temporal_depth,
                dim_head=dim_head,
                heads=heads,
            )
        elif encoder_type == 'nnunet':
            from .nnunet import NNUNet3DEncoder
            self.encoder = NNUNet3DEncoder(
                depth=nnunet_depth,
                channels=nnunet_channels,
                stride_list=nnunet_stride_list,
            )
            self.dim = nnunet_channels[nnunet_depth]  # output channel dim
        else:
            raise ValueError(f"Unknown encoder_type: {encoder_type}. Use 'ctvit' or 'nnunet'.")

        # Load checkpoint
        if checkpoint_path is not None:
            self._load_checkpoint(checkpoint_path)

        # Freeze the entire encoder
        self.encoder.requires_grad_(False)
        self.encoder.eval()

        total_params = sum(p.numel() for p in self.encoder.parameters())
        print_log(
            f"CTVolumeEncoder ({encoder_type}): loaded with {total_params:,} params (all frozen)",
            logger='current'
        )

    def _load_checkpoint(self, checkpoint_path: str):
        """Load encoder weights from a CT-CLIP checkpoint."""
        path = Path(checkpoint_path)
        if not path.exists():
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

        if path.suffix == '.safetensors':
            from safetensors import safe_open
            f = safe_open(str(path), framework='pt', device='cpu')
            all_keys = list(f.keys())

            encoder_state = {}
            for prefix in ['ctclip.visual_transformer.', 'visual_transformer.']:
                for k in all_keys:
                    if k.startswith(prefix):
                        encoder_state[k[len(prefix):]] = f.get_tensor(k)
                if encoder_state:
                    break

            if not encoder_state:
                for k in all_keys:
                    encoder_state[k] = f.get_tensor(k)
        else:
            # Standard PyTorch checkpoint
            ckpt = torch.load(str(path), map_location='cpu')
            if isinstance(ckpt, dict) and 'state_dict' in ckpt:
                ckpt = ckpt['state_dict']

            encoder_state = {}
            for prefix in ['ctclip.visual_transformer.', 'visual_transformer.']:
                for k, v in ckpt.items():
                    if k.startswith(prefix):
                        encoder_state[k[len(prefix):]] = v
                if encoder_state:
                    break

            if not encoder_state:
                encoder_state = ckpt

        # Load with strict=False to skip decoder keys (to_pixels*, dec_*)
        missing, unexpected = self.encoder.load_state_dict(encoder_state, strict=False)
        print_log(
            f"CTVolumeEncoder: loaded checkpoint from {checkpoint_path}\n"
            f"  Matched keys: {len(encoder_state) - len(unexpected)}\n"
            f"  Missing keys (expected for encoder-only): {len(missing)}\n"
            f"  Unexpected keys (decoder, skipped): {len(unexpected)}",
            logger='current'
        )

    def train(self, mode=True):
        """Override train to keep encoder always in eval mode."""
        super().train(mode)
        self.encoder.eval()
        return self

    def forward(self, volume: torch.Tensor) -> torch.Tensor:
        """
        Args:
            volume: (B, 1, D, H, W) — preprocessed 3D CT volume

        Returns:
            tokens: (B, N, dim) — flattened tokens where N = T*H'*W'
                    For CTViT with (180, 360, 360): N = 18*18*18 = 5832, dim = 512
                    For NNUNet with (180, 360, 360): N varies, dim = channels[-1]
        """

        # Query the encoder's actual parameter dtype at runtime.
        # .to(bf16) in __init__ may not persist (DeepSpeed can reset
        # frozen-parameter dtypes), so we always match dynamically.
        encoder_dtype = next(self.encoder.parameters()).dtype
        volume = volume.to(dtype=encoder_dtype)

        with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
            tokens = self.encoder(volume, return_encoded_tokens=True)
            # CTViT: (B, T, H', W', dim) e.g. (B, 18, 18, 18, 512)
            # NNUNet: (B, D', H', W', dim) e.g. (B, 11, 11, 11, 512)

        B = tokens.shape[0]
        tokens = tokens.reshape(B, -1, self.dim)  # (B, N, dim)
        return tokens