"""
GCTBase - Base class for GCT model implementations.

Provides shared functionality:
- Prediction heads (camera, depth, point)
- Forward pass structure
- Model hub mixin (PyTorchModelHubMixin)
"""

import logging
import numpy as np
import torch
import torch.nn as nn
from abc import ABC, abstractmethod
from typing import Optional, Dict, Any, List, Union
from huggingface_hub import PyTorchModelHubMixin

from lingbot_map.heads.dpt_head import DPTHead
from lingbot_map.utils.pose_enc import pose_encoding_to_extri_intri
from lingbot_map.utils.geometry import closed_form_inverse_se3

logger = logging.getLogger(__name__)


class GCTBase(nn.Module, PyTorchModelHubMixin, ABC):
    """
    Base class for GCT model implementations.

    Handles shared components:
    - Prediction heads (camera, depth, point)
    - Forward pass structure
    - Input normalization

    Subclasses must implement:
    - _build_aggregator(): Create mode-specific aggregator
    - _build_camera_head(): Create mode-specific camera head
    """

    def __init__(
        self,
        # Architecture parameters
        img_size: int = 518,
        patch_size: int = 14,
        embed_dim: int = 1024,
        patch_embed: str = 'dinov2_vitl14_reg',
        disable_global_rope: bool = False,
        # Head configuration
        enable_camera: bool = True,
        enable_point: bool = True,
        enable_local_point: bool = False,
        enable_depth: bool = True,
        enable_track: bool = False,
        # Camera head sliding window
        enable_camera_sliding_window: bool = False,
        # 3D RoPE
        enable_3d_rope: bool = False,
        # Context Parallelism (kept for checkpoint compatibility but not used)
        enable_ulysses_cp: bool = False,
        # Normalization
        enable_normalize: bool = False,
        # Prediction normalization
        pred_normalization: bool = False,
        pred_normalization_detach_scale: bool = False,
        # Gradient checkpointing
        use_gradient_checkpoint: bool = True,
    ):
        super().__init__()

        # Store configuration
        self.img_size = img_size
        self.patch_size = patch_size
        self.embed_dim = embed_dim
        self.patch_embed = patch_embed
        self.disable_global_rope = disable_global_rope

        self.enable_ulysses_cp = False  # CP disabled in standalone package
        self.enable_normalize = enable_normalize
        self.pred_normalization = pred_normalization
        self.pred_normalization_detach_scale = pred_normalization_detach_scale
        self.use_gradient_checkpoint = use_gradient_checkpoint

        # Head flags
        self.enable_camera = enable_camera
        self.enable_point = enable_point
        self.enable_local_point = enable_local_point
        self.enable_depth = enable_depth
        self.enable_track = enable_track
        self.enable_camera_sliding_window = enable_camera_sliding_window
        self.enable_3d_rope = enable_3d_rope

        # Build aggregator (subclass-specific)
        self.aggregator = self._build_aggregator()

        # Build prediction heads (subclass-specific)
        self.camera_head = self._build_camera_head() if enable_camera else None
        self.point_head = self._build_point_head() if enable_point else None
        self.local_point_head = self._build_local_point_head() if enable_local_point else None
        self.depth_head = self._build_depth_head() if enable_depth else None

    @abstractmethod
    def _build_aggregator(self) -> nn.Module:
        pass

    @abstractmethod
    def _build_camera_head(self) -> nn.Module:
        pass

    def _build_depth_head(self) -> nn.Module:
        return DPTHead(
            dim_in=2 * self.embed_dim,
            patch_size=self.patch_size,
            output_dim=2,
            activation="exp",
            conf_activation="expp1"
        )

    def _build_point_head(self) -> nn.Module:
        return DPTHead(
            dim_in=2 * self.embed_dim,
            patch_size=self.patch_size,
            output_dim=4,
            activation="inv_log",
            conf_activation="expp1"
        )

    def _build_local_point_head(self) -> nn.Module:
        return DPTHead(
            dim_in=2 * self.embed_dim,
            patch_size=self.patch_size,
            output_dim=4,
            activation="inv_log",
            conf_activation="expp1"
        )

    def _normalize_input(self, images: torch.Tensor, query_points=None):
        if len(images.shape) == 4:
            images = images.unsqueeze(0)
        if query_points is not None and len(query_points.shape) == 2:
            query_points = query_points.unsqueeze(0)
        return images, query_points

    @abstractmethod
    def _aggregate_features(
        self,
        images: torch.Tensor,
        num_frame_for_scale: Optional[int] = None,
        sliding_window_size: Optional[int] = None,
        num_frame_per_block: int = 1,
        view_graphs: Optional[torch.Tensor] = None,
        causal_graphs: Optional[Union[torch.Tensor, List[np.ndarray]]] = None,
        ordered_video: Optional[torch.Tensor] = None,
        is_cp_sliced: bool = False,
    ) -> tuple:
        pass

    def _predict_camera(
        self,
        aggregated_tokens_list: list,
        mask: Optional[torch.Tensor] = None,
        causal_inference: bool = False,
        num_frame_for_scale: Optional[int] = None,
        sliding_window_size: Optional[int] = None,
        num_frame_per_block: int = 1,
        gather_outputs: bool = True,
    ) -> Dict[str, torch.Tensor]:
        if self.camera_head is None:
            return {}

        aggregated_tokens_list_fp32 = [t.float() for t in aggregated_tokens_list]

        camera_sliding_window = sliding_window_size if self.enable_camera_sliding_window else -1

        with torch.amp.autocast('cuda', enabled=False):
            pose_enc_list = self.camera_head(
                aggregated_tokens_list_fp32,
                mask=mask,
                causal_inference=causal_inference,
                num_frame_for_scale=num_frame_for_scale if num_frame_for_scale is not None else -1,
                sliding_window_size=camera_sliding_window,
                num_frame_per_block=num_frame_per_block,
            )

        return {
            "pose_enc": pose_enc_list[-1],
            "pose_enc_list": pose_enc_list,
        }

    def _predict_depth(
        self,
        aggregated_tokens_list: list,
        images: torch.Tensor,
        patch_start_idx: int,
        gather_outputs: bool = True,
    ) -> Dict[str, torch.Tensor]:
        if self.depth_head is None:
            return {}

        aggregated_tokens_list_fp32 = [t.float() for t in aggregated_tokens_list]
        images_fp32 = images.float()

        with torch.amp.autocast('cuda', enabled=False):
            depth, depth_conf = self.depth_head(
                aggregated_tokens_list_fp32,
                images=images_fp32,
                patch_start_idx=patch_start_idx
            )

        return {"depth": depth, "depth_conf": depth_conf}

    def _predict_points(
        self,
        aggregated_tokens_list: list,
        images: torch.Tensor,
        patch_start_idx: int,
        gather_outputs: bool = True,
    ) -> Dict[str, torch.Tensor]:
        if self.point_head is None:
            return {}

        aggregated_tokens_list_fp32 = [t.float() for t in aggregated_tokens_list]
        images_fp32 = images.float()

        with torch.amp.autocast('cuda', enabled=False):
            pts3d, pts3d_conf = self.point_head(
                aggregated_tokens_list_fp32,
                images=images_fp32,
                patch_start_idx=patch_start_idx
            )

        return {"world_points": pts3d, "world_points_conf": pts3d_conf}

    def _predict_local_points(
        self,
        aggregated_tokens_list: list,
        images: torch.Tensor,
        patch_start_idx: int,
        gather_outputs: bool = True,
    ) -> Dict[str, torch.Tensor]:
        if self.local_point_head is None:
            return {}

        aggregated_tokens_list_fp32 = [t.float() for t in aggregated_tokens_list]
        images_fp32 = images.float()

        with torch.amp.autocast('cuda', enabled=False):
            pts3d, pts3d_conf = self.local_point_head(
                aggregated_tokens_list_fp32,
                images=images_fp32,
                patch_start_idx=patch_start_idx
            )

        return {"cam_points": pts3d, "cam_points_conf": pts3d_conf}

    def _unproject_depth_to_world(
        self,
        depth: torch.Tensor,
        pose_enc: torch.Tensor,
    ) -> torch.Tensor:
        B, S, H, W, _ = depth.shape
        device = depth.device
        dtype = depth.dtype

        image_size_hw = (H, W)
        extrinsics, intrinsics = pose_encoding_to_extri_intri(
            pose_enc, image_size_hw=image_size_hw, build_intrinsics=True
        )

        extrinsics_flat = extrinsics.view(B * S, 3, 4)
        extrinsics_4x4 = torch.zeros(B * S, 4, 4, device=device, dtype=dtype)
        extrinsics_4x4[:, :3, :] = extrinsics_flat
        extrinsics_4x4[:, 3, 3] = 1.0
        c2w = closed_form_inverse_se3(extrinsics_4x4).view(B, S, 4, 4)

        y_grid, x_grid = torch.meshgrid(
            torch.arange(H, device=device, dtype=dtype),
            torch.arange(W, device=device, dtype=dtype),
            indexing='ij'
        )
        pixel_coords = torch.stack([x_grid, y_grid, torch.ones_like(x_grid)], dim=-1)

        intrinsics_inv = torch.inverse(intrinsics)
        camera_coords = torch.einsum('bsij,hwj->bshwi', intrinsics_inv, pixel_coords)
        camera_points = camera_coords * depth

        ones = torch.ones_like(camera_points[..., :1])
        camera_points_h = torch.cat([camera_points, ones], dim=-1)
        world_points_h = torch.einsum('bsij,bshwj->bshwi', c2w, camera_points_h)

        return world_points_h[..., :3]

    def forward(
        self,
        images: torch.Tensor,
        query_points: Optional[torch.Tensor] = None,
        num_frame_for_scale: Optional[int] = None,
        sliding_window_size: Optional[int] = None,
        num_frame_per_block: int = 1,
        mask: Optional[torch.Tensor] = None,
        causal_inference: bool = False,
        ordered_video: Optional[torch.Tensor] = None,
        gather_outputs: bool = True,
        point_masks: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> Dict[str, torch.Tensor]:
        """
        Forward pass of the GCT model.

        Args:
            images: Input images [S, 3, H, W] or [B, S, 3, H, W], in range [0, 1]
            query_points: Optional query points [N, 2] or [B, N, 2]

        Returns:
            Dictionary containing predictions:
                - pose_enc: Camera pose encoding [B, S, 9]
                - depth: Depth maps [B, S, H, W, 1]
                - depth_conf: Depth confidence [B, S, H, W]
                - world_points: 3D world coordinates [B, S, H, W, 3]
                - world_points_conf: Point confidence [B, S, H, W]
        """
        images, query_points = self._normalize_input(images, query_points)

        aggregated_tokens_list, patch_start_idx = self._aggregate_features(
            images,
            num_frame_for_scale=num_frame_for_scale,
            sliding_window_size=sliding_window_size,
            num_frame_per_block=num_frame_per_block,
        )

        predictions = {}

        predictions.update(self._predict_camera(
            aggregated_tokens_list,
            mask=ordered_video,
            causal_inference=causal_inference,
            num_frame_for_scale=num_frame_for_scale,
            sliding_window_size=sliding_window_size,
            num_frame_per_block=num_frame_per_block,
            gather_outputs=gather_outputs,
        ))

        predictions.update(self._predict_depth(
            aggregated_tokens_list, images, patch_start_idx,
            gather_outputs=gather_outputs,
        ))

        predictions.update(self._predict_points(
            aggregated_tokens_list, images, patch_start_idx,
            gather_outputs=gather_outputs,
        ))

        predictions.update(self._predict_local_points(
            aggregated_tokens_list, images, patch_start_idx,
            gather_outputs=gather_outputs,
        ))

        if not self.training:
            predictions["images"] = images

        return predictions