"""
GPU-Optimized NLF SMPLX Processor for H200 Cluster

This module provides a highly optimized pipeline for extracting SMPLX parameters
from video frames using NLF model with YOLO person detection.

Architecture:
    Video Frames (GPU) -> YOLO Person Detection -> Person Crop Preprocessing -> NLF Inference -> Output

Key Features:
    - Full GPU pipeline with minimal CPU transfers
    - CUDA streams for overlapped execution
    - torch.compile optimization for model inference
    - Memory-efficient batched processing
    - Support for multi-worker processing

Author: Optimized for sign language video processing pipeline
Target Hardware: NVIDIA H200 (141GB HBM3)
"""

import warnings
# Suppress YOLO false positive warnings
warnings.filterwarnings('ignore', message='.*torch.Tensor inputs should be normalized.*')
warnings.filterwarnings('ignore', message='.*inference results will accumulate in RAM.*')

import torch
import torch.nn.functional as F
import torchvision  # Required for torchvision ops in TorchScript model
import numpy as np
from typing import Dict, List, Optional, Tuple, Union
from dataclasses import dataclass, field


@dataclass
class NLFConfig:
    """Configuration for NLF processing pipeline."""

    # Model paths
    nlf_model_path: str = "models/nlf_l_multi_0.3.2.torchscript"
    smplx_model_path: str = "models/SMPLX_NEUTRAL.npz"
    yolo_model_path: str = "pretrained_models/yolov8x.pt"

    # Detection parameters
    detection_confidence: float = 0.05
    detector_name: str = "yolov8x"

    # Batch sizes (tuned for H200)
    img_batch_size: int = 64        # Frames to decode at once (increased)
    det_batch_size: int = 64        # Person detections for NLF at once (increased)
    person_prep_chunk_size: int = 64  # Person preprocessing chunk size (increased)

    # Performance options
    use_compile: bool = False       # Use torch.compile for model
    use_amp: bool = True           # Use automatic mixed precision
    compile_mode: str = "max-autotune"  # torch.compile mode

    # YOLO specific
    yolo_img_size: int = 640       # YOLO input resolution

    # Model parameters
    fov: float = 55.0
    num_betas: int = 10
    num_expressions: int = 0


class NLFProcessorOptimized:
    """
    Optimized NLF SMPLX processor for H200 GPU.

    This class provides a complete pipeline for:
    1. Person detection using YOLO
    2. Person crop preprocessing with GPU-accelerated operations
    3. SMPLX parameter estimation using NLF
    4. Minimal CPU-GPU transfers

    All operations are performed on GPU with minimal CPU transfers.

    Example:
        >>> config = NLFConfig()
        >>> processor = NLFProcessorOptimized(config, device)
        >>>
        >>> # Process frames
        >>> frames_gpu = torch.from_numpy(frames).to(device)
        >>> results = processor.process_frames(frames_gpu, base_frame_idx=0)

    Attributes:
        device: CUDA device for processing
        config: NLFConfig with processing parameters
        nlf_model: NLF SMPLX estimation model
        detector: YOLO person detection model
        person_preprocessor: GPU-accelerated person crop preprocessor
    """

    def __init__(
        self,
        config: NLFConfig,
        device: torch.device,
    ):
        """
        Initialize the optimized NLF processor.

        Args:
            config: NLFConfig with paths and parameters
            device: Target CUDA device
        """
        self.config = config
        self.device = device

        # Set CUDA device
        torch.cuda.set_device(device)

        # Load NLF model
        self._load_nlf_model()

        # Load YOLO detector
        self._load_detector()

        # Create CUDA streams for overlapped execution
        self.detection_stream = torch.cuda.Stream(device=device)
        self.inference_stream = torch.cuda.Stream(device=device)

        print(f"[NLF] Processor initialized on {device}")
        print(f"[NLF] Batch sizes: img={config.img_batch_size}, det={config.det_batch_size}")

    def _load_nlf_model(self):
        """Load NLF model from TorchScript checkpoint."""
        self.nlf_model = torch.jit.load(self.config.nlf_model_path).to(self.device).eval()

        # Optionally compile model for faster inference
        if self.config.use_compile and hasattr(torch, 'compile'):
            try:
                self.nlf_model = torch.compile(
                    self.nlf_model,
                    mode=self.config.compile_mode,
                    fullgraph=False,  # Allow graph breaks for compatibility
                )
                print(f"[NLF] Model compiled with mode={self.config.compile_mode}")
            except Exception as e:
                warnings.warn(f"torch.compile failed, using eager mode: {e}")

    def _load_detector(self):
        """Load and configure YOLO detector for person detection."""
        from ultralytics import YOLO

        # Load YOLO with explicit device specification
        device_str = str(self.device).replace('cuda:', '')  # YOLO wants '0' not 'cuda:0'
        self.detector = YOLO(self.config.yolo_model_path)
        self.detector.to(self.device)

        # Force model to GPU (newer ultralytics may need this)
        if hasattr(self.detector, 'model'):
            self.detector.model = self.detector.model.to(self.device)

        print(f"[NLF] YOLO device: {self.detector.device if hasattr(self.detector, 'device') else 'unknown'}")

        # Warmup detector
        # YOLO expects: BCHW, float32, [0,1] normalized
        dummy = torch.zeros(
            1, 3, self.config.yolo_img_size, self.config.yolo_img_size,
            dtype=torch.float32, device=self.device
        )
        with torch.no_grad():
            _ = self.detector(dummy, verbose=False, device=self.device.index if hasattr(self.device, 'index') else 0)

    @torch.no_grad()
    def detect_persons(
        self,
        frames: torch.Tensor,
        base_frame_idx: int = 0,
    ) -> Tuple[Optional[torch.Tensor], ...]:
        """
        Detect persons in frames using YOLO.

        Args:
            frames: (N, H, W, 3) uint8 tensor on GPU (NHWC format from video decode)
            base_frame_idx: Starting frame index (used for final output, not internal indexing)

        Returns:
            Tuple of:
                - boxes: (M, 4) float32 tensor of [x1, y1, x2, y2]
                - img_indices: (M,) int64 tensor - RELATIVE indices (0 to N-1) within this batch
                - confidences: (M,) float32 tensor of detection confidences
            All tensors are on GPU. Returns (None, None, None) if no detections.
        """
        N, H, W, C = frames.shape

        # Convert NHWC uint8 -> NCHW float32 normalized
        # IMPORTANT: .contiguous() needed after permute to avoid view() errors downstream
        frames_nchw = (
            frames
            .permute(0, 3, 1, 2)        # NHWC -> NCHW (changes strides, not contiguous)
            .contiguous()               # Make contiguous in memory
            .to(dtype=torch.float32)    # uint8 -> float32
            .div_(255.0)                # [0,255] -> [0,1]
        )

        # YOLO requires dimensions divisible by stride (32)
        # Apply letterbox padding on GPU
        frames_padded, scale, pad_w, pad_h = self._letterbox_gpu(
            frames_nchw,
            target_size=self.config.yolo_img_size,
            stride=32,
        )

        # Run detection
        with torch.cuda.stream(self.detection_stream):
            detections = self.detector(
                frames_padded,
                conf=self.config.detection_confidence,
                classes=[0],  # Person class only
                verbose=False,
                device=self.device.index if hasattr(self.device, 'index') else 0,
            )

        # Synchronize detection stream
        self.detection_stream.synchronize()

        # Collect all detections and scale boxes back to original coordinates
        all_boxes = []
        all_img_indices = []
        all_confidences = []

        for i, det in enumerate(detections):
            if det.boxes is None or len(det.boxes) == 0:
                continue

            # Extract boxes, classes, and confidences (already on GPU)
            boxes = det.boxes.xyxy.clone()  # (n_det, 4)
            confs = det.boxes.conf          # (n_det,)

            # Scale boxes back to original image coordinates
            boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / scale  # x coordinates
            boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / scale  # y coordinates

            # Clamp to valid image bounds
            boxes[:, [0, 2]] = boxes[:, [0, 2]].clamp(0, W)
            boxes[:, [1, 3]] = boxes[:, [1, 3]].clamp(0, H)

            n_det = boxes.shape[0]

            all_boxes.append(boxes)
            # IMPORTANT: Use RELATIVE index (i) not absolute (base_frame_idx + i)
            all_img_indices.append(
                torch.full((n_det,), i, dtype=torch.int64, device=self.device)
            )
            all_confidences.append(confs)

        # Handle no detections case
        if len(all_boxes) == 0:
            return None, None, None

        # Concatenate all detections
        boxes = torch.cat(all_boxes, dim=0)
        img_indices = torch.cat(all_img_indices, dim=0)
        confidences = torch.cat(all_confidences, dim=0)

        return boxes, img_indices, confidences

    @torch.no_grad()
    def _letterbox_gpu(
        self,
        frames: torch.Tensor,
        target_size: int = 640,
        stride: int = 32,
    ) -> Tuple[torch.Tensor, float, float, float]:
        """
        GPU-native letterbox resize/pad for YOLO.

        Resizes frames to fit within target_size while maintaining aspect ratio,
        then pads to make dimensions divisible by stride.

        Args:
            frames: (N, C, H, W) float32 tensor
            target_size: Target size for longest edge
            stride: Stride alignment requirement (YOLO uses 32)

        Returns:
            Tuple of:
                - padded: (N, C, new_H, new_W) padded tensor
                - scale: Scale factor applied
                - pad_w: Padding added to width (left side)
                - pad_h: Padding added to height (top side)
        """
        N, C, H, W = frames.shape

        # Calculate scale to fit in target_size
        scale = min(target_size / H, target_size / W)

        # New dimensions after scaling
        new_h = int(round(H * scale))
        new_w = int(round(W * scale))

        # Resize using bilinear interpolation (GPU-native)
        if scale != 1.0:
            frames_resized = F.interpolate(
                frames,
                size=(new_h, new_w),
                mode='bilinear',
                align_corners=False,
            )
        else:
            frames_resized = frames

        # Calculate padding to make dimensions divisible by stride
        pad_h_total = (stride - new_h % stride) % stride
        pad_w_total = (stride - new_w % stride) % stride

        # Divide padding evenly (top/bottom, left/right)
        pad_top = pad_h_total // 2
        pad_bottom = pad_h_total - pad_top
        pad_left = pad_w_total // 2
        pad_right = pad_w_total - pad_left

        # Apply padding (F.pad uses reverse order: left, right, top, bottom)
        if pad_h_total > 0 or pad_w_total > 0:
            frames_padded = F.pad(
                frames_resized,
                (pad_left, pad_right, pad_top, pad_bottom),
                mode='constant',
                value=0.5,  # YOLO uses gray padding
            )
        else:
            frames_padded = frames_resized

        return frames_padded, scale, pad_left, pad_top

    @torch.no_grad()
    def process_detections(
        self,
        frames: torch.Tensor,
        boxes: torch.Tensor,
        img_indices: torch.Tensor,
        confidences: torch.Tensor,
        base_frame_idx: int = 0,
    ) -> Dict[str, Union[np.ndarray, Dict]]:
        """
        Process detected persons through NLF pipeline.

        Args:
            frames: (N, H, W, 3) uint8 tensor on GPU
            boxes: (M, 4) bounding boxes
            img_indices: (M,) RELATIVE source frame indices (0 to N-1 within batch)
            confidences: (M,) detection confidences
            base_frame_idx: Base frame index to add for absolute frame numbering

        Returns:
            Dictionary with SMPLX estimation results:
                - img_idx: (M,) int32 - ABSOLUTE source frame indices
                - boxes: (M, 5) float32 - detection boxes with confidence
                - joints3d: (M, 55, 3) float32 - 3D joints
                - joints2d: (M, 55, 2) float32 - 2D joints
                - joint_uncertainties: (M, 55) float32 - uncertainties
                - smplx: dict with pose parameters
        """
        num_detections = boxes.shape[0]
        det_batch_size = self.config.det_batch_size

        # Pre-process all detections using GPU person preprocessor
        # img_indices here are RELATIVE (0 to N-1) for correct indexing into frames
        preprocessed = self.person_preprocessor(
            frames, boxes, img_indices, confidences
        )

        # Output accumulators (GPU tensors)
        all_boxes_gpu = []
        all_confidences_gpu = []
        all_joints3d_gpu = []
        all_joints2d_gpu = []
        all_joint_uncertainties_gpu = []
        all_img_idx_gpu = []
        all_smplx_params = {
            'global_orient': [],
            'body_pose': [],
            'betas': [],
            'trans': [],
            'left_hand_pose': [],
            'right_hand_pose': [],
            'jaw_pose': [],
            'leye_pose': [],
            'reye_pose': [],
        }

        # Process in batches for memory efficiency
        for batch_start in range(0, num_detections, det_batch_size):
            batch_end = min(batch_start + det_batch_size, num_detections)

            # Slice preprocessed batch
            batch_imgs = preprocessed['img'][batch_start:batch_end]


            # Run NLF inference with AMP
            with torch.cuda.stream(self.inference_stream):
                try:
                    if self.config.use_amp:
                        with torch.cuda.amp.autocast():
                            # NLF expects (N, 3, H, W) float32 normalized images
                            preds = self.nlf_model.detect_smpl_batched(batch_imgs, model_name='smplx')
                    else:
                        preds = self.nlf_model.detect_smpl_batched(batch_imgs, model_name='smplx')
                except Exception as e:
                    print(f"[ERROR] NLF inference failed on batch {batch_start}-{batch_end}: {type(e).__name__}: {e}")
                    import traceback
                    traceback.print_exc()
                    return None

            # Wait for inference to complete
            self.inference_stream.synchronize()

            # Accumulate results (stay on GPU)
            # NLF returns dict with keys: pose, betas, trans, joints3d, joints2d, joint_uncertainties, etc.
            # Note: boxes come from YOLO detection, not from NLF
            # NLF returns lists with batch dimension: pred[key][0] contains the actual batch data
            if 'pose' in preds and 'joints3d' in preds:
                # Use YOLO detection boxes and confidences (passed as parameters, sliced for this batch)
                batch_boxes = boxes[batch_start:batch_end]
                batch_confidences = confidences[batch_start:batch_end]
                all_boxes_gpu.append(batch_boxes)
                all_confidences_gpu.append(batch_confidences)

                # Extract NLF outputs - unwrap batch dimension [0]
                all_joints3d_gpu.append(preds['joints3d'][0])
                all_joints2d_gpu.append(preds['joints2d'][0])
                all_joint_uncertainties_gpu.append(preds['joint_uncertainties'][0])

                # Image indices (relative, will convert to absolute later)
                batch_img_idx = preprocessed['img_idx'][batch_start:batch_end]
                all_img_idx_gpu.append(batch_img_idx)

                # SMPLX parameters - parse pose tensor
                pose = preds['pose'][0]  # Unwrap batch: [0] -> (N, 165) - full pose
                all_smplx_params['global_orient'].append(pose[:, :3])
                all_smplx_params['body_pose'].append(pose[:, 3:66])
                all_smplx_params['jaw_pose'].append(pose[:, 66:69])
                all_smplx_params['leye_pose'].append(pose[:, 69:72])
                all_smplx_params['reye_pose'].append(pose[:, 72:75])
                all_smplx_params['left_hand_pose'].append(pose[:, 75:120])
                all_smplx_params['right_hand_pose'].append(pose[:, 120:165])

                all_smplx_params['betas'].append(preds['betas'][0])
                all_smplx_params['trans'].append(preds['trans'][0])

        # Handle case of no detections
        if len(all_boxes_gpu) == 0:
            return None

        # Concatenate all batches (on GPU)
        boxes_gpu = torch.cat(all_boxes_gpu, dim=0)
        confidences_gpu = torch.cat(all_confidences_gpu, dim=0)
        joints3d_gpu = torch.cat(all_joints3d_gpu, dim=0)
        joints2d_gpu = torch.cat(all_joints2d_gpu, dim=0)
        joint_uncertainties_gpu = torch.cat(all_joint_uncertainties_gpu, dim=0)
        img_idx_gpu = torch.cat(all_img_idx_gpu, dim=0)

        smplx_params_gpu = {
            key: torch.cat(vals, dim=0)
            for key, vals in all_smplx_params.items()
        }

        # Convert RELATIVE img_idx to ABSOLUTE by adding base_frame_idx
        # Single transfer to CPU at end
        img_idx_absolute = img_idx_gpu.cpu().numpy().astype(np.int32) + base_frame_idx

        # Add confidence scores to boxes
        boxes_with_conf = torch.cat([
            boxes_gpu,
            confidences_gpu.unsqueeze(1)
        ], dim=1)

        return {
            'img_idx': img_idx_absolute,
            'boxes': boxes_with_conf.cpu().numpy().astype(np.float32),
            'joints3d': joints3d_gpu.cpu().numpy().astype(np.float32),
            'joints2d': joints2d_gpu.cpu().numpy().astype(np.float32),
            'joint_uncertainties': joint_uncertainties_gpu.cpu().numpy().astype(np.float32),
            'smplx': {
                key: val.cpu().numpy().astype(np.float32)
                for key, val in smplx_params_gpu.items()
            },
        }

    @torch.no_grad()
    def process_frames(
        self,
        frames: torch.Tensor,
        base_frame_idx: int = 0,
    ) -> Optional[Dict[str, Union[np.ndarray, Dict]]]:
        """
        Full pipeline: Send frames to NLF which does detection + SMPLX estimation internally.

        OPTIMIZED: Batch multiple frames together to reduce overhead.

        This is the main entry point for processing video frames.

        Args:
            frames: (N, H, W, 3) uint8 tensor on GPU
            base_frame_idx: Starting frame index for output img_idx

        Returns:
            Dictionary with all SMPLX estimation results, or None if no persons detected.
        """
        # Output accumulators
        raw_outputs = {'boxes': [], 'joints3d': [], 'joints2d': [], 'joint_uncertainties': [],
                      'pose': [], 'betas': [], 'trans': [], 'img_idx': []}

        num_frames = frames.shape[0]

        # OPTIMIZATION: Batch multiple frames together (reduces model call overhead)
        # NLF's detect_smpl_batched can process multiple frames at once
        # RTX 3090 (24GB) can handle ~32 frames, H200 can handle more
        nlf_batch_size = min(32, num_frames)  # Process as many frames as possible

        for batch_start in range(0, num_frames, nlf_batch_size):
            batch_end = min(batch_start + nlf_batch_size, num_frames)
            batch_frames = frames[batch_start:batch_end]

            # Convert batch to NLF format: (B, H, W, 3) uint8 → (B, 3, H, W) float32
            batch_tensor = batch_frames.permute(0, 3, 1, 2).float() / 255.0

            # Run NLF inference on batch (NLF doesn't support AMP, needs float32)
            with torch.inference_mode():
                try:
                    # Process entire batch at once - MUCH faster than per-frame!
                    preds = self.nlf_model.detect_smpl_batched(batch_tensor, model_name='smplx')
                except Exception as e:
                    print(f"[ERROR] NLF inference failed on batch {batch_start}-{batch_end}: {type(e).__name__}: {e}")
                    continue

            # Process results for each frame in the batch
            # NLF returns list with one entry per input frame
            for local_idx in range(batch_end - batch_start):
                frame_idx = batch_start + local_idx

                # Check if this frame has detections
                if 'boxes' not in preds or len(preds['boxes']) <= local_idx:
                    continue

                frame_boxes = preds['boxes'][local_idx]
                if len(frame_boxes) == 0:
                    continue

                num_detections = len(frame_boxes)

                # Accumulate results - batch the GPU->CPU transfer
                for key in ['boxes', 'joints3d', 'joints2d', 'joint_uncertainties', 'pose', 'betas', 'trans']:
                    if key in preds and len(preds[key]) > local_idx:
                        frame_data = preds[key][local_idx]
                        if num_detections > 0:
                            # Transfer entire frame's data at once (more efficient than per-detection)
                            frame_np = frame_data[:num_detections].detach().cpu().numpy()
                            for j in range(num_detections):
                                raw_outputs[key].append(frame_np[j])

                # Add frame indices for each detection
                for j in range(num_detections):
                    raw_outputs['img_idx'].append(base_frame_idx + frame_idx)

        # Handle case of no detections
        if len(raw_outputs['boxes']) == 0:
            return None

        # Convert lists to numpy arrays
        clean_outputs = {key: np.array(val) for key, val in raw_outputs.items()}

        # Reformat SMPLX parameters to match schema
        smplx_params = {
            'global_orient': clean_outputs['pose'][:, :3],
            'body_pose': clean_outputs['pose'][:, 3:66],
            'jaw_pose': clean_outputs['pose'][:, 66:69],
            'leye_pose': clean_outputs['pose'][:, 69:72],
            'reye_pose': clean_outputs['pose'][:, 72:75],
            'left_hand_pose': clean_outputs['pose'][:, 75:120],
            'right_hand_pose': clean_outputs['pose'][:, 120:165],
            'betas': clean_outputs['betas'],
            'trans': clean_outputs['trans'],
        }

        return {
            'boxes': clean_outputs['boxes'],
            'joints3d': clean_outputs['joints3d'],
            'joints2d': clean_outputs['joints2d'],
            'joint_uncertainties': clean_outputs['joint_uncertainties'],
            'img_idx': clean_outputs['img_idx'].astype(np.int32),
            'smplx': smplx_params,
        }
