'''
Author       : ZHP
Date         : 2021-05-20 12:44:27
LastEditors  : ZHP
LastEditTime : 2021-05-29 00:07:23
FilePath     : \Projects\ptzhs\point_transformer\models\ShapePartModel.py
Description  : 项目说明
Copyright    : ZHP
'''
import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

import os
# import sys
# sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from data import Indoor3DSemSeg, PartNormalDataset
# from point_transformer.data import Indoor3DSemSeg, PartNormalDataset
from models.base import BaseClassification
from utils.timer import Timer

from point_transformer_ops.point_transformer_modules import PointTransformerBlock, TransitionDown, TransitionUp
from torchsummary import summary


class PointTransformerPartSegmentationContextOnlyEnd(nn.Module):
    '''
    Point Transformer Segmantic Segmentation
    Context module at end  /  no context
    '''
    # [3,32,64,128,256,512] part  output_dim=50
    # [6,32,64,128,256,512] sem ouput_dim=13
    def __init__(self, dim=[3,32,64,128,256,512], npoints=4096, output_dim=50, pos_mlp_hidden=64, attn_mlp_hidden=4, k=16, sampling_ratio=0.25, context=False):
        '''
        @params:
            npoints : N,初始采样点数
            sampling_ratio : 降采样比例，每次 /4
            context : if True, context module at end;if False, no context module,so originale point transformer
        '''
        super().__init__()
        self.context = context
        # Encoder
        self.Encoder = nn.ModuleList()
        for i in range(len(dim)-1):
            if i == 0:
                self.Encoder.append(nn.Linear(dim[i], dim[i+1], bias=False))
            else:
                self.Encoder.append(TransitionDown(dim[i], dim[i+1], k, sampling_ratio, fast=True))
            self.Encoder.append(PointTransformerBlock(dim[i+1], k))

        # Decoder
        self.Decoder = nn.ModuleList()
        for i in range(5,0,-1):
            if i == 5:
                self.Decoder.append(nn.Linear(dim[i], dim[i], bias=False))
            else:
                self.Decoder.append(TransitionUp(dim[i+1], dim[i]))
            self.Decoder.append(PointTransformerBlock(dim[i], k))
        # context prior module
        if context:
            # add one context module at the last transformer
            self.context_end = Context(dim[1])
            out_dim = dim[1] * 3
        else:
            out_dim = dim[1]
        # fc
        self.fc_layer = nn.Sequential(
            nn.Conv1d(out_dim, out_dim, kernel_size=1, bias=False),
            nn.BatchNorm1d(out_dim),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Conv1d(out_dim, output_dim, kernel_size=1),
        )


    def break_up_pc(self, pc):
        xyz = pc[..., 0:3].contiguous()  # coordinate
        features = pc[..., 3:].transpose(1, 2).contiguous() if pc.size(-1) > 3 else None

        # [B, N, 3]  [B, 3, N]
        return xyz, features


    def forward(self, pointcloud):
        # timer = Timer("forward")
        # timer.tic()
        xyz, features = self.break_up_pc(pointcloud)
        features = features.transpose(1,2).contiguous()

        l_xyz, l_features = [xyz], [features]
        # Encode stage
        for i in range(int(len(self.Encoder)/2)):
            if i == 0:
                li_features = self.Encoder[2*i](l_features[i])
                li_xyz = l_xyz[i]
            else:
                li_features, li_xyz = self.Encoder[2*i](l_features[i], l_xyz[i])
            li_features = self.Encoder[2*i+1](li_features, li_xyz)
            
            l_features.append(li_features)
            l_xyz.append(li_xyz)
            del li_features, li_xyz            
        D_n = int(len(self.Decoder)/2)
        
        # Decode stage
        for i in range(D_n):
            if i == 0:
                l_features[D_n-i] = self.Decoder[2*i](l_features[D_n-i])
                l_features[D_n-i] = self.Decoder[2*i+1](l_features[D_n-i], l_xyz[D_n-i])
            else:
                l_features[D_n-i], l_xyz[D_n-i] = self.Decoder[2*i](l_features[D_n-i+1], l_xyz[D_n-i+1], l_features[D_n-i], l_xyz[D_n-i])
                l_features[D_n-i] = self.Decoder[2*i+1](l_features[D_n-i], l_xyz[D_n-i])
        # 只保留 l_features[1]  -- f_out
        f_out = l_features[1]  # [B, 2500, 32]
        del l_features, l_xyz
        if self.context:
            context_map, f_out = self.context_end(f_out.permute(0, 2, 1))
            f_out = f_out.permute(0, 2, 1)
        else:
            context_map = None
        # del l_features[0], l_features[1:], l_xyz
        out = self.fc_layer(f_out.transpose(1,2).contiguous())
        # timer.toc()
        return out, context_map  # (B, 50, N)        

class PointTransformerPartSegmentation(nn.Module):
    '''
    Point Transformer Segmantic Segmentation with two context module
    behind the first tranformer and the last tranformer
    '''

    # [3,32,64,128,256,512] part  output_dim=50
    # [6,32,64,128,256,512] sem ouput_dim=13
    def __init__(self, dim=[3,32,64,128,256,512], npoints=4096, output_dim=50, pos_mlp_hidden=64, attn_mlp_hidden=4, k=16, sampling_ratio=0.25, context=False):
        '''
        @params:
            npoints : N,初始采样点数
            sampling_ratio : 降采样比例，每次 /4
            context : if true,represent has context module at the first tranformer and the last tranformer
        '''
        super().__init__()
        self.context = context
        # Encoder
        self.Encoder = nn.ModuleList()
        for i in range(len(dim)-1):
            if i == 0:
                self.Encoder.append(nn.Linear(dim[i], dim[i+1], bias=False))
            else:
                self.Encoder.append(TransitionDown(dim[i], dim[i+1], k, sampling_ratio, fast=True))
            self.Encoder.append(PointTransformerBlock(dim[i+1], k))

        # Decoder
        self.Decoder = nn.ModuleList()
        for i in range(5,0,-1):
            if i == 5:
                self.Decoder.append(nn.Linear(dim[i], dim[i], bias=False))
            else:
                self.Decoder.append(TransitionUp(dim[i+1], dim[i]))

            self.Decoder.append(PointTransformerBlock(dim[i], k))
        # context prior module
        if context:
            # add two context module behind the first tranformer and the last transformer
            self.context_start = Context(dim[1])
            
            # 1x1 convolution to reduce channels
            self.reduce_conv = nn.Conv1d(dim[1] * 3, dim[1], kernel_size=1, bias=False)
            self.context_end = Context(dim[1])
            out_dim = dim[1] * 3
        else:
            out_dim = dim[1]
        # fc
        self.fc_layer = nn.Sequential(
            nn.Conv1d(out_dim, out_dim, kernel_size=1, bias=False),
            nn.BatchNorm1d(out_dim),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Conv1d(out_dim, output_dim, kernel_size=1),
        )


    def break_up_pc(self, pc):
        xyz = pc[..., 0:3].contiguous()  # coordinate
        features = pc[..., 3:].transpose(1, 2).contiguous() if pc.size(-1) > 3 else None

        # [B, N, 3]  [B, 3, N]
        return xyz, features


    def forward(self, pointcloud):
        # timer = Timer("forward")
        # timer.tic()
        xyz, features = self.break_up_pc(pointcloud)
        features = features.transpose(1,2).contiguous()

        l_xyz, l_features = [xyz], [features]
        # Encode stage
        for i in range(int(len(self.Encoder)/2)):
            if i == 0:
                li_features = self.Encoder[2*i](l_features[i])
                li_xyz = l_xyz[i]
            else:
                li_features, li_xyz = self.Encoder[2*i](l_features[i], l_xyz[i])
            li_features = self.Encoder[2*i+1](li_features, li_xyz)
            if i == 0 and self.context:
                context_map_start, li_features = self.context_start(li_features.permute(0, 2, 1).contiguous())
                # print("2-->{} || {}".format(li_features.shape, context_map_start.shape)) # 2-->torch.Size([8, 96, 4096]) || torch.Size([8, 4096, 4096])
                li_features = self.reduce_conv(li_features)
                # print("3-->{} || {}".format(li_features.shape, context_map_start.shape)) # 3-->torch.Size([8, 32, 4096]) || torch.Size([8, 4096, 4096])
                li_features = li_features.permute(0, 2, 1)
                # print("4-->{} || {}".format(li_features.shape, context_map_start.shape)) # 4-->torch.Size([8, 4096, 32]) || torch.Size([8, 4096, 4096])
            l_features.append(li_features)
            l_xyz.append(li_xyz)
            del li_features, li_xyz            
        D_n = int(len(self.Decoder)/2)
        
        # Decode stage
        for i in range(D_n):
            if i == 0:
                l_features[D_n-i] = self.Decoder[2*i](l_features[D_n-i])
                l_features[D_n-i] = self.Decoder[2*i+1](l_features[D_n-i], l_xyz[D_n-i])
            else:
                l_features[D_n-i], l_xyz[D_n-i] = self.Decoder[2*i](l_features[D_n-i+1], l_xyz[D_n-i+1], l_features[D_n-i], l_xyz[D_n-i])
                l_features[D_n-i] = self.Decoder[2*i+1](l_features[D_n-i], l_xyz[D_n-i])
        # 只保留 l_features[1]  -- f_out
        f_out = l_features[1]  # [B, 2500, 32]
        del l_features, l_xyz
        if self.context:
            context_map, f_out = self.context_end(f_out.permute(0, 2, 1))
            
            # composite 1
            # context_map = context_map + context_map_start
            # context_map = torch.sigmoid(context_map)
            
            # composite 2
            context_map = torch.mean(torch.cat([context_map.unsqueeze(0), context_map_start.unsqueeze(0)], dim=0), dim=0)
            f_out = f_out.permute(0, 2, 1)
        else:
            context_map = None
        # del l_features[0], l_features[1:], l_xyz
        out = self.fc_layer(f_out.transpose(1,2).contiguous())
        # timer.toc()
        return out, context_map  # (B, 50, N)


class PointTransformerPartSegmentation_v1(nn.Module):
    '''
    将attention map作为context map监督
    '''
    # [3,32,64,128,256,512] part  output_dim=50
    # [6,32,64,128,256,512] sem ouput_dim=13
    def __init__(self, dim=[3,32,64,128,256,512], npoints=4096, output_dim=50, pos_mlp_hidden=64, attn_mlp_hidden=4, k=16, sampling_ratio=0.25, context=False):
        '''
        @params:
            npoints : N,初始采样点数
            sampling_ratio : 降采样比例，每次 /4
            context : if true,represent has context module at the first tranformer and the last tranformer
        '''
        super().__init__()
        self.context = context
        # Encoder
        self.Encoder = nn.ModuleList()
        for i in range(len(dim)-1):
            if i == 0:
                self.Encoder.append(nn.Linear(dim[i], dim[i+1], bias=False))
            else:
                self.Encoder.append(TransitionDown(dim[i], dim[i+1], k, sampling_ratio, fast=True))
            self.Encoder.append(PointTransformerBlock(dim[i+1], k))

        # Decoder
        self.Decoder = nn.ModuleList()
        for i in range(5,0,-1):
            if i == 5:
                self.Decoder.append(nn.Linear(dim[i], dim[i], bias=False))
            else:
                self.Decoder.append(TransitionUp(dim[i+1], dim[i]))

            self.Decoder.append(PointTransformerBlock(dim[i], k))
        # context prior module
        if context:
            # add two context module behind the first tranformer and the last transformer
            self.context_start = Context(dim[1])
            
            # 1x1 convolution to reduce channels
            self.reduce_conv = nn.Conv1d(dim[1] * 3, dim[1], kernel_size=1, bias=False)
            self.context_end = Context(dim[1])
            out_dim = dim[1] * 3
        else:
            out_dim = dim[1]
        # fc
        self.fc_layer = nn.Sequential(
            nn.Conv1d(out_dim, out_dim, kernel_size=1, bias=False),
            nn.BatchNorm1d(out_dim),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Conv1d(out_dim, output_dim, kernel_size=1),
        )

    def break_up_pc(self, pc):
        xyz = pc[..., 0:3].contiguous()  # coordinate
        features = pc[..., 3:].transpose(1, 2).contiguous() if pc.size(-1) > 3 else None

        # [B, N, 3]  [B, 3, N]
        return xyz, features

    def forward(self, pointcloud):
        # timer = Timer("forward")
        # timer.tic()
        xyz, features = self.break_up_pc(pointcloud)
        features = features.transpose(1,2).contiguous()

        l_xyz, l_features = [xyz], [features]
        # Encode stage
        for i in range(int(len(self.Encoder)/2)):
            if i == 0:
                li_features = self.Encoder[2*i](l_features[i])
                li_xyz = l_xyz[i]
            else:
                li_features, li_xyz = self.Encoder[2*i](l_features[i], l_xyz[i])
            li_features = self.Encoder[2*i+1](li_features, li_xyz)
            if i == 0 and self.context:
                context_map_start, li_features = self.context_start(li_features.permute(0, 2, 1).contiguous())
                # print("2-->{} || {}".format(li_features.shape, context_map_start.shape)) # 2-->torch.Size([8, 96, 4096]) || torch.Size([8, 4096, 4096])
                li_features = self.reduce_conv(li_features)
                # print("3-->{} || {}".format(li_features.shape, context_map_start.shape)) # 3-->torch.Size([8, 32, 4096]) || torch.Size([8, 4096, 4096])
                li_features = li_features.permute(0, 2, 1)
                # print("4-->{} || {}".format(li_features.shape, context_map_start.shape)) # 4-->torch.Size([8, 4096, 32]) || torch.Size([8, 4096, 4096])
            l_features.append(li_features)
            l_xyz.append(li_xyz)
            del li_features, li_xyz            
        D_n = int(len(self.Decoder)/2)
        
        # Decode stage
        for i in range(D_n):
            if i == 0:
                l_features[D_n-i] = self.Decoder[2*i](l_features[D_n-i])
                l_features[D_n-i] = self.Decoder[2*i+1](l_features[D_n-i], l_xyz[D_n-i])
            else:
                l_features[D_n-i], l_xyz[D_n-i] = self.Decoder[2*i](l_features[D_n-i+1], l_xyz[D_n-i+1], l_features[D_n-i], l_xyz[D_n-i])
                l_features[D_n-i] = self.Decoder[2*i+1](l_features[D_n-i], l_xyz[D_n-i])
        # 只保留 l_features[1]  -- f_out
        f_out = l_features[1]  # [B, 2500, 32]
        del l_features, l_xyz
        if self.context:
            context_map, f_out = self.context_end(f_out.permute(0, 2, 1))
            context_map = context_map + context_map_start
            context_map = torch.sigmoid(context_map)
            f_out = f_out.permute(0, 2, 1)
        else:
            context_map = None
        # del l_features[0], l_features[1:], l_xyz
        out = self.fc_layer(f_out.transpose(1,2).contiguous())
        # timer.toc()
        return out, context_map  # (B, 50, N)  
    
    
class Context(nn.Module):
    '''
    上下文先验层
    '''
    def __init__(self, in_channel, dim=[8, 16]):
        super(Context, self).__init__()
        '''
        self.conv_1 = pt_utils.Conv1d(in_channel, in_channel // 8, bn=True)
        self.conv_2 = pt_utils.Conv1d(in_channel // 8,
                                      in_channel // 16,
                                      bn=True,
                                      activation=None)
        self.sigmoid = nn.Sigmoid()
        '''
        self.conv_1 = conv_with_bn(in_channel, in_channel // dim[0], bn=True)
        self.conv_2 = conv_with_bn(in_channel // dim[0], in_channel // dim[1], bn=True, activation=None)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # print("x.sahpe ：", x.shape)
        out = self.conv_1(x)            # [B, in_channel // 8, N]
        out_1 = self.conv_2(out)   # [B, in_channel // 16, N]
        context = torch.bmm(out_1.permute(0, 2, 1), out_1)  # [B, N, N]
        context = self.sigmoid(context)
        p = torch.bmm(context, x.permute(0, 2, 1))  # P  <==>[B, N, in_channel]
        ones = torch.ones_like(context)  
        p_1 = torch.bmm(ones - context, x.permute(0, 2, 1)) # [B, N, in_channel] # 1 - P
        # cat  x  x*f(x)  (1-f(x))*x
        out = torch.cat((x, p.permute(0, 2, 1), p_1.permute(0, 2, 1)), dim=1) # [B, in_channel * 3, N]
        return context, out

def conv_with_bn(in_channel, out_channel, bn=True, activation="ReLU", kernel_size=1):
    '''
    @description: construct 1x1 convolution
    @param {
        in_channel : input channels
        out_channel : output channels
        bn : Batch Normalization
    }
    @return {*}
    '''    
    model = nn.Sequential()
    model.add_module("conv", nn.Conv1d(in_channel, out_channel, kernel_size))
    if bn:
        model.add_module("BN", nn.BatchNorm1d(out_channel))
    if activation is not None:
        model.add_module("activation", getattr(nn, activation)(inplace=True))
    return model

def make_one_hot(y, num_classes):
    """ 1-hot encodes a tensor """
    new_y = torch.eye(num_classes)[y.cpu().data.numpy(),]
    if (y.is_cuda):
        return new_y.cuda()
    return new_y

class get_part_loss(torch.nn.Module):
    def __init__(self, mat_diff_loss_scale=0.001):
        super(get_part_loss, self).__init__()
        self.mat_diff_loss_scale = mat_diff_loss_scale

    # part seg 50类
    def forward(self, pred, target, context, context_gt, num_classes=50):
        '''
        pred :[B, classes, N]
        target : [B, N]
        context : [B, N, N]
        context_gt : [B, N]  (target)
        '''
        context_gt = make_one_hot(context_gt, num_classes)  # [B, N, num_classes]
        # get LL^T  --- Ideal Affinity Map\
        context_gt = torch.bmm(context_gt,context_gt.permute(0, 2, 1)) # [B, N, N]
        
        
        loss = F.cross_entropy(pred+1e-8, target)
        # print('loss_.shape : ', loss.shape, "loss", torch.mean(loss))
        # context prior loss
        loss_u=F.binary_cross_entropy_with_logits(context.float(),context_gt.float())  # 二分类交叉熵
        # print('loss_u.shape : ', loss_u.shape, "loss_u", torch.mean(loss_u))
        
        # precision
        # loss_g_p = torch.log((context.mul(context_gt)).sum(2)+1e-8) - torch.log(context.sum(2)+1e-8)
        loss_g_p = torch.log((context.mul(context_gt)).sum(2)+1e-8) - torch.log(context.sum(2)+1e-8)

        # recall
        loss_g_r = torch.log((context.mul(context_gt)).sum(2)+1e-8) - torch.log(context_gt.sum(2)+1e-8)
        
        # specificity  
        # loss_g_s = torch.log(((torch.ones_like(context) - context).mul((torch.ones_like(context_gt) - context_gt))).sum(2)+1e-8) - torch.log((1 - context_gt).sum(2)+1e-8)
        loss_g_s = torch.log(((torch.ones_like(context) - context).mul((torch.ones_like(context_gt) - context_gt))).sum(2)+1e-8) - torch.log((1 - context_gt).sum(2)+1e-8)
        loss_g = - torch.mean(loss_g_p + loss_g_r + loss_g_s)
       
        '''
        print('loss_g_p.shape : ', loss_g_p.shape, "loss_g_p", torch.mean(loss_g_p))
        print('loss_g_r.shape : ', loss_g_r.shape, 'loss_g_r',torch.mean(loss_g_r))
        print('loss_g_s.shape : ', loss_g_s.shape, 'loss_g_s',torch.mean(loss_g_s))
        print('loss_g',loss_g)
        '''
        total_loss = loss + loss_u + loss_g   #+loss_fc
        return total_loss

class get_loss(torch.nn.Module):
    def __init__(self, mat_diff_loss_scale=0.001):
        super(get_loss, self).__init__()
        self.mat_diff_loss_scale = mat_diff_loss_scale

    # part seg 50类
    def forward(self, pred, target, context, context_gt, num_classes=50):
        '''
        pred :[B, classes, N]
        target : [B, N]
        context : [B, N, N]
        context_gt : [B, N]  (target)
        '''
        #print(context_gt.shape)
        print(pred.shape,target.shape,context.shape,context_gt.shape)
        context_gt = make_one_hot(context_gt, num_classes)  # [B, N, num_classes]
        # get LL^T  --- Ideal Affinity Map
        context_gt = torch.bmm(context_gt,context_gt.permute(0, 2, 1)) # [B, N, N]
        loss = F.cross_entropy(pred+1e-8, target)
        print('------')
        #print(context,context_gt)
        #print(pred.shape,target.shape)
        print(context_gt.shape)
        print('cr',loss)
        #loss=F.nll_loss(pred,target)
        
        # context prior loss
        loss_u=F.binary_cross_entropy_with_logits(context,context_gt)  # 二分类交叉熵
        print('single cr',loss_u)
        # precision
        loss_g_p = torch.log((context.mul(context_gt)).sum(2)+1e-8) - torch.log(context.sum(2)+1e-8)
        # recall
        loss_g_r = torch.log((context.mul(context_gt)).sum(2)+1e-8) - torch.log(context_gt.sum(2)+1e-8)
        
        # specificity  
        loss_g_s = torch.log(((torch.ones_like(context) - context).mul((torch.ones_like(context_gt) - context_gt))).sum(2)+1e-8) - torch.log((1 - context_gt).sum(2)+1e-8)
        loss_g = - torch.mean(loss_g_p + loss_g_r + loss_g_s)
        print('loss_g_p.shape : ', loss_g_p.shape, "log_g_p", torch.mean(log_g_p))
        print('loss_g_r.shape : ', loss_g_r.shape, 'loss_g_r',torch.mean(loss_g_r))
        print('loss_g_s.shape : ', loss_g_s.shape, 'loss_g_s',torch.mean(loss_g_s))
        print('loss_g',loss_g)
        #loss_f=F.nll_loss(pred,target)
        #one_hot=make_one_hot(target,pred.size()[1])
        #inv_probs = 1 - pred.exp()
        #focal_weights = (inv_probs * one_hot).sum(dim=1) ** 2
        #loss_fc=loss_f*focal_weights
        #loss_fc=loss_fc.mean()
        total_loss = loss + loss_u + loss_g #+loss_fc
        print('total',total_loss)
        print('----------')
        return total_loss

if __name__ == "__main__":
    model = PointTransformerPartSegmentation(context=True)
    # print(model)
    model.cuda()
    pts = torch.randn((8, 2500, 6))
    outs, context_map = model(pts.cuda())
    print("outs.shape : {} ||  context_map.shape : {}".format(outs.shape, context_map.shape))