import imp
import torch
import torch.nn as nn
from ..ball_query import ball_query
from torch.autograd import Function
import re
from glob import glob
import os
from perception3d.utils import build_ext
import numpy as np
import warnings

has_ext = True
if torch.cuda.is_available():
    try:
        from . import group_points_ext
    except (ImportError, ModuleNotFoundError):
        import cython
        build_ext('group_points_ext', os.path.dirname(__file__))
        try:
            from . import group_points_ext
        except (ImportError, ModuleNotFoundError):
            warnings.warn('Error building extension for {}'.format(os.path.basename(os.path.dirname(__file__))))
            has_ext = False


class GroupingOperation(Function):
    @staticmethod
    def forward(ctx, features, idx):
        # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor
        r"""
        Parameters
        ----------
        features : torch.Tensor
            (B, C, N) tensor of features to group
        idx : torch.Tensor
            (B, npoint, nsample) tensor containing the indicies of features to group with
        Returns
        -------
        torch.Tensor
            (B, C, npoint, nsample) tensor
        """
        ctx.save_for_backward(idx, features)

        return group_points_ext.group_points_forward(features, idx)

    @staticmethod
    def backward(ctx, grad_out):
        # type: (Any, torch.tensor) -> Tuple[torch.Tensor, torch.Tensor]
        r"""
        Parameters
        ----------
        grad_out : torch.Tensor
            (B, C, npoint, nsample) tensor of the gradients of the output from forward
        Returns
        -------
        torch.Tensor
            (B, C, N) gradient of the features
        None
        """
        idx, features = ctx.saved_tensors
        N = features.shape[2]

        grad_features = group_points_ext.group_points_backward(grad_out.contiguous(), idx, N)

        return grad_features, torch.zeros_like(idx)


def index_points_nocuda(points, idx):
    """
    Input:
        points: input points data, [B, C, N]
        idx: sample index data, [B, S]
    Return:
        new_points:, indexed points data, [B, C, S]
    """
    B = points.shape[0]
    view_shape = list(idx.shape)
    view_shape[1:] = [1] * (len(view_shape) - 1)
    repeat_shape = list(idx.shape)
    repeat_shape[0] = 1
    batch_indices = torch.arange(B, dtype=torch.long).to(points.device).view(view_shape).repeat(repeat_shape)
    new_points = points[batch_indices, :, idx]
    new_points = new_points.permute(0, -1, *np.arange(1, len(list(new_points.shape)) - 1))
    return new_points


grouping_operation = GroupingOperation.apply if has_ext else index_points_nocuda


class QueryAndGroup(nn.Module):
    r"""
    Groups with a ball query of radius
    Parameters
    ---------
    radius : float32
        Radius of ball
    nsample : int32
        Maximum number of features to gather in the ball
    """

    def __init__(self, radius, nsample, use_xyz=True, normalize_xyz=False):
        # type: (float, float, int, bool) -> None
        super().__init__()
        self.radius, self.nsample, self.use_xyz, self.normalize_xyz = radius, nsample, use_xyz, normalize_xyz

    def forward(self, xyz, center_xyz, features=None):
        # type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[Torch.Tensor]
        r"""
        Parameters
        ----------
        xyz : torch.Tensor
            xyz coordinates of the features (B, N, 3)
        center_xyz : torch.Tensor
            centriods (B, npoint, 3)
        features : torch.Tensor
            Descriptors of the features (B, C, N)
        Returns
        -------
        new_features : torch.Tensor
            (B, 3 + C, npoint, nsample) tensor
        """

        idx = ball_query(xyz, center_xyz, self.radius, self.nsample)
        xyz_trans = xyz.transpose(1, 2).contiguous()
        grouped_xyz = grouping_operation(xyz_trans, idx)  # (B, 3, npoint, nsample)
        grouped_xyz -= center_xyz.transpose(1, 2).unsqueeze(-1)
        if self.normalize_xyz:
            grouped_xyz /= self.radius

        if features is not None:
            grouped_features = grouping_operation(features, idx)
            if self.use_xyz:
                center_features = torch.cat(
                    [grouped_xyz, grouped_features], dim=1
                )  # (B, C + 3, npoint, nsample)
            else:
                center_features = grouped_features
        else:
            assert (
                self.use_xyz
            ), "Cannot have not features and not use xyz as a feature!"
            center_features = grouped_xyz

        return center_features


class GroupAll(nn.Module):
    r"""
    Groups all features
    Parameters
    ---------
    """

    def __init__(self, use_xyz=True):
        # type: (GroupAll, bool) -> None
        super(GroupAll, self).__init__()
        self.use_xyz = use_xyz

    def forward(self, xyz, center_xyz, features=None):
        # type: (GroupAll, torch.Tensor, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor]
        r"""
        Parameters
        ----------
        xyz : torch.Tensor
            xyz coordinates of the features (B, N, 3)
        new_xyz : torch.Tensor
            Ignored
        features : torch.Tensor
            Descriptors of the features (B, C, N)
        Returns
        -------
        new_features : torch.Tensor
            (B, C + 3, 1, N) tensor
        """

        grouped_xyz = xyz.transpose(1, 2).unsqueeze(2)
        if features is not None:
            grouped_features = features.unsqueeze(2)
            if self.use_xyz:
                center_features = torch.cat(
                    [grouped_xyz, grouped_features], dim=1
                )  # (B, 3 + C, 1, N)
            else:
                center_features = grouped_features
        else:
            center_features = grouped_xyz

        return center_features
    
__all__ = ['GroupAll', 'QueryAndGroup', 'grouping_operation']