import torch


def camera_rays(camK, W=None, H=None, c2w=None, graphics_coordinate=True,
                center=False):
    """shoot viewing rays from camera parameters.
    Args:
        camK: Tensor of shape `[3,3]`, the intrinsic matrix.
        W: Integer, if set None, then `W` is calculated as `2*cx`.
        H: Integer, if set None, then `H` is calculated as `2*cy`.
        c2w: Tensor of shape `[4,4]` or `[3,4]` camera view matrix. 
        If `None`, c2w is set as `[I,0]`
        graphics_coordinate: bool. Where or not use graphics coordinate 
        (pointing negative z into screen). Default: `True`.
        center: bool. Where or set 0.5 offset for pixels Default: `False`.
    Returns:
        rays_o: tensor of shape `[W,H,3]`  origins of the rays.
        rays_d: tensor of shape `[W,H,3]`  directions of the rays.
    """
    if c2w is None:
        c2w = torch.hstack((torch.eye(3), torch.zeros((3, 1))))
    if W is None:
        W = camK[0, 2]*2
    if H is None:
        H = camK[1, 2]*2
    W = int(W)
    H = int(H)

    invK = torch.inverse(camK)
    u, v = torch.meshgrid(torch.linspace(0, W-1, W), torch.linspace(0, H-1, H))
    u = u.t()
    v = v.t()

    if center:
        u = u + 0.5
        v = v + 0.5
    dirs = torch.stack([u, v, torch.ones_like(u)], dim=-1)
    dirs = torch.matmul(dirs, invK.T)
    # use graphics coordinate. negtive z pointing into screen.
    if graphics_coordinate:
        dirs[..., 1] *= -1
        dirs[..., 2] *= -1

    rays_d = torch.matmul(dirs, c2w[:3, :3].T)
    rays_o = c2w[:3, -1].expand(rays_d.shape)
    return torch.cat([rays_o, rays_d], dim=-1)