from abc import ABC
from copy import deepcopy
from typing import Tuple

import torch
from torch import Tensor as Ts
import torch.nn as nn

from anydata.geometry.camera_utils import multiply_extrinsics, invert_extrinsics
from anydata.utils.data import pixel_grid, norm_pixel_grid, unnorm_pixel_grid
from anydata.utils.types import is_tensor

try:
    from torch_scatter import scatter_min  #@IgnoreException
except (OSError, ImportError):
    def scatter_min(*args, **kwargs) -> Tuple[Ts, Ts]:
        raise ModuleNotFoundError(
            "torch_scatter is required to use scatter-based operations in CameraFTheta "
            "(e.g., project_pointcloud). Please install the 'torch_scatter' package."
        )


class CameraBase(nn.Module, ABC):
    """Base camera class. Stores _K [B,3,3], _Tcw [B,4,4] (cam-to-world), _hw (H,W).
    Subclasses implement lift() and unlift() for geometry-specific projection."""

    def __init__(self, hw=None, Twc=None, Tcw=None):
        super().__init__()
        assert Twc is None or Tcw is None
        self.geometry = 'base'
        self._K: Ts = None
        self._D: Ts = None
        # Extrinsics
        if Twc is None and Tcw is None:
            self._Tcw = torch.eye(
                4, dtype=self._K.dtype, device=self._K.device).unsqueeze(0).repeat(self._K.shape[0], 1, 1)
        else:
            self._Tcw = Tcw if Twc is None else invert_extrinsics(Twc)
        # Resolution
        self._hw: Tuple[int, int] = hw
        if is_tensor(self._hw):
            self._hw = self._hw.shape[-2:]

    ### GETTERS

    @property
    def hw(self):
        """Image resolution as (height, width)."""
        return self._hw

    @property
    def wh(self):
        """Image resolution as (width, height)."""
        return self._hw[::-1]

    @property
    def b(self):
        """Batch size."""
        return self._K.shape[0]

    @property
    def K(self):
        """Intrinsics [B,3,3]."""
        return self._K

    @property
    def D(self):
        """Distortion coefficients (None for pinhole)."""
        return self._D

    @property
    def Tcw(self):
        """Cam-to-world extrinsics [B,4,4] (stored directly)."""
        return self._Tcw

    @property
    def Twc(self):
        """World-to-cam extrinsics [B,4,4] (computed as inverse of Tcw)."""
        return None if self._Tcw is None else invert_extrinsics(self._Tcw)

    @property
    def device(self):
        return self._K.device

    @property
    def dtype(self):
        return self._K.dtype

    ### SETTERS

    @K.setter
    def K(self, K):
        self._K = K

    @D.setter
    def D(self, D):
        self._D = D

    @Tcw.setter
    def Tcw(self, Tcw):
        self._Tcw = Tcw

    ### DUNDER METHODS

    def __len__(self):
        return self._K.shape[0]

    def __getitem__(self, idx):
        if isinstance(idx, int):
            idx = slice(idx, idx + 1)
        return type(self)(K=self._K[idx], hw=self._hw, Tcw=self._Tcw[idx])

    def __str__(self):
        B = self.b
        H, W = self._hw if self._hw is not None else (0, 0)
        dev = self.device
        dt = self.dtype
        center = self.Tcw[:, :3, -1]
        kx, ky = self._K[:, 0, 2], self._K[:, 1, 2]
        per_b = [
            f'pos=[{", ".join(f"{v:.3f}" for v in center[i].tolist())}] '
            f'cx={kx[i].item():.1f} cy={ky[i].item():.1f}'
            for i in range(B)
        ]
        body = ' | '.join(per_b) if B > 1 else per_b[0]
        return (f'{type(self).__name__}(B={B}, {H}x{W}, {self.geometry}, '
                f'{body}, {dt}, {dev})')

    def __repr__(self):
        return self.__str__()

    ### UTILITY METHODS

    def clone(self):
        """Deep-copy camera (new tensors, same device)."""
        return type(self)(K=self._K.clone(), hw=self._hw,
                          Tcw=self._Tcw.clone() if self._Tcw is not None else None)

    def to(self, *args, **kwargs):
        """Move all tensor attributes to device/dtype. Mutates in-place, returns self."""
        for attr_name in list(vars(self)):
            val = getattr(self, attr_name)
            if isinstance(val, torch.Tensor):
                setattr(self, attr_name, val.to(*args, **kwargs))
        return self

    def float(self):
        """Cast to float32."""
        return self.to(dtype=torch.float32)

    def detach(self):
        """Return a clone with all tensors detached from the compute graph."""
        cam = self.clone()
        cam._K = cam._K.detach()
        if cam._Tcw is not None:
            cam._Tcw = cam._Tcw.detach()
        if cam._D is not None:
            cam._D = cam._D.detach()
        return cam

    @staticmethod
    def from_list(cams):
        """Batch cameras by concatenating along batch dim.
        All cameras must have the same type and hw."""
        K = torch.cat([c.K for c in cams], 0)
        Tcw = torch.cat([c.Tcw for c in cams], 0)
        return type(cams[0])(K=K, hw=cams[0].hw, Tcw=Tcw)

    def Pwc(self, from_world=True):
        """Projection matrix [B,4,4]: K (padded to 4x4) @ Twc.
        If from_world=False, returns just K padded to 4x4."""
        K4 = torch.eye(4, device=self.device, dtype=self.dtype).unsqueeze(0).repeat(self.b, 1, 1)
        K4[:, :3, :3] = self._K
        if not from_world or self.Twc is None:
            return K4
        return torch.bmm(K4, self.Twc)

    ### RECONSTRUCT DEPTH MAP

    def e2z(self, e_depth):
        """Convert euclidean depth to z-buffer depth.
        e_depth: [B,1,H,W]. Returns [B,1,H,W]."""
        points = self.reconstruct_euclidean(e_depth, to_world=False)
        return self.project_points(points, from_world=False, return_z=True)[1]

    def reconstruct_depth_map(self, depth: Ts, to_world=False, scene_flow=None, world_scene_flow=None, euclidean=False):
        """Reconstruct 3D points from a depth map.
        depth: [B,1,H,W]. Returns [B,3,H,W].
        Dispatches to reconstruct_z_buffer or reconstruct_euclidean."""
        if not euclidean:
            return self.reconstruct_z_buffer(
                depth=depth, to_world=to_world,
                scene_flow=scene_flow, world_scene_flow=world_scene_flow
            )
        else:
            return self.reconstruct_euclidean(
                depth=depth, to_world=to_world,
                scene_flow=scene_flow, world_scene_flow=world_scene_flow
            )

    def reconstruct_z_buffer(self, depth: Ts, to_world=False, scene_flow=None, world_scene_flow=None):
        """Reconstruct 3D points from z-buffer (planar) depth.
        depth: [B,1,H,W]. Returns [B,3,H,W]."""
        if depth is None: return None
        b, _, h, w = depth.shape
        grid = pixel_grid(depth, with_ones=True, device=self.device).view(b, 3, -1)
        points = self.lift(grid, depth.view(depth.shape[0], 1, -1))
        if scene_flow is not None:
            points = points + scene_flow.view(b, 3, -1)
        if to_world and self.Tcw is not None:
            points = multiply_extrinsics(points, self.Tcw)
            if world_scene_flow is not None:
                points = points + world_scene_flow.view(b, 3, -1)
        return points.view(b, 3, h, w)

    def reconstruct_euclidean(self, depth: Ts, to_world=False, scene_flow=None, world_scene_flow=None):
        """Reconstruct 3D points from euclidean (distance) depth.
        depth: [B,1,H,W]. Returns [B,3,H,W]."""
        if depth is None: return None
        b, _, h, w = depth.shape
        rays = self.get_viewdirs(normalize=True, to_world=False).view(b, 3, -1)
        points = rays * depth.view(depth.shape[0], 1, -1)
        if scene_flow is not None:
            points = points + scene_flow.view(b, 3, -1)
        if to_world and self.Tcw is not None:
            points = multiply_extrinsics(points, self.Tcw)
            if world_scene_flow is not None:
                points = points + world_scene_flow.view(b, 3, -1)
        return points.view(b, 3, h, w)

    def reconstruct_depth_map_rays(self, depth: Ts, to_world=False):
        """Reconstruct 3D points using normalized view rays * depth.
        depth: [B,1,H,W]. Returns [B,3,H,W]."""
        if depth is None:
            return None
        b, _, h, w = depth.shape
        rays = self.get_viewdirs(normalize=True, to_world=False)
        points = (rays * depth).view(b, 3, -1)
        if to_world and self.Tcw is not None:
            points = self.Tcw * points
        return points.view(b, 3, h, w)
        
    def project_points(self, points: Ts, from_world=True, normalize=True,
                       return_z=False, return_e=False, flag_invalid=True, max_depth=None):
        """Project 3D points to 2D pixel coordinates.
        points: [B,3,N] or [B,3,H,W].
        Returns coords [B,N,2] or [B,H,W,2] (normalized to [-1,1] if normalize=True).
        If return_z/return_e: also returns depth [B,N] or [B,1,H,W]."""
        is_depth_map = points.dim() == 4                            # Poincloud or point map
        hw = self._hw if not is_depth_map else points.shape[-2:]    # Get height and width for projection
        return_depth = return_z or return_e                         # Return zbuffer or euclidean depth

        if is_depth_map: # Flatten depth map
            points = points.reshape(points.shape[0], 3, -1)
        b, _, n = points.shape

        # Camera-specific unlift operation
        coords, depth = self.unlift(points, from_world=from_world, euclidean=return_e)

        if not is_depth_map:
            if normalize:
                coords = norm_pixel_grid(coords, hw=self._hw, in_place=True)
                if flag_invalid:
                    invalid = (coords[:, 0] < -1) | (coords[:, 0] > 1) | \
                              (coords[:, 1] < -1) | (coords[:, 1] > 1) | (depth <= 0)
                    if max_depth is not None:
                        invalid = invalid | (depth > max_depth)
                    coords[invalid.unsqueeze(1).repeat(1, 2, 1)] = -2
            else:
                if flag_invalid:
                    invalid = depth <= 0
                    if max_depth is not None:
                        invalid = invalid | (depth > max_depth)
                    coords[invalid.unsqueeze(1).repeat(1, 2, 1)] = -2
            if return_depth:
                return coords.permute(0, 2, 1), depth
            else:
                return coords.permute(0, 2, 1)

        coords = coords.view(b, 2, *hw)
        depth = depth.view(b, 1, *hw)

        if normalize:
            coords = norm_pixel_grid(coords, hw=self._hw, in_place=True)
            if flag_invalid:
                invalid = (coords[:, 0] < -1) | (coords[:, 0] > 1) | \
                          (coords[:, 1] < -1) | (coords[:, 1] > 1) | (depth[:, 0] <= 0)
                if max_depth is not None:
                    invalid = invalid | (depth > max_depth)
                coords[invalid.unsqueeze(1).repeat(1, 2, 1, 1)] = -2
        else:
            if flag_invalid:
                invalid = (depth[:, 0] <= 0)
                if max_depth is not None:
                    invalid = invalid | (depth > max_depth)
                coords[invalid.unsqueeze(1).repeat(1, 2, 1, 1)] = -2

        if return_depth:
            return coords.permute(0, 2, 3, 1), depth
        else:
            return coords.permute(0, 2, 3, 1)

    def to_world(self, points: Ts):
        """Transform points from camera to world coords.
        points: [B,3,N] or [B,3,H,W]. Returns same shape."""
        if points.dim() > 3:
            shape = points.shape
            points = points.reshape(points.shape[0], 3, -1)
        else:
            shape = None
        world_points = points if self.Tcw is None else multiply_extrinsics(points, self.Tcw)
        return world_points if shape is None else world_points.view(shape)

    def from_world(self, points: Ts):
        """Transform points from world to camera coords.
        points: [B,3,N] or [B,3,H,W]. Returns same shape."""
        if points.dim() > 3:
            shape = points.shape
            points = points.reshape(points.shape[0], 3, -1)
        else:
            shape = None
        cam_points = points if self.Twc is None else multiply_extrinsics(points, self.Twc)
        return cam_points if shape is None else cam_points.view(shape)

    def get_center(self):
        """Camera center in world coords [B,3]."""
        return self.Tcw[:, :3, -1]

    def get_origin(self, flatten=False):
        """Camera center broadcast to image grid.
        Returns [B,3,H,W] or [B,HW,3] if flatten=True."""
        orig = self.get_center().view(self.b, 3, 1, 1).repeat(1, 1, *self.hw)
        if flatten:
            orig = orig.reshape(self.b, 3, -1).permute(0, 2, 1)
        return orig

    def get_viewdirs(self, normalize=None, to_world=None, flatten=False):
        """View direction rays [B,3,H,W] or [B,HW,3] if flatten=True.
        normalize: True/'unit' for unit-length, 'plane' for z=1 plane."""

        # Get rays from a 1-unit depth map
        ones = torch.ones((self.b, 1, *self.hw), dtype=self.dtype, device=self.device)
        rays = self.reconstruct_depth_map(ones, to_world=False)

        # Normalize (unit is 1-unit radius, plane is 1-unit distnace plane)
        if normalize is True or normalize == 'unit':
            rays = rays / torch.norm(rays, dim=1).unsqueeze(1)
        elif normalize == 'plane':
            rays = rays / torch.norm(rays, dim=1).unsqueeze(1)
            rays = rays / rays[:, [2]]

        # Move viewing rays to world if required
        if to_world:
            # rays = self.to_world(rays).reshape(self.b, 3, *self.hw)
            rays = self.no_translation().to_world(rays).reshape(self.b, 3, *self.hw)

        # Flatten resulting ray map
        if flatten:
            rays = rays.reshape(self.b, 3, -1).permute(0, 2, 1)

        return rays
        
    def get_plucker(self, flatten=False, with_orig=False):
        """Plucker ray coordinates.
        Returns [B,6,H,W] (or [B,9,H,W] with_orig), [B,HW,6/9] if flatten.
        Channels: [ray_dir(3), cross(3), (origin(3))]."""

        b = len(self)
        ones = torch.ones((b, 1, *self.hw), dtype=self.dtype, device=self.device)  # (B, 1, H, W)
        rays = self.no_translation().reconstruct_depth_map(ones, to_world=True)
        
        # NOTE(bvh): without clamp this can give NaN if invalid intrinsics
        rays = rays / torch.norm(rays, dim=1).unsqueeze(1).clamp(min=1e-8)  # (B, 3, H, W)
        orig = self.get_center().view(b, 3, 1, 1).repeat(1, 1, *self.hw)  # (B, 3, H, W)

        # NOTE(bvh): critical bugfix applied here on 9/3/2025
        orig = orig.view(b, 3, -1).permute(0, 2, 1)  # (B, H*W, 3)
        rays = rays.view(b, 3, -1).permute(0, 2, 1)  # (B, H*W, 3)

        cross = torch.cross(orig, rays, dim=-1)  # (B, H*W, 3)
        if with_orig:
            plucker1 = torch.cat((rays, cross, orig), dim=-1)  # (B, H*W, 9)
        else:
            plucker1 = torch.cat((rays, cross), dim=-1)  # (B, H*W, 6)

        # NOTE(bvh): When you scale the translation component self.pose[:, 0:3, 3] by a scalar s,
        # then cross and orig are also scaled by s, but rays stay the same.

        if flatten:
            plucker2 = plucker1  # (B, H*W, 6/9)
        else:
            # plucker = plucker.permute(0, 2, 1).view(b, -1, *self.hw)  # can crash
            plucker2 = plucker1.permute(0, 2, 1).reshape(b, -1, *self.hw)  # (B, 6/9, H, W)

        return plucker2

    def project_pointcloud(self, pcl_src: Ts, from_world=True, euclidean=False, thr=1):
        """Creates a depth map from a pointcloud (requires torch-scatter)"""

        # Get projected coordinates and depth values
        res = self.project_points(pcl_src, from_world=from_world, return_z=not euclidean, return_e=euclidean)
        uv_all: Ts = res[0]
        z_all: Ts = res[1]

        # Loop over point-cloud batch-wise
        depths_tgt = []
        for i in range(pcl_src.shape[0]):

            # Flatten coordinates and depth values
            uv, z = uv_all[i].reshape(-1, 2), z_all[i].reshape(-1, 1)

            # Remove out-of-bounds coordinates and points behind the camera
            idx = (uv[:, 0] > -1) & (uv[:, 0] < 1) & \
                  (uv[:, 1] > -1) & (uv[:, 1] < 1) & (z[:, 0] > 0.0)

            if idx.sum() == 0:
                # Empty depth map
                depth_tgt = torch.zeros((self.hw[0] * self.hw[1], 1), device=z.device, dtype=z.dtype)
            else:
                # Unormalize and stack coordinates for scatter operation
                uv = unnorm_pixel_grid(uv[idx], self.hw).floor().long()
                uv = uv[:, 0] + uv[:, 1] * self.hw[1]

                # Min scatter operation (only keep the closest depth)
                depth_tgt = 1e10 * torch.ones((self.hw[0] * self.hw[1], 1), device=z.device, dtype=z.dtype)
                depth_tgt, argmin = scatter_min(src=z[idx], index=uv.unsqueeze(1), dim=0, out=depth_tgt)
                depth_tgt[depth_tgt == 1e10] = 0.
                depth_tgt[depth_tgt < 0] = 0.

                num_valid = (depth_tgt > 0).sum()
                if num_valid > thr:
                    # Substitute invalid values with zero
                    invalid = argmin == argmin.max()
                    argmin[invalid] = 0
                else:
                    # Too sparse, make empty depth map instead
                    depth_tgt = torch.zeros((self.hw[0] * self.hw[1], 1), device=z.device, dtype=z.dtype)

            # Reshape outputs
            depth_tgt = depth_tgt.reshape(1, 1, self.hw[0], self.hw[1])
            depths_tgt.append(depth_tgt)

        # Return stacked depth maps 
        return torch.cat(depths_tgt, 0)
    
    def no_translation(self):
        """Return new camera with zeroed translation (rotation only).
        Uses deepcopy to preserve subclass state (e.g. distortion params)."""
        cam = deepcopy(self)
        Tcw = cam.Tcw.clone()
        Tcw[:, :-1, -1] = 0
        cam.Tcw = Tcw
        return cam

    def lift(self, grid: Ts, depth: Ts) -> Ts:
        """Lift 2D pixel grid and depth map to 3D points"""
        raise NotImplementedError
    
    def unlift(self, points: Ts, from_world=True, euclidean=False) -> Tuple[Ts, Ts]:
        """Unlift 3D points to 2D"""
        raise NotImplementedError