# Warsaw University of Technology

from layers.eca_block import ECABasicBlock
from models.minkgl import MinkHead, MinkTrunk, MinkGL


from models.minkloc import MinkLoc
from third_party.minkloc3d.minkloc import MinkLoc3D
from misc.utils import ModelParams


def model_factory(model_params: ModelParams):
    in_channels = 1

    if model_params.model == 'MinkLoc':
        model = MinkLoc(in_channels=in_channels,  feature_size=model_params.feature_size,
                        output_dim=model_params.output_dim, planes=model_params.planes,
                        layers=model_params.layers, num_top_down=model_params.num_top_down,
                        conv0_kernel_size=model_params.conv0_kernel_size, block=model_params.block,
                        pooling_method=model_params.pooling)
    elif model_params.model == 'MinkLoc3D':
        model = MinkLoc3D()
    elif 'egonn' in model_params.model:
        model = create_egonn_model(model_params)
    else:
        raise NotImplementedError('Model not implemented: {}'.format(model_params.model))

    return model


def create_egonn_model(model_params: ModelParams):
    model_name = model_params.model

    global_normalize = False
    local_normalize = True

    if model_name == 'egonn':
        # THIS IS OUR BEST MODEL
        block = ECABasicBlock
        planes = [32, 64, 64, 128, 128, 128, 128]
        layers = [1, 1, 1, 1, 1, 1, 1]

        global_in_levels = [5, 6, 7]
        global_map_channels = 128
        global_descriptor_size = 256

        local_in_levels = [3, 4]
        local_map_channels = 64
        local_descriptor_size = 128

    else:
        raise NotImplementedError(f'Unknown model: {model_name}')

    # Planes list number of channels for level 1 and above
    global_in_channels = [planes[i-1] for i in global_in_levels]
    head_global = MinkHead(global_in_levels, global_in_channels, global_map_channels)

    if len(local_in_levels) > 0:
        local_in_channels = [planes[i-1] for i in local_in_levels]
        head_local = MinkHead(local_in_levels, local_in_channels, local_map_channels)
    else:
        head_local = None

    min_out_level = len(planes)
    if len(global_in_levels) > 0:
        min_out_level = min(min_out_level, min(global_in_levels))
    if len(local_in_levels) > 0:
        min_out_level = min(min_out_level, min(local_in_levels))

    trunk = MinkTrunk(in_channels=1, planes=planes, layers=layers, conv0_kernel_size=5, block=block,
                      min_out_level=min_out_level)

    net = MinkGL(trunk, local_head=head_local, local_descriptor_size=local_descriptor_size,
                 local_normalize=local_normalize, global_head=head_global,
                 global_descriptor_size=global_descriptor_size, global_pool_method='GeM',
                 global_normalize=global_normalize, quantizer=model_params.quantizer)

    return net
[((1835, 1902), 'models.minkgl.MinkHead', 'MinkHead', (['global_in_levels', 'global_in_channels', 'global_map_channels'], {}), '(global_in_levels, global_in_channels, global_map_channels)\n', (1843, 1902), False, 'from models.minkgl import MinkHead, MinkTrunk, MinkGL\n'), ((2370, 2491), 'models.minkgl.MinkTrunk', 'MinkTrunk', ([], {'in_channels': '(1)', 'planes': 'planes', 'layers': 'layers', 'conv0_kernel_size': '(5)', 'block': 'block', 'min_out_level': 'min_out_level'}), '(in_channels=1, planes=planes, layers=layers, conv0_kernel_size=5,\n    block=block, min_out_level=min_out_level)\n', (2379, 2491), False, 'from models.minkgl import MinkHead, MinkTrunk, MinkGL\n'), ((2521, 2821), 'models.minkgl.MinkGL', 'MinkGL', (['trunk'], {'local_head': 'head_local', 'local_descriptor_size': 'local_descriptor_size', 'local_normalize': 'local_normalize', 'global_head': 'head_global', 'global_descriptor_size': 'global_descriptor_size', 'global_pool_method': '"""GeM"""', 'global_normalize': 'global_normalize', 'quantizer': 'model_params.quantizer'}), "(trunk, local_head=head_local, local_descriptor_size=\n    local_descriptor_size, local_normalize=local_normalize, global_head=\n    head_global, global_descriptor_size=global_descriptor_size,\n    global_pool_method='GeM', global_normalize=global_normalize, quantizer=\n    model_params.quantizer)\n", (2527, 2821), False, 'from models.minkgl import MinkHead, MinkTrunk, MinkGL\n'), ((381, 716), 'models.minkloc.MinkLoc', 'MinkLoc', ([], {'in_channels': 'in_channels', 'feature_size': 'model_params.feature_size', 'output_dim': 'model_params.output_dim', 'planes': 'model_params.planes', 'layers': 'model_params.layers', 'num_top_down': 'model_params.num_top_down', 'conv0_kernel_size': 'model_params.conv0_kernel_size', 'block': 'model_params.block', 'pooling_method': 'model_params.pooling'}), '(in_channels=in_channels, feature_size=model_params.feature_size,\n    output_dim=model_params.output_dim, planes=model_params.planes, layers=\n    model_params.layers, num_top_down=model_params.num_top_down,\n    conv0_kernel_size=model_params.conv0_kernel_size, block=model_params.\n    block, pooling_method=model_params.pooling)\n', (388, 716), False, 'from models.minkloc import MinkLoc\n'), ((2025, 2089), 'models.minkgl.MinkHead', 'MinkHead', (['local_in_levels', 'local_in_channels', 'local_map_channels'], {}), '(local_in_levels, local_in_channels, local_map_channels)\n', (2033, 2089), False, 'from models.minkgl import MinkHead, MinkTrunk, MinkGL\n'), ((856, 867), 'third_party.minkloc3d.minkloc.MinkLoc3D', 'MinkLoc3D', ([], {}), '()\n', (865, 867), False, 'from third_party.minkloc3d.minkloc import MinkLoc3D\n')]