# modified from https://github.com/autonomousvision/convolutional_occupancy_networks/blob/master/src/encoder/pointnet.py
from dataclasses import dataclass
import torch
import torch.nn as nn
from torch_scatter import scatter_mean, scatter_max

from TriplaneGaussian.utils.base import BaseModule
from TriplaneGaussian.models.networks import ResnetBlockFC
from TriplaneGaussian.utils.ops import scale_tensor

class LocalPoolPointnet(BaseModule):
    ''' PointNet-based encoder network with ResNet blocks for each point.
        Number of input points are fixed.
    
    Args:
        c_dim (int): dimension of latent code c
        dim (int): input points dimension
        hidden_dim (int): hidden dimension of the network
        scatter_type (str): feature aggregation when doing local pooling
        plane_resolution (int): defined resolution for plane feature
        padding (float): conventional padding paramter of ONet for unit cube, so [-0.5, 0.5] -> [-0.55, 0.55]
        n_blocks (int): number of blocks ResNetBlockFC layers
    '''

    @dataclass
    class Config(BaseModule.Config):
        input_channels: int = 3
        c_dim: int = 128
        hidden_dim: int = 128
        scatter_type: str = "max"
        plane_size: int = 32
        n_blocks: int = 5
        radius: float = 1.

    cfg: Config

    def configure(self) -> None:
        super().configure()
        self.fc_pos = nn.Linear(self.cfg.input_channels, 2 * self.cfg.hidden_dim)
        self.blocks = nn.ModuleList([
            ResnetBlockFC(2 * self.cfg.hidden_dim, self.cfg.hidden_dim) for i in range(self.cfg.n_blocks)
        ])
        self.fc_c = nn.Linear(self.cfg.hidden_dim, self.cfg.c_dim)

        self.actvn = nn.ReLU()

        if self.cfg.scatter_type == 'max':
            self.scatter = scatter_max
        elif self.cfg.scatter_type == 'mean':
            self.scatter = scatter_mean
        else:
            raise ValueError('incorrect scatter type')


    def generate_plane_features(self, index, c):
        # acquire indices of features in plane
        # xy = normalize_coordinate(p.clone(), plane=plane, padding=self.padding) # normalize to the range of (0, 1)
        # index = self.coordinate2index(x, self.cfg.plane_size)

        # scatter plane features from points
        fea_plane = c.new_zeros(index.shape[0], self.cfg.c_dim, self.cfg.plane_size ** 2)
        c = c.permute(0, 2, 1) # B x 512 x T
        fea_plane = scatter_mean(c, index, out=fea_plane) # B x 512 x reso^2
        fea_plane = fea_plane.reshape(index.shape[0], self.cfg.c_dim, self.cfg.plane_size, self.cfg.plane_size) # sparce matrix (B x 512 x reso x reso)

        return fea_plane

    def pool_local(self, xy, index, c):
        bs, fea_dim = c.shape[0], c.shape[2]
        keys = xy.keys()

        c_out = 0
        for key in keys:
            # scatter plane features from points
            fea = self.scatter(c.permute(0, 2, 1), index[key], dim_size=self.cfg.plane_size ** 2)
            if self.scatter == scatter_max:
                fea = fea[0]
            # gather feature back to points
            fea = fea.gather(dim=2, index=index[key].expand(-1, fea_dim, -1))
            c_out += fea
        return c_out.permute(0, 2, 1)

    def coordinate2index(self, x):
        x = (x * self.cfg.plane_size).long()
        index = x[..., 0] + self.cfg.plane_size * x[..., 1]
        assert index.max() < self.cfg.plane_size ** 2
        return index[:, None, :]

    def forward(self, p):
        batch_size, T, D = p.shape

        # acquire the index for each point
        coord = {}
        index = {}

        position = torch.clamp(p[..., :3], -self.cfg.radius + 1e-6, self.cfg.radius - 1e-6)
        position_norm = scale_tensor(position, (-self.cfg.radius, self.cfg.radius), (0, 1))
        coord["xy"] = position_norm[..., [0, 1]]
        coord["xz"] = position_norm[..., [0, 2]]
        coord["yz"] = position_norm[..., [1, 2]]
        index["xy"] = self.coordinate2index(coord["xy"])
        index["xz"] = self.coordinate2index(coord["xz"])
        index["yz"] = self.coordinate2index(coord["yz"])
        
        net = self.fc_pos(p)

        net = self.blocks[0](net)
        for block in self.blocks[1:]:
            pooled = self.pool_local(coord, index, net)
            net = torch.cat([net, pooled], dim=2)
            net = block(net)

        c = self.fc_c(net)

        features = torch.stack([
            self.generate_plane_features(index["xy"], c),
            self.generate_plane_features(index["xz"], c),
            self.generate_plane_features(index["yz"], c)
        ], dim=1)
        
        return features