import torch
import torch.nn as nn
import torch.nn.functional as F
import functools

from .pv_module import SharedMLP, PVConv 

def create_pointnet_components(
        blocks, in_channels, with_se=False, normalize=True, eps=0,
        width_multiplier=1, voxel_resolution_multiplier=1, scale_pvcnn=False, device='cuda'):
    r, vr = width_multiplier, voxel_resolution_multiplier
    layers, concat_channels = [], 0
    for out_channels, num_blocks, voxel_resolution in blocks:
        out_channels = int(r * out_channels)
        if voxel_resolution is None:
            block = functools.partial(SharedMLP, device=device)
        else:
            block = functools.partial(
                PVConv, kernel_size=3, resolution=int(vr * voxel_resolution),
                with_se=with_se, normalize=normalize, eps=eps, scale_pvcnn=scale_pvcnn, device=device)
        for _ in range(num_blocks):
            layers.append(block(in_channels, out_channels))
            in_channels = out_channels
            concat_channels += out_channels
    return layers, in_channels, concat_channels

class PCMerger(nn.Module):
# merge surface sampled PC and rendering backprojected PC (w/ 2D features):
    def __init__(self, in_channels=204, device="cuda"):
        super(PCMerger, self).__init__()
        self.mlp_normal = SharedMLP(3, [128, 128], device=device)
        self.mlp_rgb = SharedMLP(3, [128, 128], device=device)
        self.mlp_sam = SharedMLP(204 - 6, [128, 128], device=device)

    def forward(self, feat, mv_feat, pc2pc_idx):
        mv_feat_normal = self.mlp_normal(mv_feat[:, :3, :])
        mv_feat_rgb = self.mlp_rgb(mv_feat[:, 3:6, :])
        mv_feat_sam = self.mlp_sam(mv_feat[:, 6:, :])

        mv_feat_normal = mv_feat_normal.permute(0, 2, 1)
        mv_feat_rgb = mv_feat_rgb.permute(0, 2, 1)
        mv_feat_sam = mv_feat_sam.permute(0, 2, 1)
        feat = feat.permute(0, 2, 1)
        
        for i in range(mv_feat.shape[0]):
            mask = (pc2pc_idx[i] != -1).reshape(-1)
            idx = pc2pc_idx[i][mask].reshape(-1)
            feat[i][mask] += mv_feat_normal[i][idx] + mv_feat_rgb[i][idx] + mv_feat_sam[i][idx]
            
        return feat.permute(0, 2, 1)


class PVCNNEncoder(nn.Module):
    def __init__(self, pvcnn_feat_dim, device='cuda', in_channels=3, use_2d_feat=False):
        super(PVCNNEncoder, self).__init__()
        self.device = device
        self.blocks = ((pvcnn_feat_dim, 1, 32), (128, 2, 16), (256, 1, 8))
        self.use_2d_feat=use_2d_feat
        if in_channels == 6:
            self.append_channel = 2
        elif in_channels == 3:
            self.append_channel = 1
        else:
            raise NotImplementedError
        layers, channels_point, concat_channels_point = create_pointnet_components(
            blocks=self.blocks, in_channels=in_channels + self.append_channel, with_se=False, normalize=False,
            width_multiplier=1, voxel_resolution_multiplier=1, scale_pvcnn=True,
            device=device
        )
        self.encoder = nn.ModuleList(layers)#.to(self.device)
        if self.use_2d_feat:
            self.merger = PCMerger()

        

    def forward(self, input_pc, mv_feat=None, pc2pc_idx=None):
        features = input_pc.permute(0, 2, 1) * 2  # make point cloud [-1, 1]
        coords = features[:, :3, :]
        out_features_list = []
        voxel_feature_list = []
        zero_padding = torch.zeros(features.shape[0], self.append_channel, features.shape[-1], device=features.device, dtype=features.dtype)
        features = torch.cat([features, zero_padding], dim=1)##################

        for i in range(len(self.encoder)):
            features, _, voxel_feature = self.encoder[i]((features, coords))
            if i == 0 and mv_feat is not None:
               features = self.merger(features, mv_feat.permute(0, 2, 1), pc2pc_idx)
            out_features_list.append(features)
            voxel_feature_list.append(voxel_feature)
        return voxel_feature_list, out_features_list