import numpy as np 

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

from ...utils import box_utils, loss_utils
from .point_head_template import PointHeadTemplate
from ...ops.torch_hash import RadiusGraph
from pcdet.models.model_utils import graph_utils
from torch_scatter import scatter


class StratHead_InsCls(PointHeadTemplate):
    """
    A simple point-based segmentation head, which are used for PV-RCNN keypoint segmentaion.
    Reference Paper: https://arxiv.org/abs/1912.13192
    PV-RCNN: Point-Voxel Feature Set Abstraction for 3D Object Detection
    """
    def __init__(self, runtime_cfg, model_cfg, **kwargs):
        num_class = runtime_cfg['num_seg_classes']
        input_channels = runtime_cfg['input_channels']
        # cr = kwargs.get('cr', 1.0)
        
        super().__init__(model_cfg=model_cfg,
                         num_class=num_class)
        self.scale = runtime_cfg.get('scale', 1.0)
        self.assign_to_point = model_cfg.get("ASSIGN_TO_POINT", False)
        self.cls_layers = self.make_fc_layers(
            fc_cfg=[int(c*self.scale) for c in self.model_cfg.CLS_FC],
            input_channels=input_channels,
            output_channels=num_class,
            dropout=self.dropout
        )

        self.embedding_channel = model_cfg.get("EMBEDDING_CHANNEL", 64)
        self.proj_channel = model_cfg.get("PROJ_CHANNEL", 64)
        self.ins_cls = model_cfg.get("INS_CLS", False)
        self.contra_weight = model_cfg.get("CONTRA_WEIGHT", 1.0)
        self.ins_cls_weight = model_cfg.get("INS_CLS_LOSS_WEIGHT", False)
        self.global_fea_op = model_cfg.get("GLOBAL_FEA_OP", "max")
        self.n_ins_cls = model_cfg.get("N_INS_CLS", 7)
        self.more_channel = model_cfg.get("MORE_CHANNEL", False)
        self.multi_scale = model_cfg.get("MULTI_SCALE", False)
        self.append_points = model_cfg.get("APPEND_POINTS", False)
        self.append_rawpoint_features = model_cfg.get("APPEND_RAWPOINT_FEATURES", False)
        self.raw_point_feature_channel = model_cfg.get("RAW_POINT_FEATURE_CHANNEL", 64)
        self.random_ins = model_cfg.get("RANDOM_INS", False)
        self.random_keep_ratio = model_cfg.get("RANDOM_KEEP_RATIO", 1.0)

        embed_input_channel = input_channels
        embed_fc = self.model_cfg.EMBED_FC
        
        if self.n_ins_cls==12:
            self.sem_cls = [1,2,3,4,5,6,7,8,14,16,18,19]
            self.ins_cls_mapping = torch.zeros(20).cuda()
            cnt = 1
            for s_i in self.sem_cls:
                self.ins_cls_mapping[s_i] = cnt
                cnt += 1
        elif self.n_ins_cls==11:
            self.sem_cls = [1,2,3,4,5,6,7,8,16,18,19]
            self.ins_cls_mapping = torch.zeros(20).cuda()
            cnt = 1
            for s_i in self.sem_cls:
                self.ins_cls_mapping[s_i] = cnt
                cnt += 1
        elif self.n_ins_cls==19:
            self.sem_cls = [1,2,3,4,5,6,7,8,9,11,12,13,14,15,16,17,18,20 ]
            self.ins_cls_mapping = torch.zeros(21).cuda()
            cnt = 1
            for s_i in self.sem_cls:
                self.ins_cls_mapping[s_i] = cnt
                cnt += 1
        else:
            self.ins_cls_mapping = None

        if self.append_points:
            embed_input_channel += 32*3 
            embed_fc = [128, 64,]
            if self.more_channel:
                embed_fc = [128, 128, 64,]

        if self.multi_scale: 
            embed_input_channel += 384 + 192 + 96
            embed_fc = [512, 256,]
            if self.more_channel:
                embed_fc = [256, 128, 64]
        

        if self.ins_cls:
            self.embed_layers = self.make_fc_layers(
                fc_cfg=[int(c*self.scale) for c in embed_fc],
                input_channels=embed_input_channel,
                output_channels=self.embedding_channel,
                dropout=False,
            )
        if self.append_rawpoint_features:
            self.raw_point_layers = self.make_fc_layers(
                fc_cfg=[int(c*self.scale) for c in self.model_cfg.RAW_POINT_FC],
                input_channels=6,
                output_channels=self.raw_point_feature_channel,
                dropout=False,
            )
            ins_cls_input_channel = self.embedding_channel+ self.raw_point_feature_channel
            ins_cls_fc = [128, 64, 64]
        else:
            ins_cls_input_channel = self.embedding_channel
            ins_cls_fc = self.model_cfg.INS_CLS_FC
            if self.more_channel:
                ins_cls_fc = [128, 64, 64, 32]

        if self.ins_cls:
            self.ins_cls_layers = self.make_fc_layers(
                fc_cfg=[int(c*self.scale) for c in ins_cls_fc],
                input_channels=ins_cls_input_channel,
                output_channels=self.n_ins_cls,
                dropout=False,
            )

        self.build_losses(self.model_cfg.LOSS_CONFIG)
        self.target_assigner_cfg = self.model_cfg.get("TARGET_ASSIGNER", None)
        if self.target_assigner_cfg is not None:
            max_num_points = self.target_assigner_cfg.get("MAX_NUM_POINTS", None)
            self.radius_graph = RadiusGraph(max_num_points=max_num_points, ndim=3)
        #if self.assign_to_point:
        self.graph = graph_utils.KNNGraph({}, dict(NUM_NEIGHBORS=1))
        if self.append_points:
            self.ins_graph = graph_utils.KNNGraph({}, dict(NUM_NEIGHBORS=32))
    
    def build_losses(self, losses_cfg):
        if not isinstance(losses_cfg['LOSS'], list):
            losses_cfg['LOSS'] = [losses_cfg['LOSS']]
        if not isinstance(losses_cfg['WEIGHT'], list):
            losses_cfg['WEIGHT'] = [losses_cfg['WEIGHT']]
        self.loss_names = losses_cfg['LOSS']
        self.losses = nn.ModuleList()
        self.loss_weight = []
        for loss, weight in zip(losses_cfg['LOSS'], losses_cfg['WEIGHT']):
            self.losses.append(
                loss_utils.LOSSES[loss](loss_cfg=losses_cfg)
            )
            self.loss_weight.append(weight)
 
    def get_cls_loss(self, preds, labels, tb_dict=None, index=0,prefix=None):

        if tb_dict is None:
            tb_dict = {}
        point_loss = 0.0
        loss_module = self.losses[index]
        loss_name = self.loss_names[index]
        loss_this = loss_module(preds, labels)
        if prefix is None:
            tb_dict[loss_name] = loss_this.item()
        else:
            tb_dict[f'{prefix}/{loss_name}'] = loss_this.item()
        point_loss += loss_this
        return point_loss, tb_dict

    def get_loss(self, tb_dict=None, prefix=None):
        tb_dict = {} if tb_dict is None else tb_dict
        point_cls_labels = self.forward_ret_dict[self.gt_seg_cls_label_key].view(-1).long()
        point_cls_preds = self.forward_ret_dict['pred_seg_cls_logits'].view(-1, self.num_class)
        point_loss_cls, tb_dict_1 = self.get_cls_loss(point_cls_preds, point_cls_labels,)

        loss_contra = torch.Tensor([0.0]).cuda()
        tb_dict_2 = {} 
            
        if self.ins_cls:
            ins_cls_labels = self.forward_ret_dict['ins_cls_label'].view(-1).long()
            ins_cls_preds = self.forward_ret_dict['ins_cls_pred'].view(-1, self.n_ins_cls)
            loss_ins_cls, _ = self.get_cls_loss(ins_cls_preds, ins_cls_labels,)
        else:
            loss_ins_cls = torch.Tensor([0.0]).cuda()

        point_loss = point_loss_cls + self.contra_weight * loss_contra +  self.ins_cls_weight * loss_ins_cls 

        tb_dict.update(tb_dict_1)
        tb_dict.update(tb_dict_2)
        
        tb_dict['loss_contra'] = loss_contra.item()
        tb_dict['loss_seg'] = point_loss_cls.item()
        tb_dict['loss_ins_cls'] = loss_ins_cls.item()

        iou_stats, _ = self.get_iou_statistics()
        ups, downs = iou_stats[0]['ups'], iou_stats[0]['downs']
        for iou_stat in iou_stats[1:]:
            ups += iou_stat['ups']
            downs += iou_stat['downs']
        ious = ups / torch.clamp(downs, min=1.0)
        for i in range(self.num_class):
            if downs[i] > 0:
                if prefix is None:
                    tb_dict.update({f'per_class/IoU_{i}': ious[i]})
                else:
                    tb_dict.update({f'{prefix}/per_class/IoU_{i}': ious[i]})
        if prefix is None:
            tb_dict.update({f'IoU_FG': ups[1:5].sum()/torch.clamp(downs[1:5].sum(), min=1.0),
                            f'IoU_BG': ups[5:].sum()/torch.clamp(downs[5:].sum(), min=1.0),
                            })
            tb_dict.update({f'mIoU': ious.mean()})
            tb_dict.update({f'loss': point_loss.item()})
        else:
            tb_dict.update({f'{prefix}/IoU_FG': ups[1:5].sum()/torch.clamp(downs[1:5].sum(), min=1.0),
                            f'{prefix}/IoU_BG': ups[5:].sum()/torch.clamp(downs[5:].sum(), min=1.0),
                            })
            tb_dict.update({f'{prefix}/mIoU': ious.mean()})
            tb_dict.update({f'{prefix}/loss': point_loss.item()})
        
        return point_loss, tb_dict


    def get_iou_statistics(self):
        pred_dicts = self.get_evaluation_results()
        iou_dicts = []
        iou_dict = dict(
            ups=None,
            downs=None,
        )
        for pred_dict in pred_dicts:
            pred_labels = pred_dict['point_wise']['pred_segmentation_label']
            gt_labels = pred_dict['point_wise']['gt_segmentation_label']
            ups = pred_labels.new_zeros(self.num_class)
            downs = pred_labels.new_zeros(self.num_class)
            pred_labels[gt_labels == 0] = 0
            for cls in range(self.num_class):
                pred_mask = pred_labels == cls
                gt_mask = gt_labels == cls
                ups[cls] = (pred_mask & gt_mask).sum()
                downs[cls] = (pred_mask | gt_mask).sum()
            
            iou_dict['ups'] = ups if iou_dict['ups'] is None else iou_dict['ups'] + ups
            iou_dict['downs'] = downs if iou_dict['downs'] is None else iou_dict['downs'] + downs
            iou_dicts.append(
                dict(
                    ups = ups,
                    downs = downs
                )
            )
        return iou_dicts, iou_dict

    def get_evaluation_results(self):
        pred_logits = self.forward_ret_dict['pred_seg_cls_logits']
        pred_scores = torch.sigmoid(pred_logits)
        batch_idx = self.forward_ret_dict['batch_idx']
        pred_dicts = []
        point_bxyz = self.forward_ret_dict['point_bxyz']
        
        for i in range(self.forward_ret_dict['batch_size']):
            bs_mask = batch_idx == i
            point_xyz = point_bxyz[bs_mask, 1:4]
            pred_confidences, pred_labels = pred_scores[bs_mask].max(-1)
            gt_labels = self.forward_ret_dict[self.gt_seg_cls_label_key][bs_mask]
            valid_mask = (gt_labels >= 0)
            pred_labels = pred_labels[valid_mask]
            gt_labels = gt_labels[valid_mask]
            record_dict = dict(
                point_wise=dict(
                    gt_segmentation_label=gt_labels,
                    pred_segmentation_label=pred_labels,
                    point_xyz=point_xyz,
                ),
                object_wise=dict(),
                voxel_wise=dict(),
                scene_wise=dict(
                    num_seg_class=self.num_class,
                ),
            )
            pred_dicts.append(record_dict)
        return pred_dicts

    def assign_targets(self, target_assigner_cfg, batch_dict):
        ref_label = batch_dict[target_assigner_cfg["REF_SEGMENTATION_LABEL"]]
        ref_bxyz = batch_dict[target_assigner_cfg["REF_POINT_BXYZ"]]
        query_bxyz = batch_dict[target_assigner_cfg["QUERY_POINT_BXYZ"]]
        query_label_key = target_assigner_cfg["QUERY_SEGMENTATION_LABEL"]

        radius = target_assigner_cfg["RADIUS"]
        er, eq = self.radius_graph(ref_bxyz, query_bxyz, radius, 1, sort_by_dist=True)

        query_label = ref_label.new_full(query_bxyz.shape[:1], 0) # by default, assuming class 0 is ignored
        query_label[eq] = ref_label[er]
        
        batch_dict[query_label_key] = query_label

    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)
        """
        point_features = batch_dict[self.point_feature_key]
        point_pred_logits = self.cls_layers(point_features)  # (total_points, num_class)
        point_pred_scores = torch.sigmoid(point_pred_logits)
        voxel_features =  batch_dict[self.point_feature_key] 
            
        if 'strat_coord_min' in batch_dict.keys():
            for ii in range(batch_dict['batch_size']):
                batch_idx = batch_dict[self.batch_key][:, 0]==ii
                batch_dict[self.batch_key][batch_idx, 1:] += batch_dict['strat_coord_min'][ii]


        if self.training:
            # instance_label: instances labels for seg label 1-7 
            # instance_label_back: instances labels for seg label 1-13

            #instance_label_or = batch_dict['voxel_instance_label'] - 1
            if not 'strat_instance_label_back' in  batch_dict.keys():
                batch_dict['strat_instance_label_back'] = batch_dict['strat_instance_label'].long()
                batch_dict['instance_label_back'] = batch_dict['instance_label']

            instance_label_or = batch_dict['strat_instance_label_back'] 
            pt_instance_label_or = batch_dict['instance_label_back']

            segmentation_label = batch_dict['strat_segmentation_label']
            batch_dim = batch_dict[self.batch_key][:, 0]
            instance_label = instance_label_or.clone()
            pt_instance_label = pt_instance_label_or.clone() 
            
            # add an offset for instance labels in different data samples 
            last_num = 0
            for b_i in range(batch_dict['batch_size']):
                ins_mask_tmp = (batch_dim==b_i) & (instance_label_or>=0)
                pt_ins_mask_tmp = (batch_dict['point_bxyz'][:,0]==b_i) & (pt_instance_label_or>=0)
                if ins_mask_tmp.sum()==0:
                    continue 
                instance_label[ins_mask_tmp] += last_num
                pt_instance_label[pt_ins_mask_tmp] += last_num
                last_num = instance_label[ins_mask_tmp].max() +1
            
            if self.random_ins:
                drop_mask = (torch.rand(instance_label.shape)<self.random_keep_ratio).int()
                drop_mask = (drop_mask == 1).nonzero(as_tuple=True)[0]
                instance_label[drop_mask] = -1

            flag = False
            if instance_label.max()==-1:
                print('$$$$$$$$$$$$$$$')
                flag = True
                instance_label[:100]+=1
                pt_instance_label[:100]+=1
            
            ins_mask = instance_label>=0
            if ins_mask.int().sum()==0:
                print('#######################??')
                ins_mask[:10] = True 
            cls_input = voxel_features[ins_mask].clone()

        if self.multi_scale and self.training:
            query_bxyz = batch_dict['strat_up4_bxyz'][ins_mask]
            unet_features = [batch_dict['strat_up3_feat'], batch_dict['strat_up2_feat'], batch_dict['strat_up1_feat']]

            for ri, ref_bxyz in enumerate([batch_dict['strat_up3_bxyz'], batch_dict['strat_up2_bxyz'], \
                    batch_dict['strat_up1_bxyz'],]):
                if ref_bxyz.shape[0]>0 and query_bxyz.shape[0]>0:
                    e_ref, e_query = self.graph(ref_bxyz, query_bxyz)
                    ref_feat = unet_features[ri]
                    query_feat = ref_feat[e_ref]
                else:
                    query_feat = torch.zeros(cls_input.shape[0], unet_features[ri].shape[1]).cuda()
                cls_input = torch.cat([cls_input, query_feat], dim=1)

        if self.append_points and self.training:
            query_bxyz = batch_dict['strat_up4_bxyz'][ins_mask]
            pt_ins_mask = pt_instance_label>=0
            if pt_ins_mask.int().sum()==0:
                pt_ins_mask[:10] = True 
            ref_bxyz = batch_dict['point_bxyz'][pt_ins_mask]

            if ref_bxyz.shape[0]>0 and query_bxyz.shape[0]>0:
                e_ref, e_query = self.ins_graph(ref_bxyz, query_bxyz)
            
                appended_points = ref_bxyz[e_ref][:, 1:].reshape(query_bxyz.shape[0], 32, 3)
                dis = torch.sqrt((torch.square(appended_points - query_bxyz[:, 1:].unsqueeze(1))).sum(2))
                dis_mask = (dis<0.5).int() 
                appended_points = appended_points - query_bxyz[:, 1:].unsqueeze(1)
                appended_points = appended_points * dis_mask.unsqueeze(-1)
                appended_points = appended_points.reshape(query_bxyz.shape[0], 32*3)
            else:
                appended_points = torch.zeros(cls_input.shape[0], 32*3).cuda()

            cls_input = torch.cat([cls_input, appended_points], dim=1)
        
        if self.training and self.ins_cls:
            aa = False
            if cls_input.shape[0]==1:
                cls_input = cls_input.repeat(2, 1)
                aa = True 
            point_embeddings = self.embed_layers(cls_input.float())
            if aa:
                point_embeddings = point_embeddings[0].unsqueeze(0)

            point_label = point_pred_scores.max(-1)[1]                
            masked_embeddings = point_embeddings 
            instance_label_few = instance_label[ins_mask.cuda()]
             
            if self.global_fea_op=='mean':
                ins_embeddings = scatter(masked_embeddings, instance_label_few,dim=0, reduce='mean')
            else:
                ins_embeddings = scatter(masked_embeddings, instance_label_few,dim=0, reduce='max')
            if ins_embeddings.shape[0]==1:
                ins_embeddings = ins_embeddings.repeat(2, 1)
            ins_cls_pred = self.ins_cls_layers(ins_embeddings)
            ins_cls_pred = ins_cls_pred[torch.unique(instance_label_few)]
            
            gt_ins_cls_label = scatter(segmentation_label[ins_mask].float(), instance_label_few, dim=0, reduce='mean')
            gt_ins_cls_label = torch.round(gt_ins_cls_label)
            gt_ins_cls_label = gt_ins_cls_label[torch.unique(instance_label_few)]

            if self.ins_cls_mapping is not None:
                gt_ins_cls_label = self.ins_cls_mapping[gt_ins_cls_label.long()]

            cls_mask = (gt_ins_cls_label>0) & (gt_ins_cls_label<=self.n_ins_cls)
            if flag:
                print('#############')
                cls_mask[0] = True 
                gt_ins_cls_label[0] = 1 
            
            if self.n_ins_cls==19:
                gt_ins_cls_label = gt_ins_cls_label[cls_mask]
            else:
                gt_ins_cls_label = gt_ins_cls_label[cls_mask] - 1 
            
            ins_cls_pred = ins_cls_pred[cls_mask]
            
            batch_dict['ins_cls_pred'] = ins_cls_pred
            batch_dict['ins_cls_label'] = gt_ins_cls_label
            
            if 0:
                    from ...utils.vis_utils import write_ply_color
                    import os
                    out_path = 'tmp_refine_vis/strat_inscls1'
                    if not os.path.exists(out_path):
                        os.makedirs(out_path)

                    if 0:
                    #for bi in range(batch_dict['batch_size']):
                        bs_mask = batch_dict[self.batch_key][:, 0]==bi
                        voxels = batch_dict[self.batch_key][bs_mask]
                        plot_ins_label = batch_dict['strat_instance_label'][bs_mask]%22
                        plot_ins_label_back = batch_dict['strat_instance_label_back'][bs_mask]%22
                        seg_label = batch_dict['strat_segmentation_label'][bs_mask]
                        name = batch_dict['frame_id'][bi]
                        #name = batch_dict['frame_id'][bi]

                        write_ply_color(voxels[:, 1:].cpu().numpy(), plot_ins_label.cpu().numpy(), os.path.join(out_path,'%s_ins.ply'%(name)))
                        #write_ply_color(voxels[:,1:].cpu().numpy(), plot_ins_label_back.cpu().numpy(), os.path.join(out_path,'%s_back.ply'%(name)))
                        write_ply_color(voxels[:,1:].cpu().numpy(), seg_label.cpu().numpy(), os.path.join(out_path,'%s_seg.ply'%(name)))
                        """
                        bs_mask = batch_dict['point_bxyz'][:, 0] ==bi
                        voxels = batch_dict['point_bxyz'][bs_mask]
                        plot_ins_label = batch_dict['instance_label'][bs_mask]%22
                        plot_ins_label_back = batch_dict['instance_label_back'][bs_mask]%22

                        write_ply_color(voxels[:, 1:].cpu().numpy(), plot_ins_label.cpu().numpy(), os.path.join(out_path,'%s_pt_ins.ply'%(name)))
                        write_ply_color(voxels[:,1:].cpu().numpy(), plot_ins_label_back.cpu().numpy(), os.path.join(out_path,'%s_pt_back.ply'%(name)))
                        """

                    voxels =  batch_dict[self.batch_key]
                    #name = batch_dict['frame_id'][0][0]
                    name = batch_dict['frame_id']
                    #if 0:
                    for ii, lb in enumerate(torch.unique(instance_label_few)[cls_mask]):
                        point_i = voxels[instance_label==lb][:, 1:]
                        seg_i = (torch.ones(point_i.shape[0])*gt_ins_cls_label[ii].cpu()).long()
                        write_ply_color(point_i.cpu().numpy(), seg_i.cpu().numpy(), os.path.join(out_path, '%s_%d_%d.ply'%(name, ii, \
                                    gt_ins_cls_label[ii])))

                    print('saved visualization')


        if self.target_assigner_cfg is not None:
            self.assign_targets(self.target_assigner_cfg, batch_dict)
        

        ret_dict = {
            'pred_seg_cls_logits': point_pred_logits,
        }

        point_pred_scores = torch.sigmoid(point_pred_logits)
        ret_dict['pred_seg_cls_confidences'], ret_dict['pred_seg_cls_labels'] = point_pred_scores.max(dim=-1)
        
        batch_dict.update(ret_dict)
        
    
        if self.gt_seg_cls_label_key in batch_dict:
            ret_dict[self.gt_seg_cls_label_key] = batch_dict[self.gt_seg_cls_label_key]

        if self.ins_cls and self.training:
            ret_dict['ins_cls_label'] = batch_dict['ins_cls_label']
            ret_dict['ins_cls_pred'] = batch_dict['ins_cls_pred']
        elif self.training:
            ret_dict['ins_cls_pred'] = torch.zeros(1,7).cuda()
            ret_dict['ins_cls_label'] = torch.zeros(1).cuda()

        ret_dict['batch_idx'] = batch_dict[self.batch_key][:, 0].round().long()
        ret_dict['point_bxyz'] = batch_dict[self.batch_key] 

        if self.assign_to_point and (not self.training):
            # assign pred_seg_cls_labels to points
            ref_bxyz = batch_dict[self.batch_key]
            ref_labels = ret_dict['pred_seg_cls_labels']
            query_bxyz = batch_dict['point_bxyz']
            e_ref, e_query = self.graph(ref_bxyz, query_bxyz)
            new_ret_dict = {}
            for key in ret_dict.keys():
                new_ret_dict[key] = scatter(ret_dict[key][e_ref], e_query, dim=0,
                                            dim_size=query_bxyz.shape[0], reduce='max')
            new_ret_dict['point_bxyz'] = batch_dict['point_bxyz']
            ret_dict = new_ret_dict
        
        ret_dict['batch_size'] = batch_dict['batch_size']
        self.forward_ret_dict = ret_dict

        return batch_dict


