from dl.src.info.base import BaseInfo


class BetaVAEInfo(BaseInfo):
    def __init__(self, args, **kwargs):
        super(BetaVAEInfo, self).__init__(args, **kwargs)
        self.beta = args.beta

        self.elbo = kwargs["elbo"]
        self.obj = kwargs['obj']
        self.reconst = kwargs["reconst"]
        self.kld = kwargs["kld"]

        self.beta_vae = kwargs["beta_vae"]
        self.fvm = kwargs["factor_disent"]
        self.mig = kwargs["mig"]
        self.sap = kwargs["sap"]
        self.dci_disent = kwargs["dci_disent"]
        self.dci_completness = kwargs["dci_comple"]