from project.utils.deepinter_modules import BasicBlock_ResNetV2_1D, BasicBlock_Inception2D_V1

# Parse and return arguments for Triangle model
def config_parse_args():

    dist_bool = False

    model_config = {
        # model hyper-parameters
        "model_args" : {
            "BasicBlock1D" : {
                #"name"        : BasicBlock_ResNetV2_1D,
                "InChannels"  : 1280,
                "Channels"    : [64, 64, 64],
                "OutChannels" : 2,
                "num_Cycle"   : 1,
                "Kernel_size" : 17,
                "Dilation"    : [1],
                "Group"       : 1,
                "Bias"        : False,
                "track_running_stats" : False
            },

            "BasicBlock2D" : {
                "name"        : BasicBlock_Inception2D_V1,
                "InChannels"  : 64,
                "Channels"    : [64 for i in range(4)],
                "OutChannels" : 64,
                "num_Cycle"   : 1,
                "Kernel_size" : 3,
                "Dilation"    : [1],
                "Group"       : 1,
                "Bias"        : False,
                "track_running_stats" : False
            },
            'dist' : dist_bool,
            'dist_bins' : 42
        },


        "triangle_conv_args" : {
            "BasicBlock1D" : {
                #"name"        : BasicBlock_ResNetV2_1D,
                "InChannels"  : 1280,
                "Channels"    : [64, 64, 64],
                "OutChannels" : 2,
                "num_Cycle"   : 1,
                "Kernel_size" : 17,
                "Dilation"    : [1],
                "Group"       : 1,
                "Bias"        : False,
                "track_running_stats" : False
            },

            "BasicBlock2D" : {
                "name"        : BasicBlock_Inception2D_V1,
                "InChannels"  : 64,
                "Channels"    : [64 for i in range(4)],
                "OutChannels" : 64,
                "num_Cycle"   : 1,
                "Kernel_size" : 3,
                "Dilation"    : [1],
                "Group"       : 1,
                "Bias"        : False,
                "track_running_stats" : False
            },
            'dist' : dist_bool,
            'dist_bins' : 42
        },
        
        # Triangle hyper-parameters
        "triangle_args" : {
            "Channel_z"     : 64,
            "num_head"      : 4,
            "Channel_c"     : 8,
            "Transition_n"  : 4,

            "num_TriangleMulti" : 3, #20,
            "num_TriangleSelfR" : 3, #20,
            "num_TriangleSelfC" : 3, #20,
            "num_Transition"    : 3, #20,

            "final" : 64,

        }
    }
    
    return model_config 
