from cg.src.info.betavae import BetaVAEInfo
from cg.src.info.betatcvae import BetaTCVAEInfo
from cg.src.info.clgvae import CLGVAEInfo
from cg.src.info.cmcs import CMCS_GT_Info
from cg.src.info.cmcs import CMCS_Super_Info
from cg.src.info.betavae_maganet import BetaMAGAVAEInfo
# from cg.src.info.cmcs import CMCS_UnSuper_Info









INFO={
    "betavae": BetaVAEInfo,
    "betatcvae": BetaTCVAEInfo,
    "clgvae": CLGVAEInfo,
    "betavae_maganet": BetaMAGAVAEInfo,
    "cmcs_gt": CMCS_GT_Info,
    "cmcs_super": CMCS_Super_Info,
    # "cmcs_unsuper": CMCS_UnSuper_Info
}

def write_info(args, results):
    info = INFO[args.model_type](args, **results)
    info.write_results()
    return