from typing import List, Optional, Tuple

import jittor as jt
import jittor.nn as nn
from jittor import init
from jittor.contrib import concat

from models.misc.ops import FurthestPointSampler
from models.misc.ops import BallQueryGrouper
from models.misc.ops import GroupAll
from models.misc.ops import PointNetFeaturePropagation


class PointNetModuleBase(nn.Module):
    def __init__(self):
        self.n_points = None
        self.sampler = None
        self.groupers = None
        self.mlps = None

    def build_mlps(self, mlp_spec: List[int], use_xyz: bool = True,
                   bn: bool = True) -> nn.Sequential:
        layers = []

        if use_xyz:
            mlp_spec[0] += 3

        for i in range(1, len(mlp_spec)):
            layers.append(nn.Conv2d(mlp_spec[i - 1], mlp_spec[i], kernel_size=1))
            if bn:
                layers.append(nn.BatchNorm2d(mlp_spec[i]))
            layers.append(nn.ReLU())

        return nn.Sequential(*layers)

    def execute(self, xyz: jt.Var, feature: Optional[jt.Var]) -> Tuple[jt.Var, jt.Var]:
        '''
        Parameters
        ----------
        xyz: jt.Var, (B, N, 3)
        feature: jt.Var, (B, C, N)

        Returns
        -------
        new_xyz: jt.Var, (B, n_points, 3)
        new_feature: jt.Var, (B, C', n_points)
        '''
        B, _, _ = xyz.shape
        new_xyz = self.sampler(xyz) if self.n_points is not None else jt.zeros((B, 1, 3))

        new_feature_list = []
        # print (self.groupers)
        for i, grouper in self.groupers.layers.items():
            new_feature = grouper(new_xyz, xyz, feature.transpose(0, 2, 1))
            # [B, n_points, n_samples, C] -> [B, C, n_points, n_samples]
            # print('\n ---------{}--------------.'.format(i+1))
            new_feature = new_feature.transpose(0, 3, 1, 2)
            # print('1. new_feature.shape', new_feature.shape)
            new_feature = self.mlps[i](new_feature)
            # print('2. new_feature.shape', new_feature.shape)
            # [B, C, n_points, n_samples] -> [B, n_points, n_samples, C]
            #new_feature = new_feature.transpose(0, 2, 3, 1)
            # print('3. new_feature.shape', new_feature.shape)
            new_feature = new_feature.argmax(dim=-1)[1]
            # print('4. new_feature.shape', new_feature.shape)

            new_feature_list.append(new_feature)

        new_feature = jt.contrib.concat(new_feature_list, dim=-1)
        # print('len(new_feature):', len(new_feature))
        return new_xyz, new_feature


class PointnetModule(PointNetModuleBase):
    def __init__(self, mlp: List[int], n_points=None, radius=None,
                 n_samples=None, bn=True, use_xyz=True):
        super().__init__()

        self.n_points = n_points

        self.groupers = nn.ModuleList()
        if self.n_points is not None:
            self.sampler = FurthestPointSampler(n_points)
            self.groupers.append(BallQueryGrouper(radius, n_samples, use_xyz))
        else:
            self.groupers.append(GroupAll(use_xyz))

        self.mlps = nn.ModuleList()
        self.mlps.append(self.build_mlps(mlp, use_xyz))


class PointnetModuleMSG(PointNetModuleBase):
    def __init__(self, n_points: int, radius: List[float], n_samples: List[int],
                 mlps: List[List[int]], bn=True, use_xyz=True):
        super().__init__()

        self.n_points = n_points
        self.sampler = FurthestPointSampler(n_points)

        self.groupers = nn.ModuleList()
        for r, s in zip(radius, n_samples):
            self.groupers.append(BallQueryGrouper(r, s, use_xyz))

        self.mlps = nn.ModuleList()
        for mlp in mlps:
            self.mlps.append(self.build_mlps(mlp, use_xyz))


class PointNet2_partseg(nn.Module):
    def __init__(self, part_num=50, use_xyz=True):
        super().__init__()
        self.part_num = part_num
        self.use_xyz = use_xyz
        self.build_model()

    def build_model(self):
        self.pointnet_modules = nn.ModuleList()
        self.pointnet_modules.append(
            PointnetModule(
                n_points=512,
                radius=0.2,
                n_samples=64,
                mlp=[3, 64, 64, 128],
                use_xyz=self.use_xyz,
            )
        )

        self.pointnet_modules.append(
            PointnetModule(
                n_points=128,
                radius=0.4,
                n_samples=64,
                mlp=[128, 128, 128, 256],
                use_xyz=self.use_xyz,
            )
        )

        self.pointnet_modules.append(
            PointnetModule(
                mlp=[256, 256, 512, 1024],
                use_xyz=self.use_xyz,
            )
        )

        self.fp3 = PointNetFeaturePropagation(in_channel=1280, mlp=[256, 256])
        self.fp2 = PointNetFeaturePropagation(in_channel=384, mlp=[256, 128])
        self.fp1 = PointNetFeaturePropagation(in_channel=128 + 16 + 6, mlp=[128, 128, 128])

        self.fc_layer = nn.Sequential(
            nn.Conv1d(128, 128, 1),
            nn.BatchNorm1d(128),
            nn.Dropout(0.5),
            nn.Conv1d(128, self.part_num, 1)
        )

    def execute(self, xyz, feature, cls_label):
        # for module in self.pointnet_modules:
        #     xyz, feature = module(xyz, feature)

        B, N, _ = xyz.shape
        l1_xyz, l1_feature = self.pointnet_modules[0](xyz, feature)
        l2_xyz, l2_feature = self.pointnet_modules[1](l1_xyz, l1_feature)
        l3_xyz, l3_feature = self.pointnet_modules[2](l2_xyz, l2_feature)
        # print ('before interpolate shape')
        # print (l2_xyz.shape, l2_feature.shape, l3_xyz.shape, l3_feature.shape)
        l2_feature = self.fp3(l2_xyz, l3_xyz, l2_feature, l3_feature)
        l1_feature = self.fp2(l1_xyz, l2_xyz, l1_feature, l2_feature)
        cls_label_one_hot = cls_label.view(B, 16, 1).repeat(1, 1, N).permute(0, 2, 1)
        # print ('before concat size ')
        # print (cls_label_one_hot.size(),xyz.size(),feature.size())
        feature = self.fp1(xyz, l1_xyz, concat([cls_label_one_hot, xyz, feature], 2), l1_feature)
        feature = feature.permute(0, 2, 1)
        # print (feature.shape)
        return self.fc_layer(feature)


class PointNetMSG(PointNet2_partseg):
    def build_model(self):
        super().build_model()

        self.pointnet_modules = nn.ModuleList()
        self.pointnet_modules.append(
            PointnetModuleMSG(
                n_points=512,
                radius=[0.1, 0.2, 0.4],
                n_samples=[16, 32, 128],
                mlps=[[3, 32, 32, 64], [3, 64, 64, 128], [3, 64, 96, 128]],
                use_xyz=self.use_xyz,
            )
        )

        input_channels = 64 + 128 + 128
        self.pointnet_modules.append(
            PointnetModuleMSG(
                n_points=128,
                radius=[0.2, 0.4, 0.8],
                n_samples=[32, 64, 128],
                mlps=[
                    [input_channels, 64, 64, 128],
                    [input_channels, 128, 128, 256],
                    [input_channels, 128, 128, 256],
                ],
                use_xyz=self.use_xyz,
            )
        )

        self.pointnet_modules.append(
            PointnetModule(
                mlp=[128 + 256 + 256, 256, 512, 1024],
                use_xyz=self.use_xyz,
            )
        )


def main():
    model = PointNet2_partseg()
    input_point = init.gauss([2, 1024, 3], 'float', mean=0.0)
    input_feature = init.gauss([2, 1024, 3], 'float', mean=0.0)
    cls_label = init.gauss([2, 16], 'float', mean=0.0)

    print(input_point.shape)
    print(input_feature.shape)
    print(cls_label.shape)
    outputs = model(input_point, input_feature, cls_label)
    print(outputs.shape)


if __name__ == '__main__':
    main()
