from dl.src.info.base import BaseInfo


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

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

        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"]