import torch
from torch import nn
from torch.nn import functional as F

from .point_head_template import PointHeadTemplate
from ...utils import box_coder_utils, box_utils


class PointHeadBoxID(PointHeadTemplate):
    """
    A simple point-based segmentation head, which are used for PointRCNN.
    Reference Paper: https://arxiv.org/abs/1812.04244
    PointRCNN: 3D Object Proposal Generation and Detection from Point Cloud
    """

    def __init__(self, num_class, input_channels, model_cfg, num_id, predict_boxes_when_training=False, **kwargs):
        super().__init__(model_cfg=model_cfg, num_class=num_class)
        self.num_id = num_id
        self.predict_boxes_when_training = predict_boxes_when_training
        self.cls_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.CLS_FC,
            input_channels=input_channels,
            output_channels=num_class
        )
        target_cfg = self.model_cfg.TARGET_CONFIG
        self.box_coder = getattr(box_coder_utils, target_cfg.BOX_CODER)(
            **target_cfg.BOX_CODER_CONFIG
        )
        self.box_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.REG_FC,
            input_channels=input_channels,
            output_channels=self.box_coder.code_size
        )
        self.reid_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.REID_FC,
            input_channels=input_channels,
            output_channels=self.model_cfg.FEAT_LEN
        )
        self.reid_bn = nn.BatchNorm1d(self.model_cfg.FEAT_LEN)
        if num_id is None:
            self.reid_weights = None
        else:
            self.reid_weights = nn.Parameter(torch.Tensor(num_id, self.model_cfg.FEAT_LEN))
            nn.init.normal_(self.reid_weights, std=0.01)

    def assign_targets(self, input_dict, point_features):
        """
        Args:
            input_dict:
                point_features: (N1 + N2 + N3 + ..., C)
                batch_size:
                point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]
                gt_boxes (optional): (B, M, 8)
        Returns:
            point_cls_labels: (N1 + N2 + N3 + ...), long type, 0:background, -1:ignored
            point_part_labels: (N1 + N2 + N3 + ..., 3)
        """
        point_coords = input_dict['point_coords']
        gt_boxes = input_dict['gt_boxes']
        gt_tracking_ids = input_dict['gt_tracking_ids']
        assert gt_boxes.shape.__len__() == 3, 'gt_boxes.shape=%s' % str(gt_boxes.shape)
        assert point_coords.shape.__len__() in [2], 'points.shape=%s' % str(point_coords.shape)

        batch_size = gt_boxes.shape[0]
        extend_gt_boxes = box_utils.enlarge_box3d(
            gt_boxes.view(-1, gt_boxes.shape[-1]), extra_width=self.model_cfg.TARGET_CONFIG.GT_EXTRA_WIDTH
        ).view(batch_size, -1, gt_boxes.shape[-1])
        targets_dict = self.assign_stack_targets(
            points=point_coords, gt_boxes=gt_boxes, gt_tracking_ids=gt_tracking_ids,
            extend_gt_boxes=extend_gt_boxes, point_features=point_features,
            ret_part_labels=False, ret_box_labels=True, ret_reid_labels=True
        )

        return targets_dict

    def get_loss(self, tb_dict=None):
        tb_dict = {} if tb_dict is None else tb_dict
        # point_loss_cls, tb_dict = self.get_cls_layer_loss(tb_dict)
        # point_loss_box, tb_dict = self.get_box_layer_loss(tb_dict)
        point_loss_reid, tb_dict = self.get_reid_layer_loss(tb_dict)
        # point_loss = point_loss_cls + point_loss_box + point_loss_reid
        return point_loss_reid, tb_dict

    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size:
                point_features: (N1 + N2 + N3 + ..., C) or (B, N, C)
                point_features_before_fusion: (N1 + N2 + N3 + ..., C)
                point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]
                point_labels (optional): (N1 + N2 + N3 + ...)
                gt_boxes (optional): (B, M, 8)
        Returns:
            batch_dict:
                point_cls_scores: (N1 + N2 + N3 + ..., 1)
                point_part_offset: (N1 + N2 + N3 + ..., 3)
        """
        if self.model_cfg.get('USE_POINT_FEATURES_BEFORE_FUSION', False):
            point_features = batch_dict['point_features_before_fusion']
        else:
            point_features = batch_dict['point_features']
        point_cls_preds = self.cls_layers(point_features)  # (total_points, num_class)
        point_reid_preds = self.reid_layers(point_features)

        point_cls_scores, pred_classes = point_cls_preds.max(dim=-1)

        if self.training:
            targets_dict = self.assign_targets(batch_dict, point_reid_preds)
            point_cls_labels = targets_dict['point_cls_labels']
            fg_flag = point_cls_labels > 0

            point_box_preds = self.box_layers(point_features[fg_flag])  # (fg_points, box_code_size)
            self.forward_ret_dict = {
                'point_cls_preds': point_cls_preds,
                'point_box_preds': point_box_preds,
                'point_cls_labels': point_cls_labels,
                'point_box_labels': targets_dict['point_box_labels'][fg_flag]
            }

            reid_features = targets_dict['reid_features']
            reid_labels = targets_dict['reid_labels']
            if len(reid_labels) >= 2:
                bn_features = self.reid_bn(reid_features)
                reid_preds = F.linear(F.normalize(bn_features), F.normalize(self.reid_weights))
                self.forward_ret_dict['reid_preds'] = reid_preds
                self.forward_ret_dict['reid_feats'] = bn_features
                self.forward_ret_dict['reid_labels'] = reid_labels

        if not self.training or self.predict_boxes_when_training:
            batch_dict['point_cls_scores'] = torch.sigmoid(point_cls_scores)  # used in roi pooling
            fg_flag = point_cls_scores >= self.model_cfg.get('SCORE_THRESH', 0.0)
            fg_point_features = point_features[fg_flag]
            if len(fg_point_features):
                point_box_preds = self.box_layers(fg_point_features)  # (fg_points, box_code_size)
                point_box_preds = self.box_coder.decode_torch(
                    point_box_preds, batch_dict['point_coords'][:, 1:4][fg_flag], pred_classes[fg_flag] + 1
                )
            else:
                point_box_preds = fg_point_features

            batch_dict['batch_cls_preds'] = point_cls_preds[fg_flag]
            batch_dict['batch_box_preds'] = point_box_preds
            batch_dict['batch_index'] = batch_dict['point_coords'][:, 0][fg_flag]
            batch_dict['batch_feat_preds'] = self.reid_bn(point_reid_preds)
            batch_dict['cls_preds_normalized'] = False

        return batch_dict

import torch
from torch import nn
from torch.nn import functional as F

from .point_head_template import PointHeadTemplate
from ...utils import box_coder_utils, box_utils


class PointHeadBoxID(PointHeadTemplate):
    """
    A simple point-based segmentation head, which are used for PointRCNN.
    Reference Paper: https://arxiv.org/abs/1812.04244
    PointRCNN: 3D Object Proposal Generation and Detection from Point Cloud
    """

    def __init__(self, num_class, input_channels, model_cfg, num_id, predict_boxes_when_training=False, **kwargs):
        super().__init__(model_cfg=model_cfg, num_class=num_class)
        self.num_id = num_id
        self.predict_boxes_when_training = predict_boxes_when_training
        self.cls_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.CLS_FC,
            input_channels=input_channels,
            output_channels=num_class
        )
        target_cfg = self.model_cfg.TARGET_CONFIG
        self.box_coder = getattr(box_coder_utils, target_cfg.BOX_CODER)(
            **target_cfg.BOX_CODER_CONFIG
        )
        self.box_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.REG_FC,
            input_channels=input_channels,
            output_channels=self.box_coder.code_size
        )
        self.reid_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.REID_FC,
            input_channels=input_channels,
            output_channels=self.model_cfg.FEAT_LEN
        )
        self.reid_bn = nn.BatchNorm1d(self.model_cfg.FEAT_LEN)
        if num_id is None:
            self.reid_weights = None
        else:
            self.reid_weights = nn.Parameter(torch.Tensor(num_id, self.model_cfg.FEAT_LEN))
            nn.init.normal_(self.reid_weights, std=0.01)

    def assign_targets(self, input_dict, point_features):
        """
        Args:
            input_dict:
                point_features: (N1 + N2 + N3 + ..., C)
                batch_size:
                point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]
                gt_boxes (optional): (B, M, 8)
        Returns:
            point_cls_labels: (N1 + N2 + N3 + ...), long type, 0:background, -1:ignored
            point_part_labels: (N1 + N2 + N3 + ..., 3)
        """
        point_coords = input_dict['point_coords']
        gt_boxes = input_dict['gt_boxes']
        gt_tracking_ids = input_dict['gt_tracking_ids']
        assert gt_boxes.shape.__len__() == 3, 'gt_boxes.shape=%s' % str(gt_boxes.shape)
        assert point_coords.shape.__len__() in [2], 'points.shape=%s' % str(point_coords.shape)

        batch_size = gt_boxes.shape[0]
        extend_gt_boxes = box_utils.enlarge_box3d(
            gt_boxes.view(-1, gt_boxes.shape[-1]), extra_width=self.model_cfg.TARGET_CONFIG.GT_EXTRA_WIDTH
        ).view(batch_size, -1, gt_boxes.shape[-1])
        targets_dict = self.assign_stack_targets(
            points=point_coords, gt_boxes=gt_boxes, gt_tracking_ids=gt_tracking_ids,
            extend_gt_boxes=extend_gt_boxes, point_features=point_features,
            ret_part_labels=False, ret_box_labels=True, ret_reid_labels=True
        )

        return targets_dict

    def get_loss(self, tb_dict=None):
        tb_dict = {} if tb_dict is None else tb_dict
        # point_loss_cls, tb_dict = self.get_cls_layer_loss(tb_dict)
        # point_loss_box, tb_dict = self.get_box_layer_loss(tb_dict)
        point_loss_reid, tb_dict = self.get_reid_layer_loss(tb_dict)
        # point_loss = point_loss_cls + point_loss_box + point_loss_reid
        return point_loss_reid, tb_dict

    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size:
                point_features: (N1 + N2 + N3 + ..., C) or (B, N, C)
                point_features_before_fusion: (N1 + N2 + N3 + ..., C)
                point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]
                point_labels (optional): (N1 + N2 + N3 + ...)
                gt_boxes (optional): (B, M, 8)
        Returns:
            batch_dict:
                point_cls_scores: (N1 + N2 + N3 + ..., 1)
                point_part_offset: (N1 + N2 + N3 + ..., 3)
        """
        if self.model_cfg.get('USE_POINT_FEATURES_BEFORE_FUSION', False):
            point_features = batch_dict['point_features_before_fusion']
        else:
            point_features = batch_dict['point_features']
        point_cls_preds = self.cls_layers(point_features)  # (total_points, num_class)
        point_reid_preds = self.reid_layers(point_features)

        point_cls_scores, pred_classes = point_cls_preds.max(dim=-1)

        if self.training:
            targets_dict = self.assign_targets(batch_dict, point_reid_preds)
            point_cls_labels = targets_dict['point_cls_labels']
            fg_flag = point_cls_labels > 0

            point_box_preds = self.box_layers(point_features[fg_flag])  # (fg_points, box_code_size)
            self.forward_ret_dict = {
                'point_cls_preds': point_cls_preds,
                'point_box_preds': point_box_preds,
                'point_cls_labels': point_cls_labels,
                'point_box_labels': targets_dict['point_box_labels'][fg_flag]
            }

            reid_features = targets_dict['reid_features']
            reid_labels = targets_dict['reid_labels']
            if len(reid_labels) >= 2:
                bn_features = self.reid_bn(reid_features)
                reid_preds = F.linear(F.normalize(bn_features), F.normalize(self.reid_weights))
                self.forward_ret_dict['reid_preds'] = reid_preds
                self.forward_ret_dict['reid_feats'] = bn_features
                self.forward_ret_dict['reid_labels'] = reid_labels

        if not self.training or self.predict_boxes_when_training:
            batch_dict['point_cls_scores'] = torch.sigmoid(point_cls_scores)  # used in roi pooling
            fg_flag = point_cls_scores >= self.model_cfg.get('SCORE_THRESH', 0.0)
            fg_point_features = point_features[fg_flag]
            if len(fg_point_features):
                point_box_preds = self.box_layers(fg_point_features)  # (fg_points, box_code_size)
                point_box_preds = self.box_coder.decode_torch(
                    point_box_preds, batch_dict['point_coords'][:, 1:4][fg_flag], pred_classes[fg_flag] + 1
                )
            else:
                point_box_preds = fg_point_features

            batch_dict['batch_cls_preds'] = point_cls_preds[fg_flag]
            batch_dict['batch_box_preds'] = point_box_preds
            batch_dict['batch_index'] = batch_dict['point_coords'][:, 0][fg_flag]
            batch_dict['batch_feat_preds'] = self.reid_bn(point_reid_preds)
            batch_dict['cls_preds_normalized'] = False

        return batch_dict

A simple point-based segmentation head, which are used for PointRCNN.     Reference Paper: https://arxiv.org/abs/1812.04244     PointRCNN: 3D Object Proposal Generation and Detection from Point Cloud Args:             input_dict:                 point_features: (N1 + N2 + N3 + ..., C)                 batch_size:                 point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]                 gt_boxes (optional): (B, M, 8)         Returns:             point_cls_labels: (N1 + N2 + N3 + ...), long type, 0:background, -1:ignored             point_part_labels: (N1 + N2 + N3 + ..., 3) # point_loss_cls, tb_dict = self.get_cls_layer_loss(tb_dict) # point_loss_box, tb_dict = self.get_box_layer_loss(tb_dict) # point_loss = point_loss_cls + point_loss_box + point_loss_reid Args:             batch_dict:                 batch_size:                 point_features: (N1 + N2 + N3 + ..., C) or (B, N, C)                 point_features_before_fusion: (N1 + N2 + N3 + ..., C)                 point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]                 point_labels (optional): (N1 + N2 + N3 + ...)                 gt_boxes (optional): (B, M, 8)         Returns:             batch_dict:                 point_cls_scores: (N1 + N2 + N3 + ..., 1)                 point_part_offset: (N1 + N2 + N3 + ..., 3) # (total_points, num_class) # (fg_points, box_code_size) # used in roi pooling # (fg_points, box_code_size)