from .crog import CROG
from .ssg import SSG
from loguru import logger


def build_crog(args):
    model = CROG(args)
    backbone = []
    head = []
    for k, v in model.named_parameters():
        if k.startswith('backbone') and 'positional_embedding' not in k:
            backbone.append(v)
        else:
            head.append(v)
    logger.info('Backbone with decay={}, Head={}'.format(len(backbone), len(head)))
    param_list = [{
        'params': backbone,
        'initial_lr': args.lr_multi * args.base_lr
    }, {
        'params': head,
        'initial_lr': args.base_lr
    }]
    return model, param_list


def build_ssg(args):
    model = SSG(args)

    return model, model.parameters()

def build_net(args):
    network_name = args.network_name.lower()
    if network_name == 'legrconvnet':
        from .legrconv import CORE
    elif network_name == 'legrconvnet4x':
        from .legrconv4x import CORE
    elif network_name == 'legrconvnet4x_update1':
        from .legrconv4x_update1 import CORE
    elif network_name == 'legrconvnet4x_coord':
        from .legrconv4x_coordconv import CORE
    elif network_name == 'legrconvnet4x_dcn':
        from .legrconv4x_dcn import CORE
    elif network_name == 'legrconvnet4x_tacmt':
        from .legrconv4x_TACMT import CORE
    elif network_name == 'legrconv4x_dcn_txtaware':
        from .legrconv4x_dcn_txtaware import CORE
    elif network_name == 'legrconv4x_film1':
        from .legrconv4x_film1 import CORE
    elif network_name == 'legrconv4x_film2':
        from .legrconv4x_film2 import CORE
    elif network_name == 'legrconv4x_film2_refine':
        from .legrconv4x_film2_refine import CORE
    elif network_name == 'legrconv4x_film3':
        from .legrconv4x_film3 import CORE
    elif network_name == 'legrconv4x_llava':
        from .legrconv4x_llava import CORE
    elif network_name == 'legrconv4x_llava_fixed':
        from .legrconv4x_llava_fixed import CORE
    elif network_name == 'legrconv4x_llava_v2':
        from .legrconv4x_llava_v2 import CORE
    elif network_name == 'legrconv4x_llava_v2_nomask':
        from .legrconv4x_llava_v2_nomask import CORE
    elif network_name == 'legrconv4x_llava_v2_new':
        from .legrconv4x_llava_v2_new import CORE
    elif network_name == 'legrconv4x_llava_v2_c':
        from .legrconv4x_llava_v2_c import CORE
    elif network_name in (
        'legrconv4x_llava_honeybee_aligned',
        'llava_honeybee_aligned',
        'honeybee_aligned_grasp',
    ):
        from .legrconv4x_llava_honeybee_aligned import CORE
    elif network_name in (
        'legrconv4x_llava_honeybee_aligned_3dcoord_newtrans',
        'llava_honeybee_aligned_3dcoord_newtrans',
        'honeybee_aligned_grasp_3dcoord_newtrans',
    ):
        from .legrconv4x_llava_honeybee_aligned_3dcoord_newtrans import CORE
    elif network_name in (
        'legrconv4x_llava_honeybee_aligned_native_rgb',
        'llava_honeybee_aligned_native_rgb',
        'honeybee_aligned_grasp_no_cabs',
    ):
        from .legrconv4x_llava_honeybee_aligned_native_rgb import CORE
    elif network_name in (
        'legrconv4x_llava_maple_generalized_rgb',
        'llava_maple_generalized_rgb',
        'maple_generalized_rgb',
    ):
        from .legrconv4x_llava_maple_generalized_rgb import CORE
    elif network_name in (
        'legrconv4x_llava_honeybee_aligned_dualsoft',
        'legrconv_llava_honeybee_aligned_dualsoft',
        'llava_honeybee_aligned_dualsoft',
        'honeybee_aligned_dualsoft_rgb',
    ):
        from .legrconv4x_llava_honeybee_aligned_dualsoft import CORE
    elif network_name in ('legrconv4x_llava_v2_c_statichead', 'v2_c_statichead'):
        from .legrconv4x_llava_v2_c_statichead import CORE
    elif network_name in ('legrconv4x_llava_v2_c_customenc_statichead', 'v2_c_customenc_statichead'):
        from .legrconv4x_llava_v2_c_customenc_statichead import CORE
    elif network_name in ('legrconv4x_llava_v2_c_standard_statichead', 'v2_c_standard_statichead'):
        from .legrconv4x_llava_v2_c_standard_statichead import CORE
    elif network_name in (
        'legrconv4x_llava_v2_c_standard_statichead_3dcoord_newtrans',
        'v2_c_standard_statichead_3dcoord_newtrans',
    ):
        from .legrconv4x_llava_v2_c_standard_statichead_3dcoord_newtrans import CORE
    elif network_name in ('legrconv4x_llava_v2_c_standard_dynamichead', 'v2_c_standard_dynamichead'):
        from .legrconv4x_llava_v2_c_standard_dynamichead import CORE
    elif network_name in (
        'legrconv4x_llava_v2_c_standard_statichead_clip_lora_3dcoord_newtrans',
        'v2_c_standard_statichead_clip_lora_3dcoord_newtrans',
    ):
        from .legrconv4x_llava_v2_c_standard_statichead_clip_lora_3dcoord_newtrans import CORE
    elif network_name in ('legrconv4x_llava_v2_c_3dcoord_newtrans', 'v2_c_3dcoord_newtrans'):
        from .legrconv4x_llava_v2_c_3dcoord_newtrans import CORE
    elif network_name in (
        'legrconv4x_llava_v2_c_standard_statichead_clip_lora',
        'v2_c_standard_statichead_clip_lora',
    ):
        from .legrconv4x_llava_v2_c_standard_statichead_clip_lora import CORE
    elif network_name == 'legrconv4x_molmo_v2':
        from .legrconv4x_molmo_v2 import CORE
    elif network_name == 'legrconv4x_molmo_v2_c':
        from .legrconv4x_molmo_v2_c import CORE
    elif network_name in ('legrconv4x_molmo_v2_c_lora', 'molmo_v2_c_lora'):
        from .legrconv4x_molmo_v2_c_lora import CORE
    elif network_name == 'legrconv4x_qwen_vl_v2':
        from .legrconv4x_qwen_vl_v2 import CORE
    elif network_name == 'legrconv4x_qwen_vl_v2_c':
        from .legrconv4x_qwen_vl_v2_c import CORE
    elif network_name == 'legrconv4x_qwen_vl_v2_c_lora':
        from .legrconv4x_qwen_vl_v2_c_lora import CORE
    elif network_name == 'legrconv4x_qwen_vl_v2_c_lora_newtrans':
        from .legrconv4x_qwen_vl_v2_c_lora_newtrans import CORE
    elif network_name in ('legrconv4x_llava_v2_c_3dcoord', 'v2_c_3dcoord'):
        from .legrconv4x_llava_v2_c_3dcoord import CORE
    elif network_name in ('legrconv4x_llava_v2_c_3dcoord_frontcoord', 'v2_c_3dcoord_frontcoord'):
        from .legrconv4x_llava_v2_c_3dcoord import CORE
    elif network_name in ('legrconv4x_llava_v2_c_multigeo', 'v2_c_multigeo'):
        from .legrconv4x_llava_v2_c_multigeo import CORE
    elif network_name in ('legrconv4x_llava_v2_c_singlegeo', 'v2_c_singlegeo'):
        from .legrconv4x_llava_v2_c_singlegeo import CORE
    elif network_name in (
        'legrconv4x_llava_v2_c_standard_statichead_clip_3dcoord_newtrans',
        'v2_c_standard_statichead_clip_3dcoord_newtrans',
    ):
        from .legrconv4x_llava_v2_c_standard_statichead_clip_3dcoord_newtrans import CORE
    elif network_name == 'legrconv4x_llava_v2_d':
        from .legrconv4x_llava_v2_d import CORE
    elif network_name == 'legrconv4x_llava_v2_refine':
        from .legrconv4x_llava_v2_refine import CORE
    elif network_name == 'legrconv4x_llava_v2_refinemix':
        from .legrconv4x_llava_v2_refinemix import CORE
    elif network_name == 'legrconv4x_llava_v2_maple':
        from .legrconv4x_llava_v2_maple import CORE
    elif network_name in (
        'legrconv4x_llava_v2_maple_standard_statichead',
        'maple_standard_statichead',
        'legrconv4x_llava_v2_maple_standard_statichead_clip',
        'maple_standard_statichead_clip',
    ):
        from .legrconv4x_llava_v2_maple_standard_statichead import CORE
    elif network_name in (
        'legrconv4x_llava_v2_maple_standard_statichead_3dcoord_newtrans',
        'maple_standard_statichead_3dcoord_newtrans',
        'legrconv4x_llava_v2_maple_standard_statichead_clip_3dcoord_newtrans',
        'maple_standard_statichead_clip_3dcoord_newtrans',
    ):
        from .legrconv4x_llava_v2_maple_standard_statichead_3dcoord_newtrans import CORE
    elif network_name in (
        'legrconv4x_llava_v2_maple_3dcoord_newtrans',
        'legrconv4x_llava_v2_maple_depthcoord_newtrans',
        'maple_3dcoord_newtrans',
    ):
        from .legrconv4x_llava_v2_maple_3dcoord_newtrans import CORE
    elif network_name == 'legrconv4x_llava_v2_maple_soft':
        from .legrconv4x_llava_v2_maple_soft import CORE
    elif network_name == 'legrconv4x_llava_v2_maple_refine':
        from .legrconv4x_llava_v2_maple_refine import CORE
    elif network_name == 'legrconv4x_llava_v2_maple_refine_split':
        from .legrconv4x_llava_v2_maple_refine_split import CORE
    elif network_name == 'legrconv4x_llava_v2_geo':
        from .legrconv4x_llava_v2_geo import CORE
    elif network_name == 'legrconv4x_llava_v2_geo2':
        from .legrconv4x_llava_v2_geo2 import CORE
    elif network_name == 'legrconv4x_llava_v2_geo3':
        from .legrconv4x_llava_v2_geo3 import CORE
    elif network_name == 'legrconv4x_llava_v2_geo4':
        from .legrconv4x_llava_v2_geo4 import CORE
    elif network_name == 'legrconv4x_llava_lora':
        from .legrconv4x_llava_lora import CORE
    elif network_name == 'legrconv4x_llava_v2_c_lora':
        from .legrconv4x_llava_v2_c_lora import CORE
    elif network_name == 'legrconv4x_llava_v2_c_lora_newtrans':
        from .legrconv4x_llava_v2_c_lora_newtrans import CORE
    elif network_name == 'legrconv4x_llava_v2_c_lora_3dcoord_newtrans':
        from .legrconv4x_llava_v2_c_lora_3dcoord_newtrans import CORE
    elif network_name == 'graspanything_withmask':
        from .legrconv4x_llava_v2_ga_wmask import CORE
    elif network_name == 'graspanything_withoutmask':
        from .legrconv4x_llava_v2_ga_womask import CORE
    elif network_name in ('crog', 'crog_r50'):
        from .crog import CROG as CORE
    else:
        raise NotImplementedError('Network {} is not implemented'.format(network_name))
    model = CORE(args)
    
    # Special handling for models with custom parameter management
    if hasattr(model, 'get_trainable_params'):
        # Models like llava_lora handle their own parameter groups
        param_list = model.get_trainable_params()
    else:
        # Default behavior for other models
        # Freeze backbone parameters
        for name, param in model.named_parameters():
            if name.startswith('backbone'):
                param.requires_grad = False
        backbone = []
        head = []
        for k, v in model.named_parameters():
            if not v.requires_grad:
                continue
            if k.startswith('backbone') and 'positional_embedding' not in k:
                backbone.append(v)
            else:
                head.append(v)
        logger.info('Backbone with decay={}, Head={}'.format(len(backbone), len(head)))
        param_list = [ {
            'params': head,
            'initial_lr': args.base_lr
        }]
    return model, param_list