import torch

from anydata.geometry.cameras.base import CameraBase
from anydata.utils.data import cat_channel_ones


class CameraFTheta(CameraBase):
    """F-theta fisheye camera model using polynomial angle mapping.
    K: [B,14] raw params (sx, sy, cx, cy, k_proj[5], k_lift[5])."""

    def __init__(self, K: torch.Tensor, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.geometry = 'ftheta'
        assert K.shape[1] == 14, 'FTheta camera requires 14 intrinsics parameters: ' + \
        '[scaling_factor_x, scaling_factor_y, cx, cy, fw0, fw1, fw2, fw3, fw4, bw0, bw1, bw2, bw3, bw4]'
        self._K = K

    def lift(self, grid: torch.Tensor, depth: torch.Tensor):
        """Lift 2D pixel coords + depth to 3D using f-theta polynomial.
        grid: [B,3,N], depth: [B,1,N]. Returns [B,3,N]."""

        # Getting camera intrinsics
        sx, sy, cx, cy = self.K[:, 0], self.K[:, 1], self.K[:, 2], self.K[:, 3]
        j01234 = self.K[:, None, 9:14]

        x = grid[:, 0] / sx - cx
        y = grid[:, 1] / sy - cy
        theta = torch.sqrt(x ** 2 + y ** 2)
        
        exponents = torch.arange(j01234.shape[-1], device=theta.device).float()[None,None]
        theta_pow = torch.pow(theta[:,:,None], exponents)
        theta_d = torch.sum(j01234 * theta_pow, dim=-1)

        ray_x = torch.sin(theta_d) * x / theta
        ray_y = torch.sin(theta_d) * y / theta
        ray_z = torch.cos(theta_d)
        rays = torch.stack([ray_x, ray_y, ray_z], axis=1)
        rays = rays / rays[:, [2]]

        return rays * depth

    def unlift(self, points: torch.Tensor, from_world=True, euclidean=False):
        """Project 3D points to 2D using inverse f-theta polynomial.
        points: [B,3,N]. Returns (coords [B,2,N], depth [B,N])."""

        sx, sy, cx, cy = self.K[:, 0], self.K[:, 1], self.K[:, 2], self.K[:, 3]
        k01234 = self.K[:, None, 4:9]
        
        points = torch.matmul(self.Twc, cat_channel_ones(points, 1))
        x, y, z = points[:, 0], points[:, 1], points[:, 2]

        r2 = torch.sqrt(x ** 2 + y ** 2)
        r3 = torch.sqrt(x ** 2 + y ** 2 + z ** 2)
        theta = torch.acos(z / r3)

        exponents = torch.arange(k01234.shape[-1], device=theta.device).float()[None,None]
        theta_pow = torch.pow(theta[:,:,None], exponents)
        theta_d = torch.sum(k01234 * theta_pow, dim=-1)
        
        u = cx + (x / r2) * theta_d
        v = cy + (y / r2) * theta_d

        coords, depth = torch.stack([u * sx, v * sy], 1), z
        return coords, depth

    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(len(self), 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(len(self), 3, -1).permute(0, 2, 1)

        return rays