import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from pathlib import Path
from ...ops.pointnet2.pointnet2_batch import pointnet2_modules
from ...ops.pointnet2.pointnet2_stack import pointnet2_modules as pointnet2_modules_stack
from ...ops.pointnet2.pointnet2_stack import pointnet2_utils as pointnet2_utils_stack
from ...ops.roiaware_pool3d import roiaware_pool3d_utils
from ...utils import loss_utils

import pdb

class PointNet2MSG(nn.Module):
    def __init__(self, model_cfg, input_channels, **kwargs):
        super().__init__()
        self.model_cfg = model_cfg

        self.SA_modules = nn.ModuleList()
        channel_in = input_channels - 3

        self.num_points_each_layer = []
        skip_channel_list = [input_channels - 3]
        for k in range(self.model_cfg.SA_CONFIG.NPOINTS.__len__()):
            mlps = self.model_cfg.SA_CONFIG.MLPS[k].copy()
            channel_out = 0
            for idx in range(mlps.__len__()):
                mlps[idx] = [channel_in] + mlps[idx]
                channel_out += mlps[idx][-1]

            self.SA_modules.append(
                pointnet2_modules.PointnetSAModuleMSG(
                    npoint=self.model_cfg.SA_CONFIG.NPOINTS[k],
                    radii=self.model_cfg.SA_CONFIG.RADIUS[k],
                    nsamples=self.model_cfg.SA_CONFIG.NSAMPLE[k],
                    mlps=mlps,
                    use_xyz=self.model_cfg.SA_CONFIG.get('USE_XYZ', True),
                )
            )
            skip_channel_list.append(channel_out)
            channel_in = channel_out

        self.FP_modules = nn.ModuleList()

        for k in range(self.model_cfg.FP_MLPS.__len__()):
            pre_channel = self.model_cfg.FP_MLPS[k + 1][-1] if k + 1 < len(self.model_cfg.FP_MLPS) else channel_out
            self.FP_modules.append(
                pointnet2_modules.PointnetFPModule(
                    mlp=[pre_channel + skip_channel_list[k]] + self.model_cfg.FP_MLPS[k]
                )
            )

        self.num_point_features = self.model_cfg.FP_MLPS[0][-1]

    def break_up_pc(self, pc):
        batch_idx = pc[:, 0]
        xyz = pc[:, 1:4].contiguous()
        features = (pc[:, 4:].contiguous() if pc.size(-1) > 4 else None)
        return batch_idx, xyz, features

    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size: int
                vfe_features: (num_voxels, C)
                points: (num_points, 4 + C), [batch_idx, x, y, z, ...]
        Returns:
            batch_dict:
                encoded_spconv_tensor: sparse tensor
                point_features: (N, C)
        """
        batch_size = batch_dict['batch_size']
        points = batch_dict['points']
        batch_idx, xyz, features = self.break_up_pc(points)

        xyz_batch_cnt = xyz.new_zeros(batch_size).int()
        for bs_idx in range(batch_size):
            xyz_batch_cnt[bs_idx] = (batch_idx == bs_idx).sum()

        assert xyz_batch_cnt.min() == xyz_batch_cnt.max()
        xyz = xyz.view(batch_size, -1, 3)
        features = features.view(batch_size, -1, features.shape[-1]).permute(0, 2, 1) if features is not None else None

        l_xyz, l_features = [xyz], [features]
        for i in range(len(self.SA_modules)):
            li_xyz, li_features = self.SA_modules[i](l_xyz[i], l_features[i])
            l_xyz.append(li_xyz)
            l_features.append(li_features)

        for i in range(-1, -(len(self.FP_modules) + 1), -1):
            l_features[i - 1] = self.FP_modules[i](
                l_xyz[i - 1], l_xyz[i], l_features[i - 1], l_features[i]
            )  # (B, C, N)

        point_features = l_features[0].permute(0, 2, 1).contiguous()  # (B, N, C)
        batch_dict['point_features'] = point_features.view(-1, point_features.shape[-1])
        batch_dict['point_coords'] = torch.cat((batch_idx[:, None].float(), l_xyz[0].view(-1, 3)), dim=1)
        return batch_dict


class PointNet2MSGInst(nn.Module):
    def __init__(self, model_cfg, input_channels, **kwargs):
        super().__init__()
        self.model_cfg = model_cfg

        self.SA_modules = nn.ModuleList()
        channel_in = input_channels - 3

        self.num_points_each_layer = []
        skip_channel_list = [input_channels - 3]
        for k in range(self.model_cfg.SA_CONFIG.NPOINTS.__len__()):
            mlps = self.model_cfg.SA_CONFIG.MLPS[k].copy()
            channel_out = 0
            for idx in range(mlps.__len__()):
                mlps[idx] = [channel_in] + mlps[idx]
                channel_out += mlps[idx][-1]

            self.SA_modules.append(
                pointnet2_modules.PointnetSAModuleMSGInst(
                    npoint=self.model_cfg.SA_CONFIG.NPOINTS[k],
                    radii=self.model_cfg.SA_CONFIG.RADIUS[k],
                    nsamples=self.model_cfg.SA_CONFIG.NSAMPLE[k],
                    mlps=mlps,
                    use_xyz=self.model_cfg.SA_CONFIG.get('USE_XYZ', True),
                )
            )
            skip_channel_list.append(channel_out)
            channel_in = channel_out

        self.FP_modules = nn.ModuleList()

        for k in range(self.model_cfg.FP_MLPS.__len__()):
            pre_channel = self.model_cfg.FP_MLPS[k + 1][-1] if k + 1 < len(self.model_cfg.FP_MLPS) else channel_out
            self.FP_modules.append(
                pointnet2_modules.PointnetFPModuleInst(
                    mlp=[pre_channel + skip_channel_list[k]] + self.model_cfg.FP_MLPS[k],
                    #inst_cls=[pre_channel + 3] + self.model_cfg.INST_FC  # 3 for relative xyz,
                    #inst_cls = [pre_channel + skip_channel_list[k] + 3] + self.model_cfg.INST_FC,  # 3 for relative xyz
                    inst_cls=[self.model_cfg.FP_MLPS[k][-1] + pre_channel + 3] + self.model_cfg.INST_FC,  # 3 for relative xyz
                    weight_mode=self.model_cfg.WEIGHT_MODE, inst_thres=self.model_cfg.INST_THRES, only_inst=self.model_cfg.get('ONLY_INST', True)
                )
            )

        self.num_point_features = self.model_cfg.FP_MLPS[0][-1]

    def break_up_pc(self, pc):
        batch_idx = pc[:, 0]
        xyz = pc[:, 1:4].contiguous()
        features = (pc[:, 4:].contiguous() if pc.size(-1) > 4 else None)
        return batch_idx, xyz, features

    def assign_instance_targets(self, points, gt_boxes):
        """
        Args:
            points: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]
            gt_boxes: (B, M, 8)

        Returns:
            point_inst_labels: (N1 + N2 + N3 + ...), long type, -1:background
        """
        if gt_boxes is None:
            return None
        assert len(points.shape) == 2 and points.shape[1] == 4, 'points.shape=%s' % str(points.shape)
        assert len(gt_boxes.shape) == 3 and gt_boxes.shape[2] == 8, 'gt_boxes.shape=%s' % str(gt_boxes.shape)

        batch_size = gt_boxes.shape[0]
        bs_idx = points[:, 0]
        point_inst_labels = points.new_zeros(points.shape[0]).long()
        for k in range(batch_size):
            bs_mask = (bs_idx == k)
            points_single = points[bs_mask][:, 1:4]
            box_idxs_of_pts = roiaware_pool3d_utils.points_in_boxes_gpu(
                points_single.unsqueeze(dim=0), gt_boxes[k:k + 1, :, 0:7].contiguous()
            ).long().squeeze(dim=0)

            point_inst_labels[bs_mask] = box_idxs_of_pts

        return point_inst_labels.view(batch_size, -1)

    def get_loss(self, tb_dict=None):
        tb_dict = {} if tb_dict is None else tb_dict
        # predictions & labels for each FP stage
        same_inst_loss = 0.0
        total_num_neigh = 0
        total_num_label_pos = 0
        total_num_label_neg = 0
        total_num_pred_pos = 0
        total_num_pred_neg = 0
        total_num_true_pos = 0
        total_num_true_neg = 0
        total_num_false_pos = 0
        total_num_false_neg = 0
        total_conf_true_pos = 0
        total_conf_true_neg = 0
        total_conf_false_pos = 0
        total_conf_false_neg = 0
        for layer_idx, l_dict in enumerate(self.forward_ret_dict):
            same_inst_preds_logit = l_dict['batch_neigh_preds']
            same_inst_preds = torch.sigmoid(same_inst_preds_logit)
            same_inst_labels = l_dict['batch_neigh_labels']
            if self.model_cfg.LOSS_INST == 'BinaryCrossEntropy':
                same_inst_loss += F.binary_cross_entropy(same_inst_preds, same_inst_labels)
            elif self.model_cfg.LOSS_INST == 'Focal':
                inst_focal_loss = loss_utils.SigmoidFocalClassificationLoss(gamma=self.model_cfg.GAMMA, alpha=self.model_cfg.ALPHA)
                neg_normalizer = torch.clamp((same_inst_labels == 0).sum(), min=1.0).float()
                #loss_weights = same_inst_preds.new_ones(tuple(same_inst_labels.shape[:-1])) / neg_normalizer
                loss_weights = same_inst_preds.new_ones(tuple(same_inst_labels.shape)) / neg_normalizer
                loss_weights[same_inst_labels == -1] = 0  # ignore loss from same background neighbor
                same_inst_loss += inst_focal_loss(same_inst_preds_logit, same_inst_labels, loss_weights).sum()
                tb_dict.update({f'layer{layer_idx}_focal_loss': inst_focal_loss(same_inst_preds_logit, same_inst_labels, loss_weights).sum()})
            elif self.model_cfg.LOSS_INST == 'LayerFocal':
                inst_focal_loss = loss_utils.SigmoidFocalClassificationLoss(gamma=(self.model_cfg.GAMMA[layer_idx]),
                                                                            alpha=self.model_cfg.ALPHA[layer_idx])
                neg_normalizer = torch.clamp((same_inst_labels == 0).sum(), min=1.0).float()
                #loss_weights = same_inst_preds.new_ones(tuple(same_inst_labels.shape[:-1])) / neg_normalizer
                loss_weights = same_inst_preds.new_ones(tuple(same_inst_labels.shape)) / neg_normalizer
                loss_weights[same_inst_labels == -1] = 0
                same_inst_loss += inst_focal_loss(same_inst_preds_logit, same_inst_labels, loss_weights).sum()
            else:
                raise NotImplementedError
            
            # >> 8/31 instance prediction metric
            inst_label_pos = (same_inst_labels == 1)  # 1: positive label, 0: else
            inst_label_neg = (same_inst_labels == 0)  # 1: negative label, 0: else
            inst_pred_pos = (same_inst_preds >= self.model_cfg.INST_THRES)  # 1: positive prediction, 0: negative prediction
            inst_pred_neg = (same_inst_preds <= (1 - self.model_cfg.INST_THRES))  # 1: negative prediction, 0: positive prediction

            inst_pred_true_pos = ((inst_pred_pos == inst_label_pos) * inst_label_pos).float()  # 1: True positive, 0: else
            inst_pred_true_neg = ((inst_pred_neg == inst_label_neg) * inst_label_neg).float()  # 1: True negative, 0: else
            inst_pred_false_pos = ((inst_pred_pos == inst_label_neg) * inst_label_neg).float()  # 1: False positive, 0: else
            inst_pred_false_neg = ((inst_pred_neg == inst_label_pos) * inst_label_pos).float()  # 1: False negative, 0: else

            inst_conf_true_pos = same_inst_preds * inst_pred_true_pos  # the higher, the better
            inst_conf_true_neg = (inst_pred_true_neg - same_inst_preds) * inst_pred_true_neg  # the higher, the better
            inst_conf_false_pos = same_inst_preds * inst_pred_false_pos  # the lower, the better
            inst_conf_false_neg = (inst_pred_false_neg - same_inst_preds) * inst_pred_false_neg  # the lower, the better

            tb_dict.update({f'layer{layer_idx}_pos_acc': inst_pred_true_pos.sum() / inst_pred_pos.sum(),
                            f'layer{layer_idx}_neg_acc': inst_pred_true_neg.sum() / inst_pred_neg.sum(),
                            f'layer{layer_idx}_pos_recall': inst_pred_true_pos.sum() / inst_label_pos.sum(),
                            f'layer{layer_idx}_neg_recall': inst_pred_true_neg.sum() / inst_label_neg.sum(),
                            f'layer{layer_idx}_neg_ratio': inst_label_neg.sum().float() / same_inst_labels.numel()})

            total_num_neigh += same_inst_labels.numel()
            total_num_label_pos += inst_label_pos.sum()
            total_num_label_neg += inst_label_neg.sum()
            total_num_pred_pos += inst_pred_pos.sum()
            total_num_pred_neg += inst_pred_neg.sum()
            total_num_true_pos += inst_pred_true_pos.sum()
            total_num_true_neg += inst_pred_true_neg.sum()
            total_num_false_pos += inst_pred_false_pos.sum()
            total_num_false_neg += inst_pred_false_neg.sum()
            total_conf_true_pos += inst_conf_true_pos.sum()
            total_conf_true_neg += inst_conf_true_neg.sum()
            total_conf_false_pos += inst_conf_false_pos.sum()
            total_conf_false_neg += inst_conf_false_neg.sum()

        same_inst_loss /= len(self.forward_ret_dict)
        #inst_pred_acc = (total_num_true_pos + total_num_true_neg).float() / total_num_neigh
        inst_pred_acc = (total_num_true_pos + total_num_true_neg).float() / (total_num_label_pos + total_num_label_neg)
        inst_pred_acc_pos = total_num_true_pos.float() / total_num_pred_pos
        inst_pred_acc_neg = total_num_true_neg.float() / total_num_pred_neg
        inst_pred_recall_pos = total_num_true_pos.float() / total_num_label_pos
        inst_pred_recall_neg = total_num_true_neg.float() / total_num_label_neg
        inst_pred_conf_true_pos = total_conf_true_pos / total_num_true_pos
        inst_pred_conf_true_neg = total_conf_true_neg / total_num_true_neg
        inst_pred_conf_false_pos = total_conf_false_pos / total_num_false_pos
        inst_pred_conf_false_neg = total_conf_false_neg / total_num_false_neg
        tb_dict.update({'same_inst_loss': same_inst_loss.item(), 'inst_pred_acc': inst_pred_acc,
                        'inst_pred_acc_positive': inst_pred_acc_pos, 'inst_pred_acc_negative': inst_pred_acc_neg,
                        'inst_pred_recall_positive': inst_pred_recall_pos, 'inst_pred_recall_negative': inst_pred_recall_neg,
                        'inst_pred_conf_true_pos': inst_pred_conf_true_pos, 'inst_pred_conf_true_neg': inst_pred_conf_true_neg,
                        'inst_pred_conf_false_pos': inst_pred_conf_false_pos, 'inst_pred_conf_false_neg': inst_pred_conf_false_neg})

        return same_inst_loss, tb_dict

    def pdb(self, **kwargs):
        import pdb
        pdb.set_trace()

    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size: int
                vfe_features: (num_voxels, C)
                points: (num_points, 4 + C), [batch_idx, x, y, z, ...]
                gt_boxes (optional): (B, M, 8)
        Returns:
            batch_dict:
                encoded_spconv_tensor: sparse tensor
                point_features: (N, C)
        """
        train_detect = batch_dict['train_detect']
        batch_size = batch_dict['batch_size']
        points = batch_dict['points']
        # 8/22 gt boxes are needed for evaluation
        #gt_boxes = batch_dict['gt_boxes'] if self.training else None
        gt_boxes = batch_dict['gt_boxes']
        instance_labels = self.assign_instance_targets(points[:, :4], gt_boxes)
        batch_idx, xyz, features = self.break_up_pc(points)

        xyz_batch_cnt = xyz.new_zeros(batch_size).int()
        for bs_idx in range(batch_size):
            xyz_batch_cnt[bs_idx] = (batch_idx == bs_idx).sum()

        assert xyz_batch_cnt.min() == xyz_batch_cnt.max()
        xyz = xyz.view(batch_size, -1, 3)
        features = features.view(batch_size, -1, features.shape[-1]).permute(0, 2, 1) if features is not None else None

        l_xyz, l_features_orig, l_features_inst = [xyz], [features], [features]
        l_labels = [instance_labels]
        initial_idx = xyz.new_zeros(torch.Size((xyz.shape[0], xyz.shape[1])), dtype=torch.int32)
        initial_idx[:, :] = torch.arange(xyz.shape[1]).unsqueeze(0).expand(batch_size, -1)
        l_samples = [initial_idx]
        for i in range(len(self.SA_modules)):
            li_xyz, li_features_orig, li_features_inst, sample_idx = self.SA_modules[i](l_xyz[i], l_features_orig[i], l_features_inst[i])
            # 8/22 gt boxes are needed for evaluation
            #li_labels = torch.gather(l_labels[i], 1, sample_idx.long()) is self.training else None  # gather((B, npoints), 1, (B, nsample)) => (B, nsample)
            li_labels = torch.gather(l_labels[i], 1, sample_idx.long())  # gather((B, npoints), 1, (B, nsample)) => (B, nsample)
            l_xyz.append(li_xyz)
            l_features_orig.append(li_features_orig)
            l_features_inst.append(li_features_inst)
            l_labels.append(li_labels)  ##
            l_samples.append(sample_idx)  ##
        l_neigh = []
        for i in range(-1, -(len(self.FP_modules) + 1), -1):
            l_features_orig[i - 1], l_features_inst[i-1], neigh_dict = self.FP_modules[i](
                l_xyz[i - 1], l_xyz[i], l_features_orig[i - 1], l_features_orig[i], l_features_inst[i - 1], l_features_inst[i], l_labels[i - 1], l_labels[i], train_detect
            )  # (B, C, N)
            if self.training:
                sample_idx = l_samples[i - 1].long()  # (B, num_unknown)
                npoints = sample_idx.shape[1]  # num_unknown
                neigh_idx = neigh_dict['batch_neigh_idx'].long()  # (B, num_unknown, 3)
                self_inst_labels = l_labels[i - 1].unsqueeze(2)
                prev_inst_labels = l_labels[i].unsqueeze(1).expand(-1, npoints, -1)  # expand (B, num_known) to (B, num_unknown, num_known) for gathering neighbor labels
                neigh_inst_labels = torch.gather(prev_inst_labels, 2, neigh_idx) # gather((B, nsample_(i-1), nsample_i), 2, (B, nsample_(i-1), 3)) => (B, nsample_(i-1), 3)
                same_inst_labels = (neigh_inst_labels == self_inst_labels).float()
                same_bg_neighbor = (self_inst_labels == -1) * (neigh_inst_labels == -1)  # will ignore loss from same background neighbor
                same_inst_labels[same_bg_neighbor] = -1
                neigh_dict.update({'batch_neigh_labels': same_inst_labels})
            l_neigh.insert(0, neigh_dict)  # neigh_dict : {'batch_neigh_idx': idx, 'batch_neigh_preds': inst_preds, 'batch_neigh_labels': same_inst_labels}

            # 9/7
            if self.model_cfg.get('SAVE_INST', True):
                npy_save_path = Path('xyz_label_preds')
                xyz_label_preds = neigh_dict['xyz_label_preds'].cpu().numpy()
                layer_index = i + len(self.FP_modules)
                for idx, frame_id in enumerate(batch_dict['frame_id']):
                    np.savetxt(npy_save_path / f'{frame_id}_{layer_index}.txt', xyz_label_preds[idx])

        self.forward_ret_dict = l_neigh

        if batch_dict['train_detect']:
            point_features = l_features_orig[0].permute(0, 2, 1).contiguous()  # (B, N, C)
            batch_dict['point_features'] = point_features.view(-1, point_features.shape[-1])
            batch_dict['point_coords'] = torch.cat((batch_idx[:, None].float(), l_xyz[0].view(-1, 3)), dim=1)
        else:
            point_features_inst = l_features_inst[0].permute(0, 2, 1).contiguous()  # (B, N, C)
            batch_dict['point_features_inst'] = point_features_inst.view(-1, point_features_inst.shape[-1])
            batch_dict['point_coords'] = torch.cat((batch_idx[:, None].float(), l_xyz[0].view(-1, 3)), dim=1)

        return batch_dict


class PointNet2Backbone(nn.Module):
    """
    DO NOT USE THIS CURRENTLY SINCE IT MAY HAVE POTENTIAL BUGS, 20200723
    """
    def __init__(self, model_cfg, input_channels, **kwargs):
        assert False, 'DO NOT USE THIS CURRENTLY SINCE IT MAY HAVE POTENTIAL BUGS, 20200723'
        super().__init__()
        self.model_cfg = model_cfg

        self.SA_modules = nn.ModuleList()
        channel_in = input_channels - 3

        self.num_points_each_layer = []
        skip_channel_list = [input_channels]
        for k in range(self.model_cfg.SA_CONFIG.NPOINTS.__len__()):
            self.num_points_each_layer.append(self.model_cfg.SA_CONFIG.NPOINTS[k])
            mlps = self.model_cfg.SA_CONFIG.MLPS[k].copy()
            channel_out = 0
            for idx in range(mlps.__len__()):
                mlps[idx] = [channel_in] + mlps[idx]
                channel_out += mlps[idx][-1]

            self.SA_modules.append(
                pointnet2_modules_stack.StackSAModuleMSG(
                    radii=self.model_cfg.SA_CONFIG.RADIUS[k],
                    nsamples=self.model_cfg.SA_CONFIG.NSAMPLE[k],
                    mlps=mlps,
                    use_xyz=self.model_cfg.SA_CONFIG.get('USE_XYZ', True),
                )
            )
            skip_channel_list.append(channel_out)
            channel_in = channel_out

        self.FP_modules = nn.ModuleList()

        for k in range(self.model_cfg.FP_MLPS.__len__()):
            pre_channel = self.model_cfg.FP_MLPS[k + 1][-1] if k + 1 < len(self.model_cfg.FP_MLPS) else channel_out
            self.FP_modules.append(
                pointnet2_modules_stack.StackPointnetFPModule(
                    mlp=[pre_channel + skip_channel_list[k]] + self.model_cfg.FP_MLPS[k]
                )
            )

        self.num_point_features = self.model_cfg.FP_MLPS[0][-1]

    def break_up_pc(self, pc):
        batch_idx = pc[:, 0]
        xyz = pc[:, 1:4].contiguous()
        features = (pc[:, 4:].contiguous() if pc.size(-1) > 4 else None)
        return batch_idx, xyz, features

    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size: int
                vfe_features: (num_voxels, C)
                points: (num_points, 4 + C), [batch_idx, x, y, z, ...]
        Returns:
            batch_dict:
                encoded_spconv_tensor: sparse tensor
                point_features: (N, C)
        """
        batch_size = batch_dict['batch_size']
        points = batch_dict['points']
        batch_idx, xyz, features = self.break_up_pc(points)

        xyz_batch_cnt = xyz.new_zeros(batch_size).int()
        for bs_idx in range(batch_size):
            xyz_batch_cnt[bs_idx] = (batch_idx == bs_idx).sum()

        # SA layers
        l_xyz, l_features, l_batch_cnt = [xyz], [features], [xyz_batch_cnt]
        for i in range(len(self.SA_modules)):
            new_xyz_list = []
            for k in range(batch_size):
                if len(l_xyz) == 1:
                    cur_xyz = l_xyz[0][batch_idx == k]
                else:
                    last_num_points = self.num_points_each_layer[i - 1]
                    cur_xyz = l_xyz[-1][k * last_num_points: (k + 1) * last_num_points]
                cur_pt_idxs = pointnet2_utils_stack.furthest_point_sample(
                    cur_xyz[None, :, :].contiguous(), self.num_points_each_layer[i]
                ).long()[0]
                if cur_xyz.shape[0] < self.num_points_each_layer[i]:
                    empty_num = self.num_points_each_layer[i] - cur_xyz.shape[1]
                    cur_pt_idxs[0, -empty_num:] = cur_pt_idxs[0, :empty_num]
                new_xyz_list.append(cur_xyz[cur_pt_idxs])
            new_xyz = torch.cat(new_xyz_list, dim=0)

            new_xyz_batch_cnt = xyz.new_zeros(batch_size).int().fill_(self.num_points_each_layer[i])
            li_xyz, li_features = self.SA_modules[i](
                xyz=l_xyz[i], features=l_features[i], xyz_batch_cnt=l_batch_cnt[i],
                new_xyz=new_xyz, new_xyz_batch_cnt=new_xyz_batch_cnt
            )

            l_xyz.append(li_xyz)
            l_features.append(li_features)
            l_batch_cnt.append(new_xyz_batch_cnt)

        # FP layers
        l_features[0] = points[:, 1:]
        for i in range(-1, -(len(self.FP_modules) + 1), -1):
            l_features[i - 1] = self.FP_modules[i](
                unknown=l_xyz[i - 1], unknown_batch_cnt=l_batch_cnt[i - 1],
                known=l_xyz[i], known_batch_cnt=l_batch_cnt[i],
                unknown_feats=l_features[i - 1], known_feats=l_features[i]
            )

        batch_dict['point_features'] = l_features[0]
        batch_dict['point_coords'] = torch.cat((batch_idx[:, None].float(), l_xyz[0]), dim=1)
        return batch_dict
