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)