import torch
import torch.nn as nn
import torch.nn.functional as F
from model.common import (MLP, Custom1x1Subm3d, ResidualBlock, UBlock)

import functools
from collections import OrderedDict
import spconv.pytorch as spconv
from spconv.pytorch.modules import SparseModule

# this file only provides the 2 modules used in VQVAE
__all__ = ['Encoder', 'Decoder',]
    
class Encoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, spconv_channels, num_blocks):
        super(Encoder, self).__init__()
        
        block = ResidualBlock
        norm_fn = functools.partial(nn.BatchNorm1d, eps=1e-4, momentum=0.1)

        self.input_conv = spconv.SparseSequential(
            spconv.SubMConv3d(
                input_dim, spconv_channels, kernel_size=3, padding=1, bias=False, indice_key='subm1'))
        block_channels = [spconv_channels * (i + 1) for i in range(num_blocks)]
        self.unet = UBlock(block_channels, norm_fn, 2, block, indice_key_id=1)
        self.output_layer = spconv.SparseSequential(norm_fn(spconv_channels), nn.ReLU())


    def forward(self, voxel_feats, voxel_coords, spatial_shape):

        # spconv encode
        batch_size = len(voxel_coords)
        voxel_feats_total = []
        voxel_coords_total = []
        voxel_batch_id_total = []
        for i in range(batch_size):
            batch_col = torch.zeros((voxel_coords[i].shape[0], 1), device=voxel_coords[i].device, dtype=torch.int32)
            voxel_batch_id_total.append(batch_col)
            voxel_coords_total.append(voxel_coords[i].int())
            voxel_coords[i] = torch.cat([batch_col, voxel_coords[i].int()], dim=1)
            input = spconv.SparseConvTensor(voxel_feats[i], voxel_coords[i], spatial_shape, 1)
            output = self.input_conv(input)
            output = self.unet(output)
            voxel_feats_total.append(output)

        return voxel_feats_total, voxel_coords_total, voxel_batch_id_total

class Decoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super(Decoder, self).__init__()
        self.num_layers = num_layers
        self.start_layer = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        self.layers = nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim),
                nn.ReLU(),
                nn.Linear(hidden_dim, hidden_dim)
            ) for _ in range(num_layers - 2)
        ])
        self.final_layer = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        x = self.start_layer(x)
        for layer in self.layers:
            x = layer(x)
        return self.final_layer(x)