from cg.src.info.base import BaseInfo
from cg.src.info.betavae import BetaVAEInfo

class CMCS_GT_Info(BaseInfo):
    def __init__(self, args, **kwargs):
        super(CMCS_GT_Info, self).__init__(args, **kwargs)

        self.alpha = args.alpha
        self.beta = args.beta
        self.gamma = args.gamma
        self.nth_root = args.nth_root
        self.elbo = kwargs["elbo"]
        self.reconst = kwargs["reconst"]
        self.kld = kwargs["kld"]
        self.code_loss = kwargs["code_loss"]
        self.canonical_loss = kwargs["canonical_loss"]
        self.decoder_equiv = kwargs["decoder_equiv"]


class CMCS_Super_Info(CMCS_GT_Info):
    def __init__(self, args, **kwargs):
        super(CMCS_Super_Info, self).__init__(args, **kwargs)
        self.label_loss = kwargs["label_loss"]
        self.label_acc = kwargs["label_acc"]


class CMCS_UnSuper_Info(BetaVAEInfo):
    def __init__(self, args, **kwargs):
        super(CMCS_UnSuper_Info, self).__init__(args, **kwargs)
        self.gamma = args.gamma
        self.abs_diff = kwargs["abs_diff"]
        self.prior_list = args.prior_list






