"""
ComfyUI-CameraPack: Convert camera parameters between formats.

Converts DepthAnythingV3 extrinsics/intrinsics to ComfyUI's Load3DCamera format.
"""

import ast
import re
import torch
import numpy as np


class DA3ToLoad3DCamera:
    """
    Convert DepthAnythingV3 camera parameters to ComfyUI Load3DCamera format.

    DA3 outputs:
    - extrinsics: [B, 4, 4] pose matrix (rotation + translation)
    - intrinsics: [B, 3, 3] matrix [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]

    Load3DCamera expects:
    - position: {x, y, z} camera position in world space
    - target: {x, y, z} look-at point
    - zoom: int (derived from focal length)
    - cameraType: "perspective" or "orthographic"
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "extrinsics": ("EXTRINSICS",),
                "intrinsics": ("INTRINSICS",),
            },
            "optional": {
                "target_distance": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.1,
                    "max": 100.0,
                    "step": 0.1,
                    "tooltip": "Distance from camera to target point along view direction"
                }),
                "zoom_scale": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.1,
                    "max": 10.0,
                    "step": 0.1,
                    "tooltip": "Scale factor for zoom (multiply focal length)"
                }),
            }
        }

    RETURN_TYPES = ("LOAD3D_CAMERA",)
    RETURN_NAMES = ("camera_info",)
    FUNCTION = "convert"
    CATEGORY = "3d/camera"

    def convert(self, extrinsics, intrinsics, target_distance=1.0, zoom_scale=1.0):
        """
        Convert DA3 camera parameters to Load3DCamera format.

        The extrinsics matrix is a 4x4 world-to-camera transform:
        [R | t]
        [0 | 1]

        Camera position in world space: pos = -R^T @ t
        Camera forward direction: forward = R^T @ [0, 0, 1]
        """
        # Handle batch dimension - take first item
        if isinstance(extrinsics, torch.Tensor):
            ext = extrinsics[0].cpu().numpy() if extrinsics.dim() > 2 else extrinsics.cpu().numpy()
        else:
            ext = np.array(extrinsics)

        if isinstance(intrinsics, torch.Tensor):
            intr = intrinsics[0].cpu().numpy() if intrinsics.dim() > 2 else intrinsics.cpu().numpy()
        else:
            intr = np.array(intrinsics)

        # Flatten and handle different extrinsics formats
        ext = ext.flatten()
        if ext.size == 16:
            ext = ext.reshape(4, 4)
        elif ext.size == 12:
            # 3x4 matrix - add [0, 0, 0, 1] row to make 4x4
            ext = ext.reshape(3, 4)
            ext = np.vstack([ext, [0, 0, 0, 1]])
        else:
            raise ValueError(f"Unexpected extrinsics size: {ext.size}, expected 12 or 16")

        # Ensure intrinsics is 3x3
        intr = intr.flatten()
        if intr.size == 9:
            intr = intr.reshape(3, 3)
        else:
            raise ValueError(f"Unexpected intrinsics size: {intr.size}, expected 9")

        # Extract rotation and translation from extrinsics
        R = ext[:3, :3]  # 3x3 rotation matrix
        t = ext[:3, 3]   # 3x1 translation vector

        # Camera position in world space: pos = -R^T @ t
        R_T = R.T
        camera_pos = -R_T @ t

        # Camera forward direction (Z-axis in camera space transformed to world)
        # In OpenGL/standard convention, camera looks down -Z, so forward is [0, 0, -1]
        # In some conventions it's [0, 0, 1]. We'll use [0, 0, 1] and let user adjust.
        forward = R_T @ np.array([0, 0, 1])

        # Target point is camera position + forward * distance
        target_pos = camera_pos + forward * target_distance

        # Extract focal length from intrinsics
        fx = float(intr[0, 0])
        fy = float(intr[1, 1])
        focal_length = (fx + fy) / 2.0

        # Compute zoom - this is somewhat arbitrary, scale to reasonable range
        zoom = int(focal_length * zoom_scale)
        zoom = max(1, min(zoom, 1000))  # Clamp to reasonable range

        # Build camera_info dict
        camera_info = {
            "position": {
                "x": float(camera_pos[0]),
                "y": float(camera_pos[1]),
                "z": float(camera_pos[2]),
            },
            "target": {
                "x": float(target_pos[0]),
                "y": float(target_pos[1]),
                "z": float(target_pos[2]),
            },
            "zoom": zoom,
            "cameraType": "perspective",
        }

        return (camera_info,)


class CreateLoad3DCamera:
    """
    Manually create a Load3DCamera from position, target, and zoom values.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "pos_x": ("FLOAT", {"default": 0.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
                "pos_y": ("FLOAT", {"default": 0.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
                "pos_z": ("FLOAT", {"default": 5.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
                "target_x": ("FLOAT", {"default": 0.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
                "target_y": ("FLOAT", {"default": 0.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
                "target_z": ("FLOAT", {"default": 0.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
                "zoom": ("INT", {"default": 50, "min": 1, "max": 1000, "step": 1}),
                "camera_type": (["perspective", "orthographic"], {"default": "perspective"}),
            }
        }

    RETURN_TYPES = ("LOAD3D_CAMERA",)
    RETURN_NAMES = ("camera_info",)
    FUNCTION = "create"
    CATEGORY = "3d/camera"

    def create(self, pos_x, pos_y, pos_z, target_x, target_y, target_z, zoom, camera_type):
        camera_info = {
            "position": {"x": pos_x, "y": pos_y, "z": pos_z},
            "target": {"x": target_x, "y": target_y, "z": target_z},
            "zoom": zoom,
            "cameraType": camera_type,
        }
        return (camera_info,)


class CameraIntrinsics:
    """
    Create an INTRINSICS tensor from a pasted 3x3 matrix.

    Accepts formats like:
      [[3.2782e+03, 0, 960], [0, 2.4555e+03, 540], [0, 0, 1]]
      tensor([[[3278.2, 0, 960], [0, 2455.5, 540], [0, 0, 1]]])
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "matrix": ("STRING", {
                    "default": "[[1000, 0, 960],\n [0, 1000, 540],\n [0, 0, 1]]",
                    "multiline": True,
                    "tooltip": "3x3 intrinsics matrix [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]"
                }),
            }
        }

    RETURN_TYPES = ("INTRINSICS",)
    RETURN_NAMES = ("intrinsics",)
    FUNCTION = "create"
    CATEGORY = "3d/camera"

    def create(self, matrix):
        text = matrix.strip()
        # Strip tensor(...) wrapper if pasted from PyTorch output
        text = re.sub(r'^tensor\(', '', text)
        text = re.sub(r'\)\s*$', '', text)
        text = text.strip()

        parsed = ast.literal_eval(text)
        arr = np.array(parsed, dtype=np.float32)

        # Handle [1, 3, 3] input
        if arr.ndim == 3 and arr.shape == (1, 3, 3):
            arr = arr[0]

        if arr.shape != (3, 3):
            raise ValueError(f"Expected 3x3 matrix, got shape {arr.shape}")

        intrinsics = torch.from_numpy(arr).unsqueeze(0)  # [1, 3, 3]
        return (intrinsics,)


class GetExtrinsicsFromBatch:
    """Slice a sub-batch (or single) of extrinsics from a batched [N, 4, 4]
    EXTRINSICS tensor — mirrors ComfyUI's native `ImageFromBatch`.

    Use to pick one camera out of a 12-face panorama split, the first view
    of an HYWM2 reconstruction, etc. Negative `batch_index` wraps from the
    end; `length=0` returns the full tail starting at `batch_index`.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "extrinsics": ("EXTRINSICS", {
                    "tooltip": "Batched w2c [N, 4, 4] (or [1, N, 4, 4]).",
                }),
                "batch_index": ("INT", {
                    "default": 0, "min": -4096, "max": 4096, "step": 1,
                    "tooltip": "Starting index. Negative wraps from the end "
                               "(e.g. -1 = last extrinsic).",
                }),
                "length": ("INT", {
                    "default": 1, "min": 0, "max": 4096, "step": 1,
                    "tooltip": "Number of consecutive extrinsics to return "
                               "starting at batch_index. 0 = take everything "
                               "from batch_index to the end. Default 1 = "
                               "single extrinsic.",
                }),
            },
        }

    RETURN_TYPES = ("EXTRINSICS",)
    RETURN_NAMES = ("extrinsics",)
    FUNCTION = "slice_batch"
    CATEGORY = "3d/camera"

    def slice_batch(self, extrinsics, batch_index: int = 0, length: int = 1):
        if not isinstance(extrinsics, torch.Tensor):
            extrinsics = torch.as_tensor(extrinsics).float()
        ext = extrinsics
        if ext.dim() == 4 and ext.shape[0] == 1:
            ext = ext[0]
        if ext.dim() != 3 or ext.shape[-2:] not in ((4, 4), (3, 4)):
            raise ValueError(
                f"GetExtrinsicsFromBatch: expected [N, 4, 4] or [N, 3, 4]; "
                f"got {tuple(extrinsics.shape)}"
            )
        N = int(ext.shape[0])
        # Negative wrap.
        start = batch_index % N if N > 0 else 0
        if length <= 0:
            end = N
        else:
            end = min(start + length, N)
        if start >= end:
            raise ValueError(
                f"GetExtrinsicsFromBatch: empty slice (start={start} >= "
                f"end={end}, N={N}, batch_index={batch_index}, length={length})"
            )
        # `.clone()` (not just `.contiguous()`) — comfy-env's shm IPC
        # serializer reads the underlying storage size, not the view's
        # logical size. A sliced view of a 12×4×4 tensor still points at
        # the 192-element backing buffer; the receiver tries to reshape
        # 192 floats into [1, 4, 4] and explodes. `.clone()` allocates
        # fresh, view-free storage matching the slice's shape exactly.
        sliced = ext[start:end].clone().contiguous()
        # Squeeze single-element batches so single-camera consumers
        # (SHARP's predict, etc.) get the [4, 4] they expect rather than
        # [1, 4, 4]. Multi-element slices stay batched.
        if sliced.shape[0] == 1:
            sliced = sliced[0].contiguous()
        return (sliced,)


class GetIntrinsicsFromBatch:
    """Slice a sub-batch (or single) of intrinsics from a batched [N, 3, 3]
    INTRINSICS tensor — mirrors ComfyUI's native `ImageFromBatch`.

    Same semantics as `GetExtrinsicsFromBatch` but for K matrices.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "intrinsics": ("INTRINSICS", {
                    "tooltip": "Batched K [N, 3, 3] (or [1, N, 3, 3]).",
                }),
                "batch_index": ("INT", {
                    "default": 0, "min": -4096, "max": 4096, "step": 1,
                    "tooltip": "Starting index. Negative wraps from the end.",
                }),
                "length": ("INT", {
                    "default": 1, "min": 0, "max": 4096, "step": 1,
                    "tooltip": "Number of consecutive intrinsics to return. "
                               "0 = take everything from batch_index to the "
                               "end. Default 1 = single K matrix.",
                }),
            },
        }

    RETURN_TYPES = ("INTRINSICS",)
    RETURN_NAMES = ("intrinsics",)
    FUNCTION = "slice_batch"
    CATEGORY = "3d/camera"

    def slice_batch(self, intrinsics, batch_index: int = 0, length: int = 1):
        if not isinstance(intrinsics, torch.Tensor):
            intrinsics = torch.as_tensor(intrinsics).float()
        intr = intrinsics
        if intr.dim() == 4 and intr.shape[0] == 1:
            intr = intr[0]
        if intr.dim() == 2:
            # Single K passed unbatched — wrap then slice.
            intr = intr.unsqueeze(0)
        if intr.dim() != 3 or intr.shape[-2:] != (3, 3):
            raise ValueError(
                f"GetIntrinsicsFromBatch: expected [N, 3, 3]; "
                f"got {tuple(intrinsics.shape)}"
            )
        N = int(intr.shape[0])
        start = batch_index % N if N > 0 else 0
        if length <= 0:
            end = N
        else:
            end = min(start + length, N)
        if start >= end:
            raise ValueError(
                f"GetIntrinsicsFromBatch: empty slice (start={start} >= "
                f"end={end}, N={N}, batch_index={batch_index}, length={length})"
            )
        # See GetExtrinsicsFromBatch.slice_batch for why `.clone()` (not
        # just `.contiguous()`) — comfy-env's shm IPC keeps the original
        # storage when the tensor is a view.
        sliced = intr[start:end].clone().contiguous()
        # Squeeze single-element batches so single-camera consumers
        # (SHARP's predict — `intrinsics[0, 0].item()` for fx, etc.)
        # get the [3, 3] they expect rather than [1, 3, 3].
        if sliced.shape[0] == 1:
            sliced = sliced[0].contiguous()
        return (sliced,)


# Node registration
NODE_CLASS_MAPPINGS = {
    "DA3ToLoad3DCamera": DA3ToLoad3DCamera,
    "CreateLoad3DCamera": CreateLoad3DCamera,
    "CameraIntrinsics": CameraIntrinsics,
    "GetExtrinsicsFromBatch": GetExtrinsicsFromBatch,
    "GetIntrinsicsFromBatch": GetIntrinsicsFromBatch,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DA3ToLoad3DCamera": "DA3 to Load3D Camera",
    "CreateLoad3DCamera": "Create Load3D Camera",
    "CameraIntrinsics": "Camera Intrinsics",
    "GetExtrinsicsFromBatch": "Get Extrinsics From Batch",
    "GetIntrinsicsFromBatch": "Get Intrinsics From Batch",
}