"""
MONAI-based 3D preprocessing pipeline for VoxTell grounding volumes.

Pipeline (matches VoxTell's nnUNet-style training preprocessing):
  1. EnsureChannelFirstd — add channel dim
  2. Orientationd(RAS)  — standard anatomical orientation
  3. Spacingd((1,1,1))  — isotropic 1 mm voxel spacing
  4. NormalizeIntensityd — z-score normalization (mean/std over all voxels)
  5. SpatialPadd(192)   — ensure volume is at least 192^3
  6. RandCropByPosNegLabeld(192^3, pos=1, neg=1) — foreground-biased crop

Replaces MonaiSAM2Preprocessor for the VoxTell segmentation backbone.
"""

from typing import List, Optional, Tuple

import numpy as np
import torch
from ._compat import print_log

from monai.data import MetaTensor
from monai.transforms import (
    Compose,
    EnsureChannelFirstd,
    NormalizeIntensityd,
    Orientationd,
    RandCropByPosNegLabeld,
    SpatialPadd,
    CenterSpatialCropd,
    Spacingd,
    Transposed,
    SpatialPad
)
from monai.utils import ensure_tuple


class RandCropByPosNegLabelWithRandomOffsetd(RandCropByPosNegLabeld):
    """
    Random Cropping with consistent output spacing and size, using random offset.
    Pipeline: Pad -> RandomCrop (with random offset) -> Resize to target_spacing -> CenterCrop to spatial_size
    """
    
    def randomize(
        self,
        label: torch.Tensor,
        fg_indices: np.ndarray | None = None,
        bg_indices: np.ndarray | None = None,
        image: torch.Tensor | None = None,
    ) -> None:
        # Call parent's randomize to set self.centers
        super().randomize(label, fg_indices, bg_indices, image)

        if self.cropper.centers is None:
            return

        # Vectorized random offset and clamping
        centers = np.array(self.cropper.centers)  # (num_centers, num_dims)
        spatial_size = np.array(ensure_tuple(self.cropper.spatial_size))
        spatial_shape = np.array(label.shape[1:])

        # Generate random offsets for all centers at once: [0, size) - size//2
        offsets = (self.R.rand(*centers.shape) * spatial_size).astype(np.int64) - spatial_size // 2

        # Apply offsets and clamp to valid bounds
        new_centers = centers + offsets
        min_centers = spatial_size // 2
        max_centers = spatial_shape - spatial_size + spatial_size // 2
        new_centers = np.clip(new_centers, min_centers, max_centers).astype(np.int64)

        # Convert back to tuple of tuples (required format)
        self.cropper.centers = tuple(map(tuple, new_centers.tolist()))
        
        
class SmartRandCropByPosNegLabeld(RandCropByPosNegLabelWithRandomOffsetd):
    """
    Random Cropping with consistent output spacing and size and random offset.
    Pipeline: Pad -> RandomCrop -> Resize to target_spacing -> CenterCrop to spatial_size
    """
    
    def __init__(
        self,
        keys,
        label_key,
        spatial_size,
        target_spacing,
        pad_mode="constant",
        resize_mode=('bilinear', 'nearest'),
        **kwargs
    ):
        """
        Args:
            keys: keys of the corresponding items to be transformed.
            label_key: key for label used for positive/negative sampling.
            spatial_size: final output spatial size after all transforms.
            target_spacing: target spacing to resample to (e.g., [1.0, 1.0, 1.0]).
            pad_mode: padding mode for SpatialPadd.
            resize_mode: interpolation mode for Spacingd ("bilinear", "nearest", etc.).
            **kwargs: additional arguments for RandCropByPosNegLabeld.
        """
        self.spatial_size = np.array(spatial_size)
        self.target_spacing = np.array(target_spacing)
        self.pad_mode = pad_mode
        
        # Initialize parent with dummy spatial_size, will be computed per-image
        super().__init__(keys=keys, label_key=label_key, spatial_size=(1, 1, 1), **kwargs)
        
        self.spacer = Spacingd(
            keys=self.keys,
            pixdim=self.target_spacing.tolist(),
            mode=resize_mode
        )
        self.center_cropper = CenterSpatialCropd(keys=self.keys, roi_size=self.spatial_size.tolist())
    
    def _get_crop_size(self, label):
        """Compute crop size in voxels based on physical size and image spacing."""        
        if isinstance(label, MetaTensor):
            spacing = torch.norm(label.affine[:3, :3], dim=0)
        else:
            print_log(f"Image has no affine, using default spacing of 1.0", logger='current')
            spacing = self.target_spacing
        
        # Compute crop voxels: physical_extent / current_spacing
        # physical_extent = spatial_size[i] * target_spacing[i]
        crop_size = [torch.ceil(self.spatial_size[i] * (self.target_spacing[i] / spacing[i])).int().item() for i in range(3)]
        crop_size = [s + 5 for s in crop_size]
        
        return tuple(crop_size)
    
    def __call__(self, data, lazy=None):
        d = dict(data)
        label = d[self.label_key]
        
        # Step 1: Compute crop size based on image spacing
        crop_size = self._get_crop_size(label)
        self.cropper.spatial_size = crop_size # This is used for random offset sampling
        
        # Step 2: Pad if volume is smaller than crop_size
        for key in self.key_iterator(d):
            img = d[key]
            if self.pad_mode == 'min_value' or self.pad_mode == 'constant':
                pad_mode = 'constant'
                
                constant_values = img.min() if self.pad_mode == 'min_value' else 0
                
                padder = SpatialPad(spatial_size=crop_size, mode=pad_mode, constant_values=constant_values)
            else:
                pad_mode = self.pad_mode
                
                padder = SpatialPad(spatial_size=crop_size, mode=pad_mode)
            
            d[key] = padder(img)
        
        # Step 3: Random crop (uses parent's __call__) with random offset
        cropped_list = super().__call__(d, lazy=lazy)
        
        # From list of dicts to dict of lists
        stacked_list = []
        for sample in cropped_list:
            # Step 4: Resample to target spacing
            sample = self.spacer(sample)
            # Step 5: Center crop to exact spatial size
            sample = self.center_cropper(sample)
            stacked_list.append(sample)
            
        return stacked_list
            

class VoxTellPreprocessor:
    """
    3D preprocessing for VoxTell grounding input.

    Takes a raw CT volume + binary mask (any orientation, any spacing)
    and produces one or more 192^3 patches ready for VoxTell's encoder.

    In **eval mode** (set via ``set_eval_mode()``), returns the full
    preprocessed volume (orient → z-score → resample 1mm → transpose)
    without any cropping, for sliding window inference.

    Args:
        patch_size: Spatial size of each output patch.
        pos: Weight for positive (foreground) patches.
        neg: Weight for negative (background) patches.
        num_samples: Number of 192^3 patches to sample per volume.
    """

    def __init__(
        self,
        patch_size: Tuple[int, int, int] = (192, 192, 192),
        pos: float = 1.0,
        neg: float = 1.0,
        num_samples: int = 1,
    ):
        self.patch_size = patch_size
        self.num_samples = num_samples
        self._eval_mode = False

        self.transform = Compose([
            # 1. Add channel dimension: (W,H,D) → (1,W,H,D)
            EnsureChannelFirstd(
                keys=['image', 'label'],
                channel_dim='no_channel',
            ),
            # 2. Reorient to RAS (Right-Anterior-Superior)
            Orientationd(
                keys=['image', 'label'],
                axcodes='RAS',
            ),
            # 3. Z-score normalization over all voxels
            NormalizeIntensityd(
                keys=['image']
            ),
            # 4. Random crop with foreground bias
            SmartRandCropByPosNegLabeld(
                keys=['image', 'label'],
                label_key='label',
                spatial_size=patch_size,
                target_spacing=(1.0, 1.0, 1.0),
                resize_mode=('bilinear', 'nearest'),
                pos=pos,
                neg=neg,
                num_samples=num_samples,
            ),
            # 5. Reorient to RAS (Right-Anterior-Superior)
            Transposed(
                keys=['image', 'label'],
                indices=[0, 3, 2, 1] # CWHD to CDHW
            ),
        ])

        # Eval pipeline: same as training but no cropping.
        # orient → normalize → resample 1mm → transpose
        self.eval_transform = Compose([
            EnsureChannelFirstd(
                keys=['image', 'label'],
                channel_dim='no_channel',
            ),
            Orientationd(
                keys=['image', 'label'],
                axcodes='RAS',
            ),
            NormalizeIntensityd(
                keys=['image'],
            ),
            Spacingd(
                keys=['image', 'label'],
                pixdim=(1.0, 1.0, 1.0),
                mode=('bilinear', 'nearest'),
            ),
            Transposed(
                keys=['image', 'label'],
                indices=[0, 3, 2, 1],  # CWHD to CDHW
            ),
        ])

    def set_eval_mode(self):
        """Switch to eval mode: full-volume preprocessing without cropping."""
        self._eval_mode = True

    def apply_image(
        self,
        volume: np.ndarray,
        mask: np.ndarray,
        affine: Optional[np.ndarray] = None,
        data_source: Optional[str] = None,
    ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
        """
        Preprocess volume + mask and return VoxTell-ready 3D patches.

        In eval mode, returns the full preprocessed volume (no cropping)
        for sliding window inference.  In training mode, returns one or
        more foreground-biased 192^3 patches.

        Args:
            volume: Raw CT volume, any axis order (matching affine).
            mask: Binary segmentation mask, same shape as volume.
            affine: 4x4 affine matrix (required for orientation/spacing).
            data_source: Unused (kept for interface compatibility).

        Returns:
            Tuple of ``(volume_tensor, mask_tensor)`` where shapes are:
              - Training: ``(1, D, H, W)`` with D=H=W=192.
              - Eval: ``(1, D', H', W')`` at 1 mm isotropic spacing.
            If ``num_samples > 1`` (training only), returns a list of
            such tuples.  Returns None on failure.
        """
        try:
            assert affine is not None, 'Affine matrix is required'

            img_tensor = MetaTensor(
                volume.astype(np.float32), affine=affine,
            )
            lbl_tensor = MetaTensor(
                mask.astype(np.float32), affine=affine,
            )

            data_dict = {
                'image': img_tensor,
                'label': lbl_tensor,
            }

            if self._eval_mode:
                result = self.eval_transform(data_dict)
                vol = result['image'].as_tensor()  # (1, D', H', W')
                msk = result['label'].as_tensor()
                return vol, msk

            result = self.transform(data_dict)

            # RandCropByPosNegLabeld returns a list of dicts
            if isinstance(result, list):
                if self.num_samples == 1:
                    r = result[0]
                    vol_patch = r['image'].as_tensor()  # (1, D, H, W)
                    msk_patch = r['label'].as_tensor()
                    return vol_patch, msk_patch
                else:
                    patches = []
                    for r in result:
                        vol_patch = r['image'].as_tensor()
                        msk_patch = r['label'].as_tensor()
                        patches.append((vol_patch, msk_patch))
                    return patches
            else:
                vol_patch = result['image'].as_tensor()
                msk_patch = result['label'].as_tensor()
                return vol_patch, msk_patch

        except Exception as e:
            print_log(
                f'VoxTellPreprocessor error: {e}', logger='current',
            )
            import traceback
            print_log(traceback.format_exc(), logger='current')
            return None