from dargs import Argument


def nvnmd_args():
    doc_net_size_file = "configuration the number of nodes of fitting_net, just can be set as 128"
    doc_map_file = "A file containing the mapping tables to replace the calculation of embedding nets"
    doc_config_file = "A file containing the parameters about how to implement the model in certain hardware"
    doc_weight_file = "a *.npy file containing the weights of the model"
    doc_enable = "enable the nvnmd training"
    doc_restore_descriptor = "enable to restore the parameter of embedding_net from weight.npy"
    doc_restore_fitting_net = "enable to restore the parameter of fitting_net from weight.npy"
    doc_quantize_descriptor = "enable the quantizatioin of descriptor"
    doc_quantize_fitting_net = "enable the quantizatioin of fitting_net"
    args = [
        Argument("net_size", int, optional=False, default=128, doc=doc_net_size_file),
        Argument("map_file", str, optional=False, default='none', doc=doc_map_file),
        Argument("config_file", str, optional=False, default='none', doc=doc_config_file),
        Argument("weight_file", str, optional=False, default='none', doc=doc_weight_file),
        Argument("enable", bool, optional=False, default=False, doc=doc_enable),
        Argument("restore_descriptor", bool, optional=False, default=False, doc=doc_restore_descriptor),
        Argument("restore_fitting_net", bool, optional=False, default=False, doc=doc_restore_fitting_net),
        Argument("quantize_descriptor", bool, optional=False, default=False, doc=doc_quantize_descriptor),
        Argument("quantize_fitting_net", bool, optional=False, default=False, doc=doc_quantize_fitting_net),
    ]

    doc_nvnmd = 'The nvnmd options.'
    return Argument("nvnmd", dict, args, [], optional=True, doc = doc_nvnmd)