'''
Decoder to obtain a coarse and fine structures for data representation
'''

import torch
from torch import nn

class PointCloudGraphDecoder(nn.Module):
    def __init__(self, n_z, n_node, n_edge_z=None, n_edge_x=1, grid_size=10, radius=0.25):
        super(PointCloudGraphDecoder, self).__init__()
        '''
        Generate point cloud data from a latent structure that considers a coarse schema of nodes, and potentially, the linkage information between nodes

        n_z   : latent context variable, could be encoded from another encoder or from a prior distribution
        n_node: number of nodes to construct the coarse representation
                some design consideration:
                1. we may assign a dimension partition of z to each node, in this case, n_z % n_node == 0. node is then somehow independent from each other
                2. we may use z as a global context variable to determine each node attribute (position, e.g.), 
                    this might be needed if we want to have other node properties that needs more global information, like garment cuff
                    actually, i think for a coarse representation, the nodes should be correlated, but at the same time it might be independent to some part of latent encoding as well
                    imagine the sleeve length that is different


        n_edge_z : encoding to get adjacency matrix, some design considerations:
                1. we might not need this, then there is no linkage encoded the decoder behaves like PCNDecoder
                2. use this encodings to construct an adjacency matrix, and use it to
                    - generate finer point clouds by concatenating rows with corresponding coarse nodes fine_decode(A_i, x_coarse_i) = a fraction of fine point clouds
                        in this case, it is like PCNDecoder that coarse point are on the surfaces, and local fine points are generated by wraping a local surface
                    - consider the distance between coarse nodes and input point clouds by accounting for edge, e.g., if a point is managed to project to the line segment between two points (if there is an edge)
                        then use that as the distance. this can allow us to build skeleton inside the clouds

        n_edge_x : feature dimension for each edge, might be useful for characterising relations
        '''
        
        self.n_node = n_node
        self.n_z = n_z
        self.n_edge_z = n_edge_z
        self.n_edge_x = n_edge_x

        self.grid_size = grid_size
        self.comp_radius = radius

        self.n_output_points = self.grid_size**2*n_node

        #fully-connected nets for node feature prediction
        self.fc_node_pos = nn.Sequential(
            nn.Linear(n_z, n_node),
            nn.LeakyReLU(negative_slope=0.2),
            nn.Linear(n_node, n_node),
            nn.LeakyReLU(negative_slope=0.2),
            nn.Linear(n_node, n_node*3)         #feature dim n_node*3 position
        )

        self.fc_node_pos_zcomp = nn.Sequential(
            nn.Linear(n_z // n_node, 32),
            nn.LeakyReLU(negative_slope=0.2),
            nn.Linear(32, 32),
            nn.LeakyReLU(negative_slope=0.2),
            nn.Linear(32, 3)                    #feature dim 3 position for each node
        )

        if n_edge_z is not None:        
            self.fc_edge_feature = nn.Sequential(
                nn.Linear(n_edge_z, n_edge_z),
                nn.LeakyReLU(negative_slope=0.2),
                nn.Linear(n_edge_z, n_node*n_edge_z)      #feature dim n_node*n_edge_z, prepare to evaluate inner-product to determine adjacency
            )
        else:
            self.fc_edge_feature = None

        self.conv_output = nn.Sequential(
            nn.Conv1d(in_channels=n_z+n_node*n_edge_x+3+2, out_channels=512, kernel_size=1, stride=1),
            nn.ReLU(),
            nn.Conv1d(in_channels=512, out_channels=512, kernel_size=1, stride=1),
            nn.ReLU(),
            nn.Conv1d(in_channels=512, out_channels=3, kernel_size=1, stride=1)
        )

        self.conv_output_noedge = nn.Sequential(
            nn.Conv1d(in_channels=n_z+3+2, out_channels=512, kernel_size=1, stride=1),
            nn.ReLU(),
            nn.Conv1d(in_channels=512, out_channels=512, kernel_size=1, stride=1),
            nn.ReLU(),
            nn.Conv1d(in_channels=512, out_channels=3, kernel_size=1, stride=1)
        )

        def init_xavier_normal(m):
            if type(m) == nn.Linear:
                nn.init.xavier_normal_(m.weight)

        self.fc_node_pos.apply(init_xavier_normal)
        self.fc_node_pos_zcomp.apply(init_xavier_normal)
        self.conv_output.apply(init_xavier_normal)
        self.conv_output_noedge.apply(init_xavier_normal)
        
        return

    def generatePointPatch(self):
        '''
        method to generate finer point clouds for subsequent warping
        One idea is like PCN, the node serves as the center of plane
        The other is using node as the centroid of a sphere
        '''
        g = torch.linspace(-self.comp_radius, self.comp_radius, self.grid_size)
            
        grid = torch.meshgrid(g, g)

        return grid
    
    def forward(self, z, edge_z=None):
        nodes = self.fc_node_pos(z).view(-1, self.n_node, 3)

        gridPatch = self.generatePointPatch()
        if z.is_cuda:
            gridPatch = gridPatch.cuda()
        
        patchFeature = torch.stack(gridPatch, 2).view(-1, 2).unsqueeze(0).repeat(z.shape[0], self.n_node, 1)    #(batch, n_output_points, 2) coordinates on the patch
        nodeFeature = nodes.unsqueeze(2).repeat(1, 1, self.grid_size**2, 1).view(-1, self.n_output_points, 3)   #(batch, n_output_points, 3) 

        adjacencyMatrix = None

        if edge_z is not None:
            #prepare a adjacency matrix here
            output = self.conv_output
            #the size of edge_z is (batch, n_edge_z)
            assert(self.fc_edge_feature is not None)

            nodeFeatureForEdge = self.fc_edge_feature(edge_z).view(-1, self.n_node, self.n_edge_z)
            
            adjacencyMatrix = nn.Sigmoid()(torch.bmm(nodeFeatureForEdge, nodeFeatureForEdge.transpose(1, 2)))
            
            adjacencyMatrixRepeat = adjacencyMatrix.unsqueeze(2).repeat(1, 1, self.grid_size**2, 1).view(-1, self.n_output_points, self.n_node)

            nodeFeature = torch.cat((nodeFeature, adjacencyMatrixRepeat), dim=2)   #(batch, n_output_points, n_node+3)
        else:
            output = self.conv_output_noedge

        globalFeature = z.unsqueeze(1).expand(-1, self.n_output_points, -1)         #(batch, n_output_points, n_z)

        concatFeature = torch.cat((patchFeature, nodeFeature, globalFeature), dim=2).transpose(1, 2)   #exchange channel and n_fine to make it consistent with NCWH  

        center = nodes.unsqueeze(2).repeat(1, 1, self.grid_size**2, 1).view(-1, self.n_output_points, 3).transpose(1, 2)    #(batch, 3, n_output_points)

        fine = output(concatFeature) + center

        # return associated adjacency matrix as well if needed
        return nodes.transpose(1, 2), adjacencyMatrix, fine
    



