# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch
from torch import nn as nn
from torch.nn import functional as F

from mmseg.core import add_prefix
from mmseg.models import SEGMENTORS
from ..builder import build_backbone, build_head, build_loss, build_neck, build_voxel_encoder
from .base import Base3DSegmentor
from mmdet3d.ops import CylinderVoxelization, FastCylinderVoxelization
import time

@SEGMENTORS.register_module()
class Cylinder3D(Base3DSegmentor):
    """3D Encoder Decoder segmentors.

    EncoderDecoder typically consists of backbone, decode_head, auxiliary_head.
    Note that auxiliary_head is only used for deep supervision during training,
    which could be thrown during inference.
    """

    def __init__(self,
                 voxel_layer,
                 voxel_encoder,
                 #middle_encoder=None,
                 backbone,
                 decode_head,
                 ignore_index=0,
                 neck=None,
                 bbox_head=None,
                 train_cfg=None,
                 test_cfg=None,
                 init_cfg=None,
                 is_fast=True
                 ):
        super(Cylinder3D, self).__init__(init_cfg=init_cfg)
        if is_fast is True:
            self.voxel_layer = FastCylinderVoxelization(**voxel_layer)
        else:
            self.voxel_layer = CylinderVoxelization(**voxel_layer)
        self.voxel_encoder = build_voxel_encoder(voxel_encoder)
        self.backbone = build_backbone(backbone)
        if neck is not None:
            self.neck = build_neck(neck)
        self.decode_head = build_head(decode_head)
        # self._init_decode_head(decode_head)
        # self._init_auxiliary_head(auxiliary_head)
        # self._init_loss_regularization(loss_regularization)

        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
        # assert self.with_decode_head, \
        #     '3D EncoderDecoder Segmentor should have a decode_head'
            
    def aug_test(self):
        pass
    def encode_decode(self):
        pass
    def extract_feat(self):
        pass
    def forward_train(self, points, img_metas, pts_semantic_mask):
        time1=time.time()
        data_tuple = self.voxel_layer([points,pts_semantic_mask])
        _, train_vox_label, train_grid, _, train_pt_fea = data_tuple
        time2=time.time()
        coords, features_3d = self.voxel_encoder(train_pt_fea, train_grid)
        time3=time.time()
        spatial_features = self.backbone(features_3d, coords, len(train_vox_label))#batch_size
        time4=time.time()
        losses =self.decode_head.forward_train(spatial_features,train_vox_label)
        time5=time.time()

        print('voxel cost',time2-time1,'s')
        print('encoder cost',time3-time2,'s')
        print('backbone cost',time4-time3,'s')
        print('deconder cost',time5-time4,'s')

        return losses

    def inference(self,points, img_meta,rescale):
        data_tuple = self.voxel_layer([points])
        _, train_grid, train_pt_fea = data_tuple
        coords, features_3d = self.voxel_encoder(train_pt_fea, train_grid)
        spatial_features = self.backbone(features_3d, coords, len(train_grid))
        seg_logits =self.decode_head.forward_test(spatial_features)
        return seg_logits, train_grid

    def simple_test(self, points, img_metas, rescale=True):
        """Simple test with single scene.

        Args:
            points (list[torch.Tensor]): List of points of shape [N, 3+C].
            img_metas (list[dict]): Meta information of each sample.
            rescale (bool): Whether transform to original number of points.
                Will be used for voxelization based segmentors.
                Defaults to True.

        Returns:
            list[dict]: The output prediction result with following keys:

                - semantic_mask (Tensor): Segmentation mask of shape [N].
        """
        # 3D segmentation requires per-point prediction, so it's impossible
        # to use down-sampling to get a batch of scenes with same num_points
        # therefore, we only support testing one scene every time


        seg_prob, train_grid = self.inference(points, img_metas,rescale)
        # for point, img_meta in zip(points, img_metas):
        #     seg_prob, train_grid = self.inference(point.unsqueeze(0), [img_meta],
        #                               rescale)
        #     seg_map = seg_prob[0].argmax(0)  # [N]

            
        #     train_grid = train_grid[0]
        #     #print(seg_map.cpu().numpy().tolist())
        #     # to cpu tensor for consistency with det3d
        #     #np.save('b.npy',seg_map.cpu().numpy()) 
        #     seg_map = seg_map.cpu()
        #     predict_labels=seg_map[train_grid[:,0],train_grid[:,1],train_grid[:,2]]
        #     np.save('b.npy',predict_labels.numpy()) 
        #     #seg_pred.append(seg_map)

        #     seg_pred.append(predict_labels)
        # warp in dict
        #seg_pred = [dict(semantic_mask=seg_map) for seg_map in seg_pred]
        seg_map = seg_prob.argmax(1).cpu()
        predict_labels=[]
        for i in range(seg_map.shape[0]):
            predict_labels.append(seg_map[i,train_grid[i][:,0],train_grid[i][:,1],train_grid[i][:,2]])

        seg_pred = [dict(semantic_mask=predict_label) for predict_label in predict_labels]

        return seg_pred
    
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch
from torch import nn as nn
from torch.nn import functional as F

from mmseg.core import add_prefix
from mmseg.models import SEGMENTORS
from ..builder import build_backbone, build_head, build_loss, build_neck, build_voxel_encoder
from .base import Base3DSegmentor
from mmdet3d.ops import CylinderVoxelization, FastCylinderVoxelization
import time

@SEGMENTORS.register_module()
class Cylinder3D(Base3DSegmentor):
    """3D Encoder Decoder segmentors.

    EncoderDecoder typically consists of backbone, decode_head, auxiliary_head.
    Note that auxiliary_head is only used for deep supervision during training,
    which could be thrown during inference.
    """

    def __init__(self,
                 voxel_layer,
                 voxel_encoder,
                 #middle_encoder=None,
                 backbone,
                 decode_head,
                 ignore_index=0,
                 neck=None,
                 bbox_head=None,
                 train_cfg=None,
                 test_cfg=None,
                 init_cfg=None,
                 is_fast=True
                 ):
        super(Cylinder3D, self).__init__(init_cfg=init_cfg)
        if is_fast is True:
            self.voxel_layer = FastCylinderVoxelization(**voxel_layer)
        else:
            self.voxel_layer = CylinderVoxelization(**voxel_layer)
        self.voxel_encoder = build_voxel_encoder(voxel_encoder)
        self.backbone = build_backbone(backbone)
        if neck is not None:
            self.neck = build_neck(neck)
        self.decode_head = build_head(decode_head)
        # self._init_decode_head(decode_head)
        # self._init_auxiliary_head(auxiliary_head)
        # self._init_loss_regularization(loss_regularization)

        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
        # assert self.with_decode_head, \
        #     '3D EncoderDecoder Segmentor should have a decode_head'
            
    def aug_test(self):
        pass
    def encode_decode(self):
        pass
    def extract_feat(self):
        pass
    def forward_train(self, points, img_metas, pts_semantic_mask):
        time1=time.time()
        data_tuple = self.voxel_layer([points,pts_semantic_mask])
        _, train_vox_label, train_grid, _, train_pt_fea = data_tuple
        time2=time.time()
        coords, features_3d = self.voxel_encoder(train_pt_fea, train_grid)
        time3=time.time()
        spatial_features = self.backbone(features_3d, coords, len(train_vox_label))#batch_size
        time4=time.time()
        losses =self.decode_head.forward_train(spatial_features,train_vox_label)
        time5=time.time()

        print('voxel cost',time2-time1,'s')
        print('encoder cost',time3-time2,'s')
        print('backbone cost',time4-time3,'s')
        print('deconder cost',time5-time4,'s')

        return losses

    def inference(self,points, img_meta,rescale):
        data_tuple = self.voxel_layer([points])
        _, train_grid, train_pt_fea = data_tuple
        coords, features_3d = self.voxel_encoder(train_pt_fea, train_grid)
        spatial_features = self.backbone(features_3d, coords, len(train_grid))
        seg_logits =self.decode_head.forward_test(spatial_features)
        return seg_logits, train_grid

    def simple_test(self, points, img_metas, rescale=True):
        """Simple test with single scene.

        Args:
            points (list[torch.Tensor]): List of points of shape [N, 3+C].
            img_metas (list[dict]): Meta information of each sample.
            rescale (bool): Whether transform to original number of points.
                Will be used for voxelization based segmentors.
                Defaults to True.

        Returns:
            list[dict]: The output prediction result with following keys:

                - semantic_mask (Tensor): Segmentation mask of shape [N].
        """
        # 3D segmentation requires per-point prediction, so it's impossible
        # to use down-sampling to get a batch of scenes with same num_points
        # therefore, we only support testing one scene every time


        seg_prob, train_grid = self.inference(points, img_metas,rescale)
        # for point, img_meta in zip(points, img_metas):
        #     seg_prob, train_grid = self.inference(point.unsqueeze(0), [img_meta],
        #                               rescale)
        #     seg_map = seg_prob[0].argmax(0)  # [N]

            
        #     train_grid = train_grid[0]
        #     #print(seg_map.cpu().numpy().tolist())
        #     # to cpu tensor for consistency with det3d
        #     #np.save('b.npy',seg_map.cpu().numpy()) 
        #     seg_map = seg_map.cpu()
        #     predict_labels=seg_map[train_grid[:,0],train_grid[:,1],train_grid[:,2]]
        #     np.save('b.npy',predict_labels.numpy()) 
        #     #seg_pred.append(seg_map)

        #     seg_pred.append(predict_labels)
        # warp in dict
        #seg_pred = [dict(semantic_mask=seg_map) for seg_map in seg_pred]
        seg_map = seg_prob.argmax(1).cpu()
        predict_labels=[]
        for i in range(seg_map.shape[0]):
            predict_labels.append(seg_map[i,train_grid[i][:,0],train_grid[i][:,1],train_grid[i][:,2]])

        seg_pred = [dict(semantic_mask=predict_label) for predict_label in predict_labels]

        return seg_pred
    
# Copyright (c) OpenMMLab. All rights reserved. 3D Encoder Decoder segmentors.      EncoderDecoder typically consists of backbone, decode_head, auxiliary_head.     Note that auxiliary_head is only used for deep supervision during training,     which could be thrown during inference. #middle_encoder=None, # self._init_decode_head(decode_head) # self._init_auxiliary_head(auxiliary_head) # self._init_loss_regularization(loss_regularization) # assert self.with_decode_head, \ #     '3D EncoderDecoder Segmentor should have a decode_head' #batch_size Simple test with single scene.          Args:             points (list[torch.Tensor]): List of points of shape [N, 3+C].             img_metas (list[dict]): Meta information of each sample.             rescale (bool): Whether transform to original number of points.                 Will be used for voxelization based segmentors.                 Defaults to True.          Returns:             list[dict]: The output prediction result with following keys:                  - semantic_mask (Tensor): Segmentation mask of shape [N]. # 3D segmentation requires per-point prediction, so it's impossible # to use down-sampling to get a batch of scenes with same num_points # therefore, we only support testing one scene every time # for point, img_meta in zip(points, img_metas): #     seg_prob, train_grid = self.inference(point.unsqueeze(0), [img_meta], #                               rescale) #     seg_map = seg_prob[0].argmax(0)  # [N] #     train_grid = train_grid[0] #     #print(seg_map.cpu().numpy().tolist()) #     # to cpu tensor for consistency with det3d #     #np.save('b.npy',seg_map.cpu().numpy()) #     seg_map = seg_map.cpu() #     predict_labels=seg_map[train_grid[:,0],train_grid[:,1],train_grid[:,2]] #     np.save('b.npy',predict_labels.numpy()) #     #seg_pred.append(seg_map) #     seg_pred.append(predict_labels) # warp in dict #seg_pred = [dict(semantic_mask=seg_map) for seg_map in seg_pred]