import megengine as mge
import megengine.module as nn
import megengine.functional as F
from model.module import Encoder, Fusion, Decoder, Regression
from common import se3, quaternion
import math


class OMNet(nn.Module):
    def __init__(self, params):
        super(OMNet, self).__init__()
        self.num_iter = params.titer
        self.encoder = [Encoder() for _ in range(self.num_iter)]
        self.fusion = [Fusion() for _ in range(self.num_iter)]
        self.decoder = [Decoder() for _ in range(self.num_iter)]
        self.regression = [Regression() for _ in range(self.num_iter)]
        self.overlap_dist = params.overlap_dist

        for m in self.modules():
            if isinstance(m, nn.Conv1d) or isinstance(m, nn.Linear):
                nn.init.msra_normal_(m.weight, a=math.sqrt(5))
                if m.bias is not None:
                    fan_in, _ = nn.init.calculate_fan_in_and_fan_out(m.weight)
                    bound = 1 / math.sqrt(fan_in)
                    nn.init.uniform_(m.bias, -bound, bound)
            # elif isinstance(m, nn.BatchNorm1d):
            #     nn.init.ones_(m.weight)
            #     nn.init.zeros_(m.bias)

    def generate_overlap_mask(self, points_src, points_ref, mask_src, mask_ref, transform_gt):
        points_src[F.logical_not(mask_src.astype("bool")), :] = 50.0
        points_ref[F.logical_not(mask_ref.astype("bool")), :] = 100.0
        points_src = se3.mge_transform(transform_gt, points_src)
        points_src = F.expand_dims(points_src, axis=2)
        points_ref = F.expand_dims(points_ref, axis=1)
        dist_matrix = F.sqrt(F.sum(F.square(points_src - points_ref), axis=-1))  # (B, N, N)
        dist_s2r = F.min(dist_matrix, axis=2)
        dist_r2s = 
F.min(dist_matrix, axis=1)