import torch.nn as nn

from . import functional as F
from .voxelization import Voxelization
from .shared_mlp import SharedMLP
import torch

__all__ = ['PVConv']


class PVConv(nn.Module):
    def __init__(
            self, in_channels, out_channels, kernel_size, resolution, with_se=False, normalize=True, eps=0, scale_pvcnn=False,
            device='cuda'):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.resolution = resolution
        self.voxelization = Voxelization(resolution, normalize=normalize, eps=eps, scale_pvcnn=scale_pvcnn)
        voxel_layers = [
            nn.Conv3d(in_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2, device=device),
            nn.InstanceNorm3d(out_channels, eps=1e-4, device=device),
            nn.LeakyReLU(0.1, True),
            nn.Conv3d(out_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2, device=device),
            nn.InstanceNorm3d(out_channels, eps=1e-4, device=device),
            nn.LeakyReLU(0.1, True),
        ]
        self.voxel_layers = nn.Sequential(*voxel_layers)
        self.point_features = SharedMLP(in_channels, out_channels, device=device)

    def forward(self, inputs):
        features, coords = inputs
        voxel_features, voxel_coords = self.voxelization(features, coords)
        voxel_features = self.voxel_layers(voxel_features)
        devoxel_features = F.trilinear_devoxelize(voxel_features, voxel_coords, self.resolution, self.training)
        fused_features = devoxel_features + self.point_features(features)
        return fused_features, coords, voxel_features