
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.autograd as autograd
from torch.autograd import Variable
from torch.optim import lr_scheduler
from torchvision import transforms
import copy
import numpy as np
import random
import higher
from collections import OrderedDict

try:
    from backpack import backpack, extend
    from backpack.extensions import BatchGrad
except:
    backpack = None

import networks
from lib.misc import (
        random_pairs_of_minibatches, split_meta_train_test, ParamDict,
        MovingAverage, l2_between_dicts, proj,Augmix
    )



ALGORITHMS = [
    'ERM',
    # 'Fish',
    'IRM',
    'GroupDRO',
    'Mixup',
    'MLDG',
    'CORAL',
    'MMD',
    'DANN',
    'CDANN',
    'MTL',
    'SagNet',
    'ARM',
    'VREx',
    'RSC',
    'SD',
    # 'DDG_AugMix',
    # 'DDG',
    # 'ANDMask',
    # 'SANDMask',
    # 'IGA',
    'SelfReg',
    # "Fishr",
    'TRM',
    'IB_ERM',
    # 'IB_IRM',
    'CAD',
    'CondCAD',
    'Transfer',
    # 'CausIRL_CORAL',
    'CausIRL_MMD',
    'MAMR',
    'MAML',
    "RSD"
]

def get_algorithm_class(algorithm_name):
    """Return the algorithm class with the given name."""
    if algorithm_name not in globals():
        raise NotImplementedError("Algorithm not found: {}".format(algorithm_name))
    return globals()[algorithm_name]


class Algorithm(torch.nn.Module):
    """
    A subclass of Algorithm implements a domain generalization algorithm.
    Subclasses should implement the following:
    - update()
    - predict()
    """
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(Algorithm, self).__init__()
        self.hparams = hparams
        self.device = device

    def update(self, minibatches, unlabeled=None):
        """
        Perform one update step, given a list of (x, y) tuples for all
        environments.
        Admits an optional list of unlabeled minibatches from the test domains,
        when task is domain_adaptation.
        """
        raise NotImplementedError

    def predict(self, x):
        raise NotImplementedError

class ERM(Algorithm):
    """
    Empirical Risk Minimization (ERM)
    """

    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(ERM, self).__init__(input_shape, num_classes, num_domains,
                                  hparams, device)
        self.featurizer = networks.Featurizer(input_shape, self.hparams)
        self.classifier = networks.Classifier(
            self.featurizer.n_outputs,
            num_classes,
            self.hparams['nonlinear_classifier'])

        self.network = nn.Sequential(self.featurizer, self.classifier)
        self.optimizer = torch.optim.Adam(
            self.network.parameters(),
            lr=self.hparams["lr"],
            weight_decay=self.hparams['weight_decay']
        )
        # self.optimizer = torch.optim.SGD(
        #     self.network.parameters(),
        #     lr=self.hparams["lr"],
        #     weight_decay=self.hparams['weight_decay']
        # )

    def update(self, minibatches, unlabeled=None):
        all_x = torch.cat([x for x,y in minibatches])
        all_y = torch.cat([y for x,y in minibatches])
        # print(self.predict(all_x).shape,all_y.shape)
        # loss = F.l1_loss(self.predict(all_x).view(-1), all_y.float().view(-1))
        loss = F.l1_loss(self.predict(all_x).view(-1), all_y.float().view(-1))

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return {'loss': loss.item()}

    def predict(self, x):
        return self.network(x)

# class ERM(Algorithm):
#     """
#     Empirical Risk Minimization (ERM)
#     """

#     def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
#         super(ERM, self).__init__(input_shape, num_classes, num_domains,
#                                   hparams, device)
#         self.featurizer = networks.Featurizer(input_shape, self.hparams)
#         self.classifier = networks.Classifier(
#             self.featurizer.n_outputs,
#             num_classes,
#             self.hparams['nonlinear_classifier'])

#         self.network = nn.Sequential(self.featurizer, self.classifier)
#         self.optimizer = torch.optim.Adam(
#             self.network.parameters(),
#             lr=self.hparams["lr"],
#             weight_decay=self.hparams['weight_decay']
#         )

#     def update(self, minibatches, unlabeled=None):
#         all_x = torch.cat([x for x,y in minibatches])
#         all_y = torch.cat([y for x,y in minibatches])
#         # print(self.predict(all_x).shape,all_y.shape)
#         loss = F.l1_loss(self.predict(all_x).view(-1), all_y.float().view(-1))

#         self.optimizer.zero_grad()
#         loss.backward()
#         self.optimizer.step()

#         return {'loss': loss.item()}

#     def predict(self, x):
#         return self.network(x)




class Fish(Algorithm):
    """
    Implementation of Fish, as seen in Gradient Matching for Domain
    Generalization, Shi et al. 2021.
    """

    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(Fish, self).__init__(input_shape, num_classes, num_domains,
                                   hparams)
        self.input_shape = input_shape
        self.num_classes = num_classes

        self.network = networks.WholeFish(input_shape, num_classes, hparams)
        self.optimizer = torch.optim.Adam(
            self.network.parameters(),
            lr=self.hparams["lr"],
            weight_decay=self.hparams['weight_decay']
        )
        self.optimizer_inner_state = None

    def create_clone(self, device):
        self.network_inner = networks.WholeFish(self.input_shape, self.num_classes, self.hparams,
                                            weights=self.network.state_dict()).to(device)
        self.optimizer_inner = torch.optim.Adam(
            self.network_inner.parameters(),
            lr=self.hparams["lr"],
            weight_decay=self.hparams['weight_decay']
        )
        if self.optimizer_inner_state is not None:
            self.optimizer_inner.load_state_dict(self.optimizer_inner_state)

    def fish(self, meta_weights, inner_weights, lr_meta):
        meta_weights = ParamDict(meta_weights)
        inner_weights = ParamDict(inner_weights)
        meta_weights += lr_meta * (inner_weights - meta_weights)
        return meta_weights

    def update(self, minibatches, unlabeled=None):
        self.create_clone(minibatches[0][0].device)

        for x, y in minibatches:
            loss = F.l1_loss(self.network_inner(x), y.view(-1))
            self.optimizer_inner.zero_grad()
            loss.backward()
            self.optimizer_inner.step()

        self.optimizer_inner_state = self.optimizer_inner.state_dict()
        meta_weights = self.fish(
            meta_weights=self.network.state_dict(),
            inner_weights=self.network_inner.state_dict(),
            lr_meta=self.hparams["meta_lr"]
        )
        self.network.reset_weights(meta_weights)

        return {'loss': loss.item()}

    def predict(self, x):
        return self.network(x)


class ARM(ERM):
    """ Adaptive Risk Minimization (ARM) """
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        original_input_shape = input_shape
        input_shape = (1 + original_input_shape[0],) + original_input_shape[1:]
        super(ARM, self).__init__(input_shape, num_classes, num_domains,
                                  hparams, device)
        self.context_net = networks.ContextNet(original_input_shape)
        self.support_size = hparams['batch_size']

    def predict(self, x):
        batch_size, c, h, w = x.shape
        if batch_size % self.support_size == 0:
            meta_batch_size = batch_size // self.support_size
            support_size = self.support_size
        else:
            meta_batch_size, support_size = 1, batch_size
        context = self.context_net(x)
        context = context.reshape((meta_batch_size, support_size, 1, h, w))
        context = context.mean(dim=1)
        context = torch.repeat_interleave(context, repeats=support_size, dim=0)
        x = torch.cat([x, context], dim=1)
        return self.network(x)


class AbstractDANN(Algorithm):
    """Domain-Adversarial Neural Networks (abstract class)"""

    def __init__(self, input_shape, num_classes, num_domains,
                 hparams, conditional, class_balance, device='cpu'):

        super(AbstractDANN, self).__init__(input_shape, num_classes, num_domains,
                                  hparams)

        self.register_buffer('update_count', torch.tensor([0]))
        self.conditional = conditional
        self.class_balance = class_balance

        # Algorithms
        self.featurizer = networks.Featurizer(input_shape, self.hparams)
        self.classifier = networks.Classifier(
            self.featurizer.n_outputs,
            num_classes,
            self.hparams['nonlinear_classifier'])
        self.discriminator = networks.MLP_(self.featurizer.n_outputs,
            num_domains, self.hparams)
        self.class_embeddings = nn.Embedding(num_classes,
            self.featurizer.n_outputs)

        # Optimizers
        self.disc_opt = torch.optim.Adam(
            (list(self.discriminator.parameters()) +
                list(self.class_embeddings.parameters())),
            lr=self.hparams["lr_d"],
            weight_decay=self.hparams['weight_decay_d'],
            betas=(self.hparams['beta1'], 0.9))

        self.gen_opt = torch.optim.Adam(
            (list(self.featurizer.parameters()) +
                list(self.classifier.parameters())),
            lr=self.hparams["lr_g"],
            weight_decay=self.hparams['weight_decay_g'],
            betas=(self.hparams['beta1'], 0.9))

    def update(self, minibatches, unlabeled=None):
        device = "cuda" if minibatches[0][0].is_cuda else "cpu"
        self.update_count += 1
        all_x = torch.cat([x for x, y in minibatches])
        all_y = torch.cat([y for x, y in minibatches])
        all_z = self.featurizer(all_x)
        if self.conditional:
            disc_input = all_z + self.class_embeddings(all_y)
        else:
            disc_input = all_z
        disc_out = self.discriminator(disc_input)
        disc_labels = torch.cat([
            torch.full((x.shape[0], ), i, dtype=torch.int64, device=device)
            for i, (x, y) in enumerate(minibatches)
        ])

        if self.class_balance:
            y_counts = F.one_hot(all_y).sum(dim=0)
            weights = 1. / (y_counts[all_y] * y_counts.shape[0]).float()
            disc_loss = F.cross_entropy(disc_out, disc_labels, reduction='none')
            disc_loss = (weights * disc_loss).sum()
        else:
            disc_loss = F.cross_entropy(disc_out, disc_labels)

        disc_softmax = F.softmax(disc_out, dim=1)
        input_grad = autograd.grad(disc_softmax[:, disc_labels].sum(),
            [disc_input], create_graph=True)[0]
        grad_penalty = (input_grad**2).sum(dim=1).mean(dim=0)
        disc_loss += self.hparams['grad_penalty'] * grad_penalty

        d_steps_per_g = self.hparams['d_steps_per_g_step']
        if (self.update_count.item() % (1+d_steps_per_g) < d_steps_per_g):

            self.disc_opt.zero_grad()
            disc_loss.backward()
            self.disc_opt.step()
            return {'disc_loss': disc_loss.item()}
        else:
            all_preds = self.classifier(all_z)
            classifier_loss = F.l1_loss(all_preds.view(-1), all_y.view(-1))
            gen_loss = (classifier_loss +
                        (self.hparams['lambda'] * -disc_loss))
            self.disc_opt.zero_grad()
            self.gen_opt.zero_grad()
            gen_loss.backward()
            self.gen_opt.step()
            return {'gen_loss': gen_loss.item()}

    def predict(self, x):
        return self.classifier(self.featurizer(x))

class DANN(AbstractDANN):
    """Unconditional DANN"""
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(DANN, self).__init__(input_shape, num_classes, num_domains,
            hparams, conditional=False, class_balance=False)


class CDANN(AbstractDANN):
    """Conditional DANN"""
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(CDANN, self).__init__(input_shape, num_classes, num_domains,
            hparams, conditional=True, class_balance=True)


class IRM(ERM):
    """Invariant Risk Minimization"""

    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(IRM, self).__init__(input_shape, num_classes, num_domains,
                                  hparams)
        self.register_buffer('update_count', torch.tensor([0]))

    @staticmethod
    def _irm_penalty(logits, y):
        device = "cuda" if logits[0][0].is_cuda else "cpu"
        scale = torch.tensor(1.).to(device).requires_grad_()
        loss_1 = F.l1_loss(logits[::2].view(-1) * scale, y[::2].view(-1))
        loss_2 = F.l1_loss(logits[1::2].view(-1) * scale, y[1::2].view(-1))
        grad_1 = autograd.grad(loss_1, [scale], create_graph=True)[0]
        grad_2 = autograd.grad(loss_2, [scale], create_graph=True)[0]
        result = torch.sum(grad_1 * grad_2)
        return result

    def update(self, minibatches, unlabeled=None):
        device = "cuda" if minibatches[0][0].is_cuda else "cpu"
        penalty_weight = (self.hparams['irm_lambda'] if self.update_count
                          >= self.hparams['irm_penalty_anneal_iters'] else
                          1.0)
        nll = 0.
        penalty = 0.

        all_x = torch.cat([x for x,y in minibatches])
        all_logits = self.network(all_x)
        all_logits_idx = 0
        for i, (x, y) in enumerate(minibatches):
            logits = all_logits[all_logits_idx:all_logits_idx + x.shape[0]]
            all_logits_idx += x.shape[0]
            nll += F.l1_loss(logits.view(-1), y.view(-1))
            penalty += self._irm_penalty(logits, y)
        nll /= len(minibatches)
        penalty /= len(minibatches)
        loss = nll + (penalty_weight * penalty)

        if self.update_count == self.hparams['irm_penalty_anneal_iters']:
            # Reset Adam, because it doesn't like the sharp jump in gradient
            # magnitudes that happens at this step.
            self.optimizer = torch.optim.Adam(
                self.network.parameters(),
                lr=self.hparams["lr"],
                weight_decay=self.hparams['weight_decay'])

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        self.update_count += 1
        return {'loss': loss.item(), 'nll': nll.item(),
            'penalty': penalty.item()}


class VREx(ERM):
    """V-REx algorithm from http://arxiv.org/abs/2003.00688"""
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(VREx, self).__init__(input_shape, num_classes, num_domains,
                                  hparams)
        self.register_buffer('update_count', torch.tensor([0]))

    def update(self, minibatches, unlabeled=None):
        if self.update_count >= self.hparams["vrex_penalty_anneal_iters"]:
            penalty_weight = self.hparams["vrex_lambda"]
        else:
            penalty_weight = 1.0

        nll = 0.

        all_x = torch.cat([x for x, y in minibatches])
        all_logits = self.network(all_x)
        all_logits_idx = 0
        losses = torch.zeros(len(minibatches))
        for i, (x, y) in enumerate(minibatches):
            logits = all_logits[all_logits_idx:all_logits_idx + x.shape[0]]
            all_logits_idx += x.shape[0]
            nll = F.l1_loss(logits.view(-1), y.view(-1))
            losses[i] = nll

        mean = losses.mean()
        penalty = ((losses - mean) ** 2).mean()
        loss = mean + penalty_weight * penalty

        if self.update_count == self.hparams['vrex_penalty_anneal_iters']:
            # Reset Adam (like IRM), because it doesn't like the sharp jump in
            # gradient magnitudes that happens at this step.
            self.optimizer = torch.optim.Adam(
                self.network.parameters(),
                lr=self.hparams["lr"],
                weight_decay=self.hparams['weight_decay'])

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        self.update_count += 1
        return {'loss': loss.item(), 'nll': nll.item(),
                'penalty': penalty.item()}


class Mixup(ERM):
    """
    Mixup of minibatches from different domains
    https://arxiv.org/pdf/2001.00677.pdf
    https://arxiv.org/pdf/1912.01805.pdf
    """
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(Mixup, self).__init__(input_shape, num_classes, num_domains,
                                    hparams)

    def update(self, minibatches, unlabeled=None):
        objective = 0

        for (xi, yi), (xj, yj) in random_pairs_of_minibatches(minibatches):
            lam = np.random.beta(self.hparams["mixup_alpha"],
                                 self.hparams["mixup_alpha"])

            x = lam * xi + (1 - lam) * xj
            predictions = self.predict(x)

            objective += lam * F.l1_loss(predictions.view(-1), yi.view(-1))
            objective += (1 - lam) * F.l1_loss(predictions.view(-1), yj.view(-1))

        objective /= len(minibatches)

        self.optimizer.zero_grad()
        objective.backward()
        self.optimizer.step()

        return {'loss': objective.item()}


class GroupDRO(ERM):
    """
    Robust ERM minimizes the error at the worst minibatch
    Algorithm 1 from [https://arxiv.org/pdf/1911.08731.pdf]
    """
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(GroupDRO, self).__init__(input_shape, num_classes, num_domains,
                                        hparams)
        self.register_buffer("q", torch.Tensor())

    def update(self, minibatches, unlabeled=None):
        device = "cuda" if minibatches[0][0].is_cuda else "cpu"

        if not len(self.q):
            self.q = torch.ones(len(minibatches)).to(device)

        losses = torch.zeros(len(minibatches)).to(device)

        for m in range(len(minibatches)):
            x, y = minibatches[m]
            losses[m] = F.l1_loss(self.predict(x).view(-1), y.view(-1))
            self.q[m] *= (self.hparams["groupdro_eta"] * losses[m].data).exp()

        self.q /= self.q.sum()

        loss = torch.dot(losses, self.q)

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return {'loss': loss.item()}

class AbstractMMD(ERM):
    """
    Perform ERM while matching the pair-wise domain feature distributions
    using MMD (abstract class)
    """
    def __init__(self, input_shape, num_classes, num_domains, hparams, gaussian, device='cpu'):
        super(AbstractMMD, self).__init__(input_shape, num_classes, num_domains,
                                  hparams)
        if gaussian:
            self.kernel_type = "gaussian"
        else:
            self.kernel_type = "mean_cov"

    def my_cdist(self, x1, x2):
        x1_norm = x1.pow(2).sum(dim=-1, keepdim=True)
        x2_norm = x2.pow(2).sum(dim=-1, keepdim=True)
        res = torch.addmm(x2_norm.transpose(-2, -1),
                          x1,
                          x2.transpose(-2, -1), alpha=-2).add_(x1_norm)
        return res.clamp_min_(1e-30)

    def gaussian_kernel(self, x, y, gamma=[0.001, 0.01, 0.1, 1, 10, 100,
                                           1000]):
        D = self.my_cdist(x, y)
        K = torch.zeros_like(D)

        for g in gamma:
            K.add_(torch.exp(D.mul(-g)))

        return K

    def mmd(self, x, y):
        if self.kernel_type == "gaussian":
            Kxx = self.gaussian_kernel(x, x).mean()
            Kyy = self.gaussian_kernel(y, y).mean()
            Kxy = self.gaussian_kernel(x, y).mean()
            return Kxx + Kyy - 2 * Kxy
        else:
            mean_x = x.mean(0, keepdim=True)
            mean_y = y.mean(0, keepdim=True)
            cent_x = x - mean_x
            cent_y = y - mean_y
            cova_x = (cent_x.t() @ cent_x) / (len(x) - 1)
            cova_y = (cent_y.t() @ cent_y) / (len(y) - 1)

            mean_diff = (mean_x - mean_y).pow(2).mean()
            cova_diff = (cova_x - cova_y).pow(2).mean()

            return mean_diff + cova_diff

    def update(self, minibatches, unlabeled=None):
        objective = 0
        penalty = 0
        nmb = len(minibatches)

        features = [self.featurizer(xi) for xi, _ in minibatches]
        classifs = [self.classifier(fi) for fi in features]
        targets = [yi for _, yi in minibatches]

        for i in range(nmb):
            objective += F.l1_loss(classifs[i].view(-1), targets[i].view(-1))
            for j in range(i + 1, nmb):
                penalty += self.mmd(features[i], features[j])

        objective /= nmb
        if nmb > 1:
            penalty /= (nmb * (nmb - 1) / 2)

        self.optimizer.zero_grad()
        (objective + (self.hparams['mmd_gamma']*penalty)).backward()
        self.optimizer.step()

        if torch.is_tensor(penalty):
            penalty = penalty.item()

        return {'loss': objective.item(), 'penalty': penalty}


class MMD(AbstractMMD):
    """
    MMD using Gaussian kernel
    """

    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(MMD, self).__init__(input_shape, num_classes,
                                          num_domains, hparams, gaussian=True)

class MLDG(ERM):
    """
    Model-Agnostic Meta-Learning
    Algorithm 1 / Equation (3) from: https://arxiv.org/pdf/1710.03463.pdf
    Related: https://arxiv.org/pdf/1703.03400.pdf
    Related: https://arxiv.org/pdf/1910.13580.pdf
    """
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(MLDG, self).__init__(input_shape, num_classes, num_domains,
                                   hparams)

    # def update(self, minibatches, unlabeled=None):
    #     """
    #     Terms being computed:
    #         * Li = Loss(xi, yi, params)
    #         * Gi = Grad(Li, params)
    #         * Lj = Loss(xj, yj, Optimizer(params, grad(Li, params)))
    #         * Gj = Grad(Lj, params)
    #         * params = Optimizer(params, Grad(Li + beta * Lj, params))
    #         *        = Optimizer(params, Gi + beta * Gj)
    #     That is, when calling .step(), we want grads to be Gi + beta * Gj
    #     For computational efficiency, we do not compute second derivatives.
    #     """
    #     num_mb = len(minibatches)
    #     objective = 0

    #     self.optimizer.zero_grad()
    #     for p in self.network.parameters():
    #         if p.grad is None:
    #             p.grad = torch.zeros_like(p)

    #     for (xi, yi), (xj, yj) in random_pairs_of_minibatches(minibatches):
    #         # fine tune clone-network on task "i"
    #         inner_net = copy.deepcopy(self.network)

    #         inner_opt = torch.optim.Adam(
    #             inner_net.parameters(),
    #             lr=self.hparams["lr"],
    #             weight_decay=self.hparams['weight_decay']
    #         )

    #         inner_obj = F.l1_loss(inner_net(xi).view(-1), yi.view(-1))

    #         inner_opt.zero_grad()
    #         inner_obj.backward()
    #         inner_opt.step()

    #         # The network has now accumulated gradients Gi
    #         # The clone-network has now parameters P - lr * Gi
    #         for p_tgt, p_src in zip(self.network.parameters(),
    #                                 inner_net.parameters()):
    #             if p_src.grad is not None:
    #                 p_tgt.grad.data.add_(p_src.grad.data / num_mb)

    #         # `objective` is populated for reporting purposes
    #         objective += inner_obj.item()

    #         # this computes Gj on the clone-network
    #         loss_inner_j = F.l1_loss(inner_net(xj).view(-1), yj.view(-1))
    #         grad_inner_j = autograd.grad(loss_inner_j, inner_net.parameters(),
    #             allow_unused=True)

    #         # `objective` is populated for reporting purposes
    #         objective += (self.hparams['mldg_beta'] * loss_inner_j).item()

    #         for p, g_j in zip(self.network.parameters(), grad_inner_j):
    #             if g_j is not None:
    #                 p.grad.data.add_(
    #                     self.hparams['mldg_beta'] * g_j.data / num_mb)

    #         # The network has now accumulated gradients Gi + beta * Gj
    #         # Repeat for all train-test splits, do .step()

    #     objective /= len(minibatches)

    #     self.optimizer.step()

    #     return {'loss': objective}

    # This commented "update" method back-propagates through the gradients of
    # the inner update, as suggested in the original MAML paper.  However, this
    # is twice as expensive as the uncommented "update" method, which does not
    # compute second-order derivatives, implementing the First-Order MAML
    # method (FOMAML) described in the original MAML paper.

    def update(self, minibatches, unlabeled=None):
        objective = 0
        beta = self.hparams["beta"]
        inner_iterations = self.hparams["inner_iterations"]

        self.optimizer.zero_grad()

        with higher.innerloop_ctx(self.network, self.optimizer,
            copy_initial_weights=False) as (inner_network, inner_optimizer):

            for (xi, yi), (xj, yj) in random_pairs_of_minibatches(minibatches):
                for inner_iteration in range(inner_iterations):
                    li = F.l1_loss(inner_network(xi).view(-1), yi.view(-1))
                    inner_optimizer.step(li)
    
                objective += F.l1_loss(self.network(xi).view(-1), yi.view(-1))
                objective += beta * F.l1_loss(inner_network(xj).view(-1), yj.view(-1))

            objective /= len(minibatches)
            objective.backward()
    
        self.optimizer.step()
        return {'loss': objective.item()}
        
class MAML(ERM):
    """
    Model-Agnostic Meta-Learning
    Algorithm 1 / Equation (3) from: https://arxiv.org/pdf/1710.03463.pdf
    Related: https://arxiv.org/pdf/1703.03400.pdf
    Related: https://arxiv.org/pdf/1910.13580.pdf
    """
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(MAML, self).__init__(input_shape, num_classes, num_domains,
                                   hparams)

    # def update(self, minibatches, unlabeled=None):
    #     """
    #     Terms being computed:
    #         * Li = Loss(xi, yi, params)
    #         * Gi = Grad(Li, params)
    #         * Lj = Loss(xj, yj, Optimizer(params, grad(Li, params)))
    #         * Gj = Grad(Lj, params)
    #         * params = Optimizer(params, Grad(Li + beta * Lj, params))
    #         *        = Optimizer(params, Gi + beta * Gj)
    #     That is, when calling .step(), we want grads to be Gi + beta * Gj
    #     For computational efficiency, we do not compute second derivatives.
    #     """
    #     num_mb = len(minibatches)
    #     objective = 0

    #     self.optimizer.zero_grad()
    #     for p in self.network.parameters():
    #         if p.grad is None:
    #             p.grad = torch.zeros_like(p)

    #     for (xi, yi), (xj, yj) in random_pairs_of_minibatches(minibatches):
    #         # fine tune clone-network on task "i"
    #         inner_net = copy.deepcopy(self.network)

    #         inner_opt = torch.optim.Adam(
    #             inner_net.parameters(),
    #             lr=self.hparams["lr"],
    #             weight_decay=self.hparams['weight_decay']
    #         )

    #         inner_obj = F.l1_loss(inner_net(xi).view(-1), yi.view(-1))

    #         inner_opt.zero_grad()
    #         inner_obj.backward()
    #         inner_opt.step()

    #         # The network has now accumulated gradients Gi
    #         # The clone-network has now parameters P - lr * Gi
    #         for p_tgt, p_src in zip(self.network.parameters(),
    #                                 inner_net.parameters()):
    #             if p_src.grad is not None:
    #                 p_tgt.grad.data.add_(p_src.grad.data / num_mb)

    #         # `objective` is populated for reporting purposes
    #         objective += inner_obj.item()

    #         if self.hparams['weight'] =="none":

    #             loss_inner_j = F.l1_loss(inner_net(xj).view(-1), yj.view(-1))

    #         elif self.hparams['weight'] =="pair":
    #             with torch.no_grad():
    #                 support_features=self.featurizer(xi)
    #                 query_features=self.featurizer(xj)
    #                 weights=torch.mean((support_features-query_features)**2,dim=-1)
    #                 weights.detach()

    #             loss_inner_j=torch.mean((inner_net(xj).view(-1)-yj.view(-1))**2 *weights)
       

    #         # this computes Gj on the clone-network
            

    #         grad_inner_j = autograd.grad(loss_inner_j, inner_net.parameters(),
    #             allow_unused=True)

    #         # `objective` is populated for reporting purposes
    #         objective += (self.hparams['mldg_beta'] * loss_inner_j).item()

    #         for p, g_j in zip(self.network.parameters(), grad_inner_j):
    #             if g_j is not None:
    #                 p.grad.data.add_(
    #                     self.hparams['mldg_beta'] * g_j.data / num_mb)

    #         # The network has now accumulated gradients Gi + beta * Gj
    #         # Repeat for all train-test splits, do .step()

    #     objective /= len(minibatches)

    #     self.optimizer.step()

    #     return {'loss': objective}

    # This commented "update" method back-propagates through the gradients of
    # the inner update, as suggested in the original MAML paper.  However, this
    # is twice as expensive as the uncommented "update" method, which does not
    # compute second-order derivatives, implementing the First-Order MAML
    # method (FOMAML) described in the original MAML paper.

    def update(self, minibatches, unlabeled=None):
        objective = 0
        beta = self.hparams["beta"]
        inner_iterations = self.hparams["inner_iterations"]
        # 
        self.optimizer.zero_grad()
        
        with higher.innerloop_ctx(self.network, self.optimizer,
            copy_initial_weights=False) as (inner_network, inner_optimizer):

            for (xi, yi), (xj, yj) in random_pairs_of_minibatches(minibatches):
                # print(xi.shape,xj.shape)
                for inner_iteration in range(inner_iterations):
                    li = F.l1_loss(inner_network(xi).view(-1), yi.view(-1))
                    inner_optimizer.step(li)
    
                # objective += F.l1_loss(self.network(xi), yi.view(-1))
                # if self.hparams['weight'] =="none":
                loss_inner_j = F.l1_loss(inner_network(xj).view(-1), yj.view(-1))
                # elif self.hparams['weight'] =="pair":
                #     with torch.no_grad():
                #         support_features=self.featurizer(xi)
                #         query_features=self.featurizer(xj)
                #         weights=torch.mean((support_features-query_features)**2,dim=-1)
                #         weights.detach()

                    # loss_inner_j=torch.mean((inner_network(xj).view(-1)-yj.view(-1))**2 *weights)
                # objective += beta * F.l1_loss(inner_network(xj).view(-1), yj.view(-1))
                objective += beta * loss_inner_j

            objective /= len(minibatches)
            objective.backward()
    
        self.optimizer.step()
    
        return {'loss': objective.item()}

class MAMR(ERM):
    """
    Model-Agnostic Meta-Learning
    Algorithm 1 / Equation (3) from: https://arxiv.org/pdf/1710.03463.pdf
    Related: https://arxiv.org/pdf/1703.03400.pdf
    Related: https://arxiv.org/pdf/1910.13580.pdf
    """
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(MAMR, self).__init__(input_shape, num_classes, num_domains,
                                   hparams)
        self.cnt=0


    def update(self, minibatches, unlabeled=None):
        self.cnt+=1
        objective = 0
        beta = self.hparams["beta"]
        inner_iterations = self.hparams["inner_iterations"]

        self.optimizer.zero_grad()

        with higher.innerloop_ctx(self.network, self.optimizer,
            copy_initial_weights=False) as (inner_network, inner_optimizer):

            for (xi, yi), (xj, yj) in random_pairs_of_minibatches(minibatches):
                for inner_iteration in range(inner_iterations):
                    li = F.l1_loss(inner_network(xi).view(-1), yi.view(-1))
                    inner_optimizer.step(li)
    
                # objective += F.l1_loss(self.network(xi), yi.view(-1))
                # if self.hparams['weight'] =="none":
                #     loss_inner_j = F.l1_loss(inner_network(xj).view(-1), yj.view(-1))
                # elif self.hparams['weight'] =="pair":
                with torch.no_grad():
                    support_features=self.featurizer(xi)
                    query_features=self.featurizer(xj)
                    # weights=torch.mean(((support_features-query_features)/(torch.abs(support_features) + 1e-8)),dim=-1)
                    weights=torch.mean(((support_features-query_features)**2),dim=-1,keepdim=True)
                    
                    # print(weights)
                    weights=F.normalize(weights, dim=0)
                    # print(weights)
                    weights=weights.detach().view(-1)
                    # weights=torch.mean(weights)
                    # print
                    weights=torch.clamp_min(weights, min=1e-2)
                    # print(weights)

                # loss_inner_j= torch.mean(((inner_network(xj).view(-1) - yj.view(-1))**2 )* weights)
                loss_inner_j= torch.mean(torch.abs(((inner_network(xj).view(-1) - yj.view(-1))))* weights)
                # objective += beta * F.l1_loss(inner_network(xj).view(-1), yj.view(-1))
                objective += beta * loss_inner_j

            objective /= len(minibatches)
            objective.backward()
        self.optimizer.step()
        return {'loss': objective.item()}

class RSD(ERM):
    """
    RSD, Representation subspace distance, Xinyang Chen, et al. ICML 2021, 

    """
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(RSD, self).__init__(input_shape, num_classes, num_domains,
                                   hparams)

        self.tradeoff2=hparams["tradeoff2"]
        self.tradeoff=hparams["tradeoff"]

    def rsd(self, Feature_s, Feature_t):
        try:
            u_s, s_s, v_s = torch.svd(Feature_s.t())
        except:
            u_s, s_s, v_s = torch.svd(Feature_s.t()+1e-5*Feature_s.t().mean()*torch.rand_like(Feature_s.t()))

        try:
            u_t, s_t, v_t =torch.svd(Feature_t.t())
        except:
            u_t, s_t, v_t =torch.svd(Feature_t.t()+1e-5*Feature_t.t().mean()*torch.rand_like(Feature_t.t()))

        v_s,v_t=v_s.t(),v_t.t()
        p_s, cospa, p_t = torch.svd(torch.mm(u_s.t(), u_t))
        sinpa = torch.sqrt(1-torch.pow(cospa,2))
        return torch.norm(sinpa,1)+ self.tradeoff2*torch.norm(torch.abs(p_s) - torch.abs(p_t), 2)

    def update(self, minibatches, unlabeled=None):

        objective = 0
        for (xi, yi), (xj, yj) in random_pairs_of_minibatches(minibatches):
            
            feature_i, feature_j = self.featurizer(xi), self.featurizer(xj)
            rsd_loss=self.rsd(feature_i,feature_j)
            l1 = F.l1_loss(self.network(xi).view(-1), yi.view(-1))
            l2 = F.l1_loss(self.network(xj).view(-1), yj.view(-1))

            objective=rsd_loss * self.tradeoff + l1 + l2 

            self.optimizer.zero_grad()
            objective.backward()
            self.optimizer.step()

        return {'loss': objective.item()}



class KMAML(AbstractMMD):
    """
    Model-Agnostic Meta-Learning
    Algorithm 1 / Equation (3) from: https://arxiv.org/pdf/1710.03463.pdf
    Related: https://arxiv.org/pdf/1703.03400.pdf
    Related: https://arxiv.org/pdf/1910.13580.pdf
    """
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(KMAML, self).__init__(input_shape, num_classes, num_domains,
                                   hparams, gaussian=True)
        # self.cnt=0
    def update(self, minibatches, unlabeled=None):
        # self.cnt+=1

        objective = 0
        beta = self.hparams["beta"]
        inner_iterations = self.hparams["inner_iterations"]

        self.optimizer.zero_grad()

        with higher.innerloop_ctx(self.network, self.optimizer,
            copy_initial_weights=False) as (inner_network, inner_optimizer):

            for (xi, yi), (xj, yj) in random_pairs_of_minibatches(minibatches):
                for inner_iteration in range(inner_iterations):
                    li = F.l1_loss(inner_network(xi).view(-1), yi.view(-1))
                    inner_optimizer.step(li)
    
                # objective += F.l1_loss(self.network(xi), yi.view(-1))
                # if self.hparams['weight'] =="none":
                #     loss_inner_j = F.l1_loss(inner_network(xj).view(-1), yj.view(-1))
                # elif self.hparams['weight'] =="pair":
                with torch.no_grad():
                    support_features=self.featurizer(xi)
                    query_features=self.featurizer(xj)
                    # weights=torch.mean(((support_features-query_features)/(torch.abs(support_features) + 1e-8)),dim=-1)
                    weights=self.mmd(support_features,query_features)

                    # weights=torch.mean(((support_features-query_features)**2),dim=-1,keepdim=True)
                    
                    # print(weights)
                    # weights=F.normalize(weights, dim=0)
                    # print(weights)
                    weights=weights.detach().view(-1)
                    # weights=torch.mean(weights)
                    
                    # weights=torch.clamp_min(weights, min=1e-5)
                # print(weights)

                loss_inner_j=torch.mean(((inner_network(xj).view(-1) - yj.view(-1))**2 )* weights)
                # objective += beta * F.l1_loss(inner_network(xj).view(-1), yj.view(-1))
                objective += beta * loss_inner_j

            objective /= len(minibatches)
            objective.backward()
    
        self.optimizer.step()
    
        return {'loss': objective.item()}

class CORAL(AbstractMMD):
    """
    MMD using mean and covariance difference
    """

    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(CORAL, self).__init__(input_shape, num_classes,
                                         num_domains, hparams, gaussian=False)


class MTL(Algorithm):
    """
    A neural network version of
    Domain Generalization by Marginal Transfer Learning
    (https://arxiv.org/abs/1711.07910)
    """

    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(MTL, self).__init__(input_shape, num_classes, num_domains,
                                  hparams)
        self.featurizer = networks.Featurizer(input_shape, self.hparams)
        self.classifier = networks.Classifier(
            self.featurizer.n_outputs * 2,
            num_classes,
            self.hparams['nonlinear_classifier'])
        self.optimizer = torch.optim.Adam(
            list(self.featurizer.parameters()) +\
            list(self.classifier.parameters()),
            lr=self.hparams["lr"],
            weight_decay=self.hparams['weight_decay']
        )

        self.register_buffer('embeddings',
                             torch.zeros(num_domains,
                                         self.featurizer.n_outputs))

        self.ema = self.hparams['mtl_ema']

    def update(self, minibatches, unlabeled=None):
        loss = 0
        for env, (x, y) in enumerate(minibatches):
            loss += F.l1_loss(self.predict(x, env).view(-1), y.view(-1))

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return {'loss': loss.item()}

    def update_embeddings_(self, features, env=None):
        return_embedding = features.mean(0)

        if env is not None:
            return_embedding = self.ema * return_embedding +\
                               (1 - self.ema) * self.embeddings[env]

            self.embeddings[env] = return_embedding.clone().detach()

        return return_embedding.view(1, -1).repeat(len(features), 1)

    def predict(self, x, env=None):
        features = self.featurizer(x)
        embedding = self.update_embeddings_(features, env).normal_()
        return self.classifier(torch.cat((features, embedding), 1))

class SagNet(Algorithm):
    """
    Style Agnostic Network
    Algorithm 1 from: https://arxiv.org/abs/1910.11645
    """

    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(SagNet, self).__init__(input_shape, num_classes, num_domains,
                                  hparams)
        # featurizer network
        self.network_f = networks.Featurizer(input_shape, self.hparams)
        # content network
        self.network_c = networks.Classifier(
            self.network_f.n_outputs,
            num_classes,
            self.hparams['nonlinear_classifier'])
        # style network
        self.network_s = networks.Classifier(
            self.network_f.n_outputs,
            num_classes,
            self.hparams['nonlinear_classifier'])

        def opt(p):
            return torch.optim.Adam(p, lr=hparams["lr"],
                    weight_decay=hparams["weight_decay"])

        self.optimizer_f = opt(self.network_f.parameters())
        self.optimizer_c = opt(self.network_c.parameters())
        self.optimizer_s = opt(self.network_s.parameters())
        self.weight_adv = hparams["sag_w_adv"]

    def forward_c(self, x):
        # learning content network on randomized style
        return self.network_c(self.randomize(self.network_f(x), "style"))

    def forward_s(self, x):
        # learning style network on randomized content
        return self.network_s(self.randomize(self.network_f(x), "content"))

    def randomize(self, x, what="style", eps=1e-5):
        device = "cuda" if x.is_cuda else "cpu"
        sizes = x.size()
        alpha = torch.rand(sizes[0], 1).to(device)

        if len(sizes) == 4:
            x = x.view(sizes[0], sizes[1], -1)
            alpha = alpha.unsqueeze(-1)

        mean = x.mean(-1, keepdim=True)
        var = x.var(-1, keepdim=True)

        x = (x - mean) / (var + eps).sqrt()

        idx_swap = torch.randperm(sizes[0])
        if what == "style":
            mean = alpha * mean + (1 - alpha) * mean[idx_swap]
            var = alpha * var + (1 - alpha) * var[idx_swap]
        else:
            x = x[idx_swap].detach()

        x = x * (var + eps).sqrt() + mean
        return x.view(*sizes)

    def update(self, minibatches, unlabeled=None):
        all_x = torch.cat([x for x, y in minibatches])
        all_y = torch.cat([y for x, y in minibatches])

        # learn content
        self.optimizer_f.zero_grad()
        self.optimizer_c.zero_grad()
        loss_c = F.l1_loss(self.forward_c(all_x).view(-1), all_y.view(-1))
        loss_c.backward()
        self.optimizer_f.step()
        self.optimizer_c.step()

        # learn style
        self.optimizer_s.zero_grad()
        loss_s = F.l1_loss(self.forward_s(all_x).view(-1), all_y.view(-1))
        loss_s.backward()
        self.optimizer_s.step()

        # learn adversary
        self.optimizer_f.zero_grad()
        loss_adv = -F.log_softmax(self.forward_s(all_x), dim=1).mean(1).mean()
        loss_adv = loss_adv * self.weight_adv
        loss_adv.backward()
        self.optimizer_f.step()

        return {'loss_c': loss_c.item(), 'loss_s': loss_s.item(),
                'loss_adv': loss_adv.item()}

    def predict(self, x):
        return self.network_c(self.network_f(x))


class RSC(ERM):
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(RSC, self).__init__(input_shape, num_classes, num_domains,
                                   hparams)
        self.drop_f = (1 - hparams['rsc_f_drop_factor']) * 100
        self.drop_b = (1 - hparams['rsc_b_drop_factor']) * 100
        self.num_classes = num_classes

    def update(self, minibatches, unlabeled=None):
        device = "cuda" if minibatches[0][0].is_cuda else "cpu"

        # inputs
        all_x = torch.cat([x for x, y in minibatches])
        # labels
        all_y = torch.cat([y for _, y in minibatches])
        # one-hot labels
        all_o = torch.nn.functional.one_hot(all_y, self.num_classes)
        # features
        all_f = self.featurizer(all_x)
        # predictions
        all_p = self.classifier(all_f)

        # Equation (1): compute gradients with respect to representation
        all_g = autograd.grad((all_p * all_o).sum(), all_f)[0]

        # Equation (2): compute top-gradient-percentile mask
        percentiles = np.percentile(all_g.cpu(), self.drop_f, axis=1)
        percentiles = torch.Tensor(percentiles)
        percentiles = percentiles.unsqueeze(1).repeat(1, all_g.size(1))
        mask_f = all_g.lt(percentiles.to(device)).float()

        # Equation (3): mute top-gradient-percentile activations
        all_f_muted = all_f * mask_f

        # Equation (4): compute muted predictions
        all_p_muted = self.classifier(all_f_muted)

        # Section 3.3: Batch Percentage
        all_s = F.softmax(all_p, dim=1)
        all_s_muted = F.softmax(all_p_muted, dim=1)
        changes = (all_s * all_o).sum(1) - (all_s_muted * all_o).sum(1)
        percentile = np.percentile(changes.detach().cpu(), self.drop_b)
        mask_b = changes.lt(percentile).float().view(-1, 1)
        mask = torch.logical_or(mask_f, mask_b).float()

        # Equations (3) and (4) again, this time mutting over examples
        all_p_muted_again = self.classifier(all_f * mask)

        # Equation (5): update
        loss = F.l1_loss(all_p_muted_again.view(-1), all_y.view(-1))
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return {'loss': loss.item()}


class SD(ERM):
    """
    Gradient Starvation: A Learning Proclivity in Neural Networks
    Equation 25 from [https://arxiv.org/pdf/2011.09468.pdf]
    """
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(SD, self).__init__(input_shape, num_classes, num_domains,
                                        hparams)
        self.sd_reg = hparams["sd_reg"]

    def update(self, minibatches, unlabeled=None):
        all_x = torch.cat([x for x,y in minibatches])
        all_y = torch.cat([y for x,y in minibatches])
        all_p = self.predict(all_x)

        loss = F.l1_loss(all_p.view(-1), all_y.view(-1))
        penalty = (all_p ** 2).mean()
        objective = loss + self.sd_reg * penalty

        self.optimizer.zero_grad()
        objective.backward()
        self.optimizer.step()

        return {'loss': loss.item(), 'penalty': penalty.item()}


class DDG_AugMix(ERM):

    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(DDG_AugMix, self).__init__(input_shape, num_classes, num_domains,
                                  hparams, device)
        self.hparams = hparams
        self.class_balance=hparams['class_balanced']
        self.iteration = 0
        self.id_featurizer = self.featurizer
        self.dis_id = self.classifier
        self.gen = networks.AdaINGen(1, self.id_featurizer.n_outputs, hparams) if not hparams['is_mnist'] else networks.VAEGen()
        self.dis_img = networks.MsImageDis(hparams=hparams) 
        
        def to_gray(half=False): #simple
            def forward(x):
                x = torch.mean(x, dim=1, keepdim=True)
                if half:
                    x = x.half()
                return x
            return forward
        self.single = to_gray(False)
        self.optimizer_gen = torch.optim.Adam([p for p in list(self.gen.parameters())  if p.requires_grad], lr=self.hparams['lr_g'], betas=(0, 0.999), weight_decay=self.hparams['weight_decay_g'])
        self.jsd = True

        self.id_criterion = nn.CrossEntropyLoss()
        

    def recon_criterion(self, input, target):
            diff = input - target.detach()
            return torch.mean(torch.abs(diff[:]))
    
    def train_bn(self, m):
        classname = m.__class__.__name__
        if classname.find('BatchNorm') != -1:
            m.train()
            print('there has bn')

    def forward(self, x_a, x_b, xp_a, xp_b):
        '''
            inpus:
                x_a, x_b: image from dataloader a,b
                xp_a, xp_b: positive pair of x_a, x_b
        '''
        b = xp_a.shape[0]
        s_a = self.gen.encode(self.single(x_a))# v for x_a
        s_b = self.gen.encode(self.single(x_b))# v for x_b
        f_a, x_fa = self.id_featurizer(x_a) # f_a: detached s for x_a, x_fa: s for x_a
        p_a = self.dis_id(x_fa)             # identity classification result for x_a
        f_b, x_fb = self.id_featurizer(x_b)
        p_b = self.dis_id(x_fb)
        fp_a, xp_fa = self.id_featurizer(xp_a)
        pp_a = self.dis_id(xp_fa)
        fp_b, xp_fb = self.id_featurizer(xp_b)
        pp_b = self.dis_id(xp_fb)
        # self-reconstruction
        x_a_recon = self.gen.decode(s_a[:b], f_a[:b]) # generate from identity and style of a
        x_b_recon = self.gen.decode(s_b[:b], f_b[:b])
        # style-preserving reconstructi
        x_a_recon_p = self.gen.decode(s_a[:b], fp_a) # generate from identity of p_a and styld of a
        x_b_recon_p = self.gen.decode(s_b[:b], fp_b)

        return x_fa, x_fb, p_a, p_b, pp_a, pp_b, x_a_recon, x_b_recon, x_a_recon_p, x_b_recon_p    
 

    def gen_update(self, xf_a, xf_b, p_a, p_b, pp_a, pp_b, x_a_recon, x_b_recon, x_a_recon_p, x_b_recon_p, x_a, x_b,  l_a, l_b, hparams):
        '''
            inputs:
                x_ab: generated from identity of b and style of a
                x_ba: generated from identity of a and style of b
                s_a, s_b: style factors for x_a, x_b
                f_a, f_b: detached semantic factors for x_a, x_b
                p_a, p_b: identity prediction results for x_a, x_b
                pp_a, pp_b: identity prediction results for the positive pair of x_a, x_b
                x_a_recon, x_b_recon: reconstruction of x_a, x_b
                x_a_recon_p, x_b_recon_p: reconstruction of the positive pair of x_a, x_b
                x_a, x_b,  l_a, l_b: images and identity labels
                hparams: parameters
        '''
        b = x_a_recon.shape[0]
        self.optimizer_gen.zero_grad()
        self.optimizer.zero_grad()

        #################################
        # auto-encoder image reconstruction
        self.loss_gen_recon_x = self.recon_criterion(x_a_recon, x_a[:b])+self.recon_criterion(x_b_recon, x_b[:b]) + self.recon_criterion(x_a_recon_p, x_a[:b])+ self.recon_criterion(x_b_recon_p, x_b[:b])

        # Emprical Loss
        self.loss_id = self.id_criterion(p_a[:b], l_a[:b]) + self.id_criterion(p_b[:b], l_b[:b]) +  self.id_criterion(pp_a, l_a[:b]) + self.id_criterion(pp_b, l_b[:b])
        self.loss_aug_id =  self.id_criterion(p_a[b:], l_a[b:]) + self.id_criterion(p_b[b:], l_b[b:])

        if self.jsd:
            logits_clean, logits_aug1, logits_aug2 = torch.split(xf_a, b)
            logits_clean_b, logits_aug1_b, logits_aug2_b = torch.split(xf_b, b)
            p_clean, p_aug1, p_aug2 = F.softmax(
                logits_clean, dim=1), F.softmax(
                    logits_aug1, dim=1), F.softmax(
                        logits_aug2, dim=1)
            p_clean_b, p_aug1_b, p_aug2_b = F.softmax(
                logits_clean_b, dim=1), F.softmax(
                    logits_aug1_b, dim=1), F.softmax(
                        logits_aug2_b, dim=1)

            # Clamp mixture distribution to avoid exploding KL divergence
            p_mixture = torch.clamp((p_clean + p_aug1 + p_aug2) / 3., 1e-7, 1).log()
            p_mixture_b = torch.clamp((p_clean_b + p_aug1_b + p_aug2_b) / 3., 1e-7, 1).log()
            self.jsd_loss = 12 * (F.kl_div(p_mixture, p_clean, reduction='batchmean') +
                            F.kl_div(p_mixture, p_aug1, reduction='batchmean') +
                            F.kl_div(p_mixture, p_aug2, reduction='batchmean') +
                            F.kl_div(p_mixture_b, p_clean_b, reduction='batchmean') +
                            F.kl_div(p_mixture_b, p_aug1_b, reduction='batchmean') +
                            F.kl_div(p_mixture_b, p_aug2_b, reduction='batchmean')) / 3.
        else:
            self.jsd_loss = torch.tensor(0)
        # total loss
        self.loss_gen_total = self.loss_id + \
                    hparams['recon_x_w'] * self.loss_gen_recon_x + \
                    hparams['recon_id_w'] * self.loss_aug_id +\
                    self.jsd_loss

        self.loss_gen_total.backward()
        self.optimizer_gen.step()
        self.optimizer.step()

    def update(self, minibatches, minibatches_neg, pretrain_model=None, unlabeled=None, iteration=0):
        images_a = torch.cat([x for x, y, pos in minibatches])
        labels_a = torch.cat([y for x, y, pos in minibatches])
        pos_a_ = torch.cat([pos for x, y, pos in minibatches])
        images_b = torch.cat([x for x, y, pos in minibatches_neg])
        labels_b = torch.cat([y for x, y, pos in minibatches_neg])
        pos_b_ = torch.cat([pos for x, y, pos in minibatches_neg])
        device = "cuda" if minibatches[0][0].is_cuda else "cpu" 
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        self.preprocess = transforms.Compose(
            [transforms.ToTensor(),
            transforms.Normalize(mean, std)])
        self.TO_pil = transforms.ToPILImage()
        if self.jsd:
            for image_a, image_b, label_a, label_b in zip(images_a, images_b, labels_a, labels_b):
                image_a, x_ba1, x_ba2= Augmix(self.TO_pil(image_b.cpu()), self.preprocess, no_jsd=not self.jsd)
                image_b, x_ab1, x_ab2= Augmix(self.TO_pil(image_a.cpu()), self.preprocess, no_jsd=not self.jsd)
                images_a = torch.cat([images_a, x_ba1.to(device).unsqueeze(0), x_ba2.to(device).unsqueeze(0)], dim=0)
                images_b = torch.cat([images_b, x_ab1.to(device).unsqueeze(0), x_ab2.to(device).unsqueeze(0)], dim=0)
                labels_a = torch.cat([labels_a, label_a.unsqueeze(0), label_a.unsqueeze(0)], 0)
                labels_b = torch.cat([labels_b, label_b.unsqueeze(0), label_b.unsqueeze(0)], 0)
        else:
            images_b = Augmix(images_b, self.preprocess, no_jsd=not self.jsd)
            images_a = Augmix(images_a, self.preprocess, no_jsd=not self.jsd)
        
        pos_a, pos_b = [], []
        for image_a, image_b in zip(pos_a_, pos_b_):
            x_pos = Augmix(self.TO_pil(image_b.cpu()), self.preprocess, no_jsd=True)
            x_pos_b = Augmix(self.TO_pil(image_a.cpu()), self.preprocess, no_jsd=True)
            pos_a.append(x_pos.to(device).unsqueeze(0))
            pos_b.append(x_pos_b.to(device).unsqueeze(0))
    
        pos_a, pos_b =torch.cat(pos_a), torch.cat(pos_b)
        xf_a, xf_b, p_a, p_b, pp_a, pp_b, x_a_recon, x_b_recon, x_a_recon_p, x_b_recon_p = self.forward(images_a, images_b, pos_a, pos_b)

        self.gen_update(xf_a, xf_b, p_a, p_b, pp_a, pp_b, x_a_recon, x_b_recon, x_a_recon_p, x_b_recon_p, images_a, images_b, labels_a, labels_b, self.hparams)
        return {'loss_total': self.loss_gen_total.item(), 
                'loss_cls': self.loss_id.item(), 
                'loss_cls_aug': self.loss_aug_id.item(), 
                'loss_recon_x': self.loss_gen_recon_x.item(),
                'loss_jsd': self.jsd_loss.item()}
                    
    def sample(self, x_a, x_b, pretrain_model=None):
        device = "cuda" if x_a.is_cuda else "cpu" 
        x_as, x_bs, x_a_aug, x_b_aug, x_a_aug1, x_b_aug1 = [], [], [], [], [], []
        for image_a, image_b in zip(x_a, x_b):
            x_b_, x_ab1, x_ab2= Augmix(self.TO_pil(image_b.cpu()), self.preprocess, no_jsd=not self.jsd)
            x_a_, x_ba1, x_ba2= Augmix(self.TO_pil(image_a.cpu()), self.preprocess, no_jsd=not self.jsd)
            x_a_aug.append(x_ba1.to(device).unsqueeze(0)); x_a_aug1.append(x_ba2.to(device).unsqueeze(0))
            x_b_aug.append(x_ab1.to(device).unsqueeze(0)); x_b_aug1.append(x_ab2.to(device).unsqueeze(0))
            x_as.append(x_a_.to(device).unsqueeze(0)); x_bs.append(x_b_.to(device).unsqueeze(0))
        x_a_aug, x_a_aug1=torch.cat(x_a_aug), torch.cat(x_a_aug1)
        x_b_aug, x_b_aug1 = torch.cat(x_b_aug), torch.cat(x_b_aug1)
        x_as, x_bs = torch.cat(x_as), torch.cat(x_bs)
        return x_as, x_a_aug, x_a_aug1, x_bs, x_b_aug, x_b_aug1

    def predict(self, x):
        return self.dis_id(self.id_featurizer(x)[-1])

    def step(self):
        self.iteration += 1

class DDG(ERM):
    
    def __init__(self, input_shape, num_classes, num_domains, hparams, device='cpu'):
        super(DDG, self).__init__(input_shape, num_classes, num_domains,
                                  hparams, device)
        self.hparams = hparams
        self.class_balance=hparams['class_balanced']
        self.iteration = 0
        self.id_featurizer = self.featurizer
        self.dis_id = self.classifier
        self.gen = networks.AdaINGen(1, self.id_featurizer.n_outputs, hparams) if not hparams['is_mnist'] else networks.VAEGen()
        self.dis_img = networks.MsImageDis(hparams=hparams) 
        self.recon_xp_w = hparams['recon_xp_w']
        self.recon_xn_w = hparams['recon_xn_w']
        self.margin = hparams['margin']
        self.eta = hparams['eta']
        
        def to_gray(half=False): #simple
            def forward(x):
                x = torch.mean(x, dim=1, keepdim=True)
                if half:
                    x = x.half()
                return x
            return forward
        self.single = to_gray(False)
        self.optimizer_gen = torch.optim.Adam([p for p in list(self.gen.parameters())  if p.requires_grad], lr=self.hparams['lr_g'], betas=(0, 0.999), weight_decay=self.hparams['weight_decay_g'])
        if self.hparams['stage'] == 0:
            # Setup the optimizers
            self.optimizer_dis_img = torch.optim.Adam(
                self.dis_img.parameters(),
                lr=self.hparams["lr_d"],
                weight_decay=self.hparams['weight_decay'])
            step = hparams['steps']*0.6
            print(step)
            self.dis_scheduler = lr_scheduler.MultiStepLR(self.optimizer_dis_img, milestones=[step, step+step//2, step+step//2+step//4],
                                        gamma=0.1)
            self.gen_scheduler = lr_scheduler.MultiStepLR(self.optimizer_gen, milestones=[step, step+step//2, step+step//2+step//4],
                                        gamma=0.1)

        self.id_criterion = nn.MSELoss() 
        self.dom_criterion = nn.CrossEntropyLoss()
        

    def recon_criterion(self, input, target, reduction=True):
            diff = input - target.detach()
            B,C,H,W = input.shape
            if reduction == False:
                return torch.mean(torch.abs(diff[:]).view(B,-1),dim=-1)
            return torch.mean(torch.abs(diff[:]))
    
    def train_bn(self, m):
        classname = m.__class__.__name__
        if classname.find('BatchNorm') != -1:
            m.train()
            print('there has bn')

    def forward(self, x_a, x_b, xp_a, xp_b):
        '''
            inpus:
                x_a, x_b: image from dataloader a,b
                xp_a, xp_b: positive pair of x_a, x_b
        '''
        s_a = self.gen.encode(self.single(x_a))# v for x_a
        s_b = self.gen.encode(self.single(x_b))# v for x_b
        s_pa = self.gen.encode(self.single(xp_a))# v for x_a
        s_pb = self.gen.encode(self.single(xp_b))# v for x_b
        f_a, x_fa = self.id_featurizer(x_a, self.hparams['stage']) # f_a: detached s for x_a, x_fa: s for x_a
        p_a = self.dis_id(x_fa)             # identity classification result for x_a
        f_b, x_fb = self.id_featurizer(x_b, self.hparams['stage'])
        p_b = self.dis_id(x_fb)
        fp_a, xp_fa = self.id_featurizer(xp_a, self.hparams['stage'])
        pp_a = self.dis_id(xp_fa)
        fp_b, xp_fb = self.id_featurizer(xp_b, self.hparams['stage'])
        pp_b = self.dis_id(xp_fb)
        # if self.hparams['stage'] == 0:
        #     # cross-style generation
        #     x_ba = self.gen.decode(s_b, f_a) # x_ba: generated from identity of a and style of b
        #     x_ab = self.gen.decode(s_a, f_b)
        #     x_a_recon = self.gen.decode(s_a, f_a) # generate from identity and style of a
        #     x_b_recon = self.gen.decode(s_b, f_b) 
        # else:
        #     x_ba = None
        #     x_ab = None
        #     x_a_recon = None
        #     x_b_recon = None
                    # cross-style generation
        x_ba = self.gen.decode(s_b, f_a) # x_ba: generated from identity of a and style of b
        x_ab = self.gen.decode(s_a, f_b)
        x_a_recon = self.gen.decode(s_a, f_a) # generate from identity and style of a
        x_b_recon = self.gen.decode(s_b, f_b) 
        x_a_recon_p = self.gen.decode(s_a, fp_a)
        x_b_recon_p = self.gen.decode(s_b, fp_b)

        return x_ab, x_ba, s_a, s_b, f_a, f_b, p_a, p_b, pp_a, pp_b, x_a_recon, x_b_recon, x_a_recon_p, x_b_recon_p    

    def dis_update(self, x_ab, x_ba, x_a, x_b, hparams):
        '''
            inpus:
                x_ab: generated from identity of b and style of a (fake)
                x_ba: generated from identity of a and style of b (fake)
                x_a, x_b: real image
        '''
        self.optimizer_dis_img.zero_grad()
        self.loss_dis_a, reg_a = self.dis_img.calc_dis_loss(self.dis_img, x_ba.detach(), x_a)
        self.loss_dis_b, reg_b = self.dis_img.calc_dis_loss(self.dis_img, x_ab.detach(), x_b)
        self.loss_dis_total = hparams['gan_w'] * self.loss_dis_a + hparams['gan_w'] * self.loss_dis_b
        self.loss_dis_total.backward() # discriminators are trained here
        self.optimizer_dis_img.step()     

    def gen_update(self, x_ab, x_ba, s_a, s_b, f_a, f_b, p_a, p_b, pp_a, pp_b, x_a_recon, x_b_recon, x_a_recon_p, x_b_recon_p, x_a, x_b,  l_a, l_b, hparams):
        '''
            inputs:
                x_ab: generated from identity of b and style of a
                x_ba: generated from identity of a and style of b
                s_a, s_b: style factors for x_a, x_b
                f_a, f_b: detached semantic factors for x_a, x_b
                p_a, p_b: identity prediction results for x_a, x_b
                pp_a, pp_b: identity prediction results for the positive pair of x_a, x_b
                x_a_recon, x_b_recon: reconstruction of x_a, x_b
                x_a_recon_p, x_b_recon_p: reconstruction of the positive pair of x_a, x_b
                x_a, x_b,  l_a, l_b: images and identity labels
                hparams: parameters
        '''
        self.optimizer_gen.zero_grad()
        self.optimizer.zero_grad()

        #################################

        # auto-encoder image reconstruction
        self.recon_a2a, self.recon_b2b = self.recon_criterion(x_a_recon_p, x_a, reduction=False), self.recon_criterion(x_b_recon_p, x_b, reduction=False)
        self.loss_gen_recon_p =  torch.mean(torch.max(self.recon_a2a-self.margin, torch.zeros_like(self.recon_a2a)))+ torch.mean(torch.max(self.recon_b2b-self.margin, torch.zeros_like(self.recon_b2b)))

        # Emprical Loss
        if not hparams['is_mnist']:
            _, x_fa_recon = self.id_featurizer(x_ab); p_a_recon = self.dis_id(x_fa_recon)
            _, x_fb_recon = self.id_featurizer(x_ba); p_b_recon = self.dis_id(x_fb_recon)
        else:
            _, x_fa_recon = self.id_featurizer(x_ba); p_a_recon = self.dis_id(x_fa_recon)
            _, x_fb_recon = self.id_featurizer(x_ab); p_b_recon = self.dis_id(x_fb_recon)            
        self.loss_id = self.id_criterion(p_a, l_a) + self.id_criterion(p_b, l_b) +  self.id_criterion(pp_a, l_a) + self.id_criterion(pp_b, l_b)
        self.loss_gen_recon_id = self.id_criterion(p_a_recon, l_a) + self.id_criterion(p_b_recon, l_b)

        self.step(torch.mean(self.recon_a2a))
        # total loss
        self.loss_gen_total = self.loss_id +\
                self.recon_xp_w * self.loss_gen_recon_p +\
                hparams['recon_id_w'] * self.loss_gen_recon_id 
        # print(self.loss_id, self.recon_xp_w,self.loss_gen_recon_p, hparams['recon_id_w'], self.loss_gen_recon_id , self.loss_gen_total)
        self.loss_gen_total.backward()
        self.optimizer_gen.step()
        self.optimizer.step()

    def gan_update(self, x_ab, x_ba, s_a, s_b, f_a, f_b, p_a, p_b, pp_a, pp_b, x_a_recon, x_b_recon, x_a_recon_p, x_b_recon_p, x_a, x_b,  l_a, l_b, hparams):
        '''
            Train the GAN
            inputs:
                x_ab: generated from identity of b and style of a
                x_ba: generated from identity of a and style of b
                s_a, s_b: style factors for x_a, x_b
                f_a, f_b: detached semantic factors for x_a, x_b
                p_a, p_b: identity prediction results for x_a, x_b
                pp_a, pp_b: identity prediction results for the positive pair of x_a, x_b
                x_a_recon, x_b_recon: reconstruction of x_a, x_b
                x_a_recon_p, x_b_recon_p: reconstruction of the positive pair of x_a, x_b
                x_a, x_b,  l_a, l_b: images and identity labels
                hparams: parameters
        '''
        self.optimizer_gen.zero_grad()
        self.optimizer.zero_grad()
 
        # no gradient
        x_ba_copy = Variable(x_ba.data, requires_grad=False)
        x_ab_copy = Variable(x_ab.data, requires_grad=False)
        f_a, f_b = f_a.detach(), f_b.detach()

        rand_num = random.uniform(0,1)
        #################################
        # encode structure
        if 0.5>=rand_num:
            # encode again (encoder is tuned, input is fixed)
            s_a_recon = self.gen.enc_content(self.single(x_ab_copy))
            s_b_recon = self.gen.enc_content(self.single(x_ba_copy))
        else:
            # copy the encoder
            self.enc_content_copy = copy.deepcopy(self.gen.enc_content)
            self.enc_content_copy = self.enc_content_copy.eval()
            # encode again (encoder is fixed, input is tuned)
            s_a_recon = self.enc_content_copy(self.single(x_ab))
            s_b_recon = self.enc_content_copy(self.single(x_ba))

        #################################
        # encode appearance
        self.id_copy = copy.deepcopy(self.id_featurizer)
        self.dis_id_copy = copy.deepcopy(self.dis_id)
        self.id_copy.eval()
        self.dis_id_copy.eval()

        # encode again (encoder is fixed, input is tuned)
        f_a_recon, _ = self.id_copy(x_ba);
        f_b_recon, _ = self.id_copy(x_ab); 

        # auto-encoder image reconstruction
        self.loss_gen_recon_x =  self.recon_criterion(x_a_recon, x_a)+self.recon_criterion(x_b_recon, x_b)

        # Emprical Loss

        x_aba, x_bab = self.gen.decode(s_a_recon, f_a_recon), self.gen.decode(s_b_recon, f_b_recon)  if hparams['recon_x_cyc_w'] > 0 else None
        self.loss_gen_cycrecon_x = self.recon_criterion(x_aba, x_a) + self.recon_criterion(x_bab, x_b) if hparams['recon_x_cyc_w'] > 0 else torch.tensor(0)

        # GAN loss
        self.loss_gen_adv = self.dis_img.calc_gen_loss(self.dis_img, x_ba) + self.dis_img.calc_gen_loss(self.dis_img, x_ab)

        self.step()
        if self.iteration > hparams['steps'] * hparams['warm_iter_r']:
            hparams['recon_x_cyc_w'] += hparams['warm_scale']
            hparams['recon_x_cyc_w'] = min(hparams['recon_x_cyc_w'], hparams['max_cyc_w'])

        # total loss
        self.loss_gen_total = hparams['gan_w'] * self.loss_gen_adv + \
                              hparams['recon_x_w'] * self.loss_gen_recon_x + \
                              hparams['recon_x_cyc_w'] * self.loss_gen_cycrecon_x

        self.loss_gen_total.backward()
        self.optimizer_gen.step()
        self.optimizer.step()


    def update(self, minibatches, minibatches_neg, pretrain_model=None, unlabeled=None, iteration=0):
        images_a = torch.cat([x for x, y, pos in minibatches])
        labels_a = torch.cat([y for x, y, pos in minibatches])
        pos_a = torch.cat([pos for x, y, pos in minibatches])
        images_b = torch.cat([x for x, y, pos in minibatches_neg])
        labels_b = torch.cat([y for x, y, pos in minibatches_neg])
        pos_b = torch.cat([pos for x, y, pos in minibatches_neg])
        

        if self.hparams['stage'] == 1 and pretrain_model is not None:
            # swap semantic factors
            s_a = pretrain_model.gen.encode(self.single(images_a))# v for x_a
            s_b = pretrain_model.gen.encode(self.single(images_b))# v for x_b
            f_a, x_fa = pretrain_model.id_featurizer(images_a) # f_a: detached s for x_a, x_fa: s for x_a
            f_b, x_fb = pretrain_model.id_featurizer(images_b)
            # cross-style generation
            x_ba = pretrain_model.gen.decode(s_b, f_a) # x_ba: generated from identity of a and style of b
            x_ab = pretrain_model.gen.decode(s_a, f_b)
            _, _, s_a, s_b, f_a, f_b, p_a, p_b, pp_a, pp_b, x_a_recon, x_b_recon, x_a_recon_p, x_b_recon_p = self.forward(images_a, images_b, pos_a, pos_b)
        else:
            x_ab, x_ba, s_a, s_b, f_a, f_b, p_a, p_b, pp_a, pp_b, x_a_recon, x_b_recon, x_a_recon_p, x_b_recon_p = self.forward(images_a, images_b, pos_a, pos_b)

        if self.hparams['stage'] == 0:
            self.dis_update(x_ab.clone(), x_ba.clone(), images_a, images_b, self.hparams)
            self.gan_update(x_ab, x_ba, s_a, s_b, f_a, f_b, p_a, p_b, pp_a, pp_b, x_a_recon, x_b_recon, x_a_recon_p, x_b_recon_p, images_a, images_b, labels_a, labels_b, self.hparams)
            self.gen_scheduler.step()
            self.dis_scheduler.step()
            return {'loss_total': self.loss_gen_total.item(), 
                    'loss_gan': self.loss_gen_adv.item(),
                    'loss_recon_x': self.loss_gen_recon_x.item(),
                    'loss_x_cyc': self.loss_gen_cycrecon_x.item()}
        else:
            self.gen_update(x_ab, x_ba, s_a, s_b, f_a, f_b, p_a, p_b, pp_a, pp_b, x_a_recon, x_b_recon, x_a_recon_p, x_b_recon_p, images_a, images_b, labels_a, labels_b, self.hparams)
            return {
                    'loss_cls': self.loss_id.item(),
                    'loss_gen_recon_id': self.loss_gen_recon_id.item(), 
                    'recon_xp_w': self.recon_xp_w,
                    'loss_recon_p': self.loss_gen_recon_p.item()}
                    
    def sample(self, x_a, x_b, pretrain_model=None):
        self.eval()
        x_a_recon, x_b_recon, x_ba1, x_ab1, x_aba, x_bab = [], [], [], [], [], []
        for i in range(x_a.size(0)):
            model = pretrain_model if pretrain_model is not None else self
            s_a = model.gen.encode( self.single(x_a[i].unsqueeze(0)) )
            s_b = model.gen.encode( self.single(x_b[i].unsqueeze(0)) )
            f_a, _ = model.id_featurizer(x_a[i].unsqueeze(0))
            f_b, _ = model.id_featurizer(x_b[i].unsqueeze(0))
            x_a_recon.append(model.gen.decode(s_a, f_a))
            x_b_recon.append(model.gen.decode(s_b, f_b))
            x_ba = model.gen.decode(s_b, f_a)
            x_ab = model.gen.decode(s_a, f_b)
            x_ba1.append(x_ba)
            x_ab1.append(x_ab)

        x_a_recon, x_b_recon = torch.cat(x_a_recon), torch.cat(x_b_recon)
        x_ba1, x_ab1 = torch.cat(x_ba1), torch.cat(x_ab1)
        self.train()

        return x_a, x_ba1, x_b, x_ab1

    def predict(self, x):
        return self.dis_id(self.id_featurizer(x)[-1])

    def step(self, recon_p=None):
        self.iteration += 1
        if recon_p is None:
            return
        self.recon_xp_w = min(max(self.recon_xp_w + self.eta * (recon_p.item() - self.margin), 0), 1)



class IGA(ERM):
    """
    Inter-environmental Gradient Alignment
    From https://arxiv.org/abs/2008.01883v2
    """

    def __init__(self, in_features, num_classes, num_domains, hparams):
        super(IGA, self).__init__(in_features, num_classes, num_domains, hparams)

    def update(self, minibatches, unlabeled=None):
        total_loss = 0
        grads = []
        for i, (x, y) in enumerate(minibatches):
            logits = self.network(x)

            env_loss = F.l1_loss(logits, y.view(-1))
            total_loss += env_loss

            env_grad = autograd.grad(env_loss, self.network.parameters(),
                                        create_graph=True)

            grads.append(env_grad)

        mean_loss = total_loss / len(minibatches)
        mean_grad = autograd.grad(mean_loss, self.network.parameters(),
                                        retain_graph=True)

        # compute trace penalty
        penalty_value = 0
        for grad in grads:
            for g, mean_g in zip(grad, mean_grad):
                penalty_value += (g - mean_g).pow(2).sum()

        objective = mean_loss + self.hparams['penalty'] * penalty_value

        self.optimizer.zero_grad()
        objective.backward()
        self.optimizer.step()

        return {'loss': mean_loss.item(), 'penalty': penalty_value.item()}


class SelfReg(ERM):
    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(SelfReg, self).__init__(input_shape, num_classes, num_domains,
                                   hparams)
        self.num_classes = num_classes
        self.MSEloss = nn.MSELoss()
        input_feat_size = self.featurizer.n_outputs
        hidden_size = input_feat_size if input_feat_size==2048 else input_feat_size*2

        self.cdpl = nn.Sequential(
                            nn.Linear(input_feat_size, hidden_size),
                            nn.BatchNorm1d(hidden_size),
                            nn.ReLU(inplace=True),
                            nn.Linear(hidden_size, hidden_size),
                            nn.BatchNorm1d(hidden_size),
                            nn.ReLU(inplace=True),
                            nn.Linear(hidden_size, input_feat_size),
                            nn.BatchNorm1d(input_feat_size)
        )

    def update(self, minibatches, unlabeled=None):

        all_x = torch.cat([x for x, y in minibatches])
        all_y = torch.cat([y for _, y in minibatches])

        lam = np.random.beta(0.5, 0.5)

        batch_size = all_y.size()[0]

        # cluster and order features into same-class group
        with torch.no_grad():
            sorted_y, indices = torch.sort(all_y)
            sorted_x = torch.zeros_like(all_x)
            for idx, order in enumerate(indices):
                sorted_x[idx] = all_x[order]
            intervals = []
            ex = 0
            for idx, val in enumerate(sorted_y):
                if ex==val:
                    continue
                intervals.append(idx)
                ex = val
            intervals.append(batch_size)

            all_x = sorted_x
            all_y = sorted_y

        feat = self.featurizer(all_x)
        proj = self.cdpl(feat)

        output = self.classifier(feat)

        # shuffle
        output_2 = torch.zeros_like(output)
        feat_2 = torch.zeros_like(proj)
        output_3 = torch.zeros_like(output)
        feat_3 = torch.zeros_like(proj)
        ex = 0
        for end in intervals:
            shuffle_indices = torch.randperm(end-ex)+ex
            shuffle_indices2 = torch.randperm(end-ex)+ex
            for idx in range(end-ex):
                output_2[idx+ex] = output[shuffle_indices[idx]]
                feat_2[idx+ex] = proj[shuffle_indices[idx]]
                output_3[idx+ex] = output[shuffle_indices2[idx]]
                feat_3[idx+ex] = proj[shuffle_indices2[idx]]
            ex = end

        # mixup
        output_3 = lam*output_2 + (1-lam)*output_3
        feat_3 = lam*feat_2 + (1-lam)*feat_3

        # regularization
        L_ind_logit = self.MSEloss(output, output_2)
        L_hdl_logit = self.MSEloss(output, output_3)
        L_ind_feat = 0.3 * self.MSEloss(feat, feat_2)
        L_hdl_feat = 0.3 * self.MSEloss(feat, feat_3)

        cl_loss = F.l1_loss(output.view(-1), all_y.view(-1))
        C_scale = min(cl_loss.item(), 1.)
        loss = cl_loss + C_scale*(lam*(L_ind_logit + L_ind_feat)+(1-lam)*(L_hdl_logit + L_hdl_feat))

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return {'loss': loss.item()}


class SANDMask(ERM):
    """
    SAND-mask: An Enhanced Gradient Masking Strategy for the Discovery of Invariances in Domain Generalization
    <https://arxiv.org/abs/2106.02266>
    """

    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(SANDMask, self).__init__(input_shape, num_classes, num_domains, hparams)

        self.tau = hparams["tau"]
        self.k = hparams["k"]
        betas = (0.9, 0.999)
        self.optimizer = torch.optim.Adam(
            self.network.parameters(),
            lr=self.hparams["lr"],
            weight_decay=self.hparams['weight_decay'],
            betas=betas
        )

        self.register_buffer('update_count', torch.tensor([0]))

    def update(self, minibatches, unlabeled=None):

        mean_loss = 0
        param_gradients = [[] for _ in self.network.parameters()]
        for i, (x, y) in enumerate(minibatches):
            logits = self.network(x)

            env_loss = F.l1_loss(logits.view(-1), y.view(-1))
            mean_loss += env_loss.item() / len(minibatches)
            env_grads = autograd.grad(env_loss, self.network.parameters(), retain_graph=True)
            for grads, env_grad in zip(param_gradients, env_grads):
                grads.append(env_grad)

        self.optimizer.zero_grad()
        # gradient masking applied here
        self.mask_grads(param_gradients, self.network.parameters())
        self.optimizer.step()
        self.update_count += 1

        return {'loss': mean_loss}

    def mask_grads(self, gradients, params):
        '''
        Here a mask with continuous values in the range [0,1] is formed to control the amount of update for each
        parameter based on the agreement of gradients coming from different environments.
        '''
        device = gradients[0][0].device
        for param, grads in zip(params, gradients):
            grads = torch.stack(grads, dim=0)
            avg_grad = torch.mean(grads, dim=0)
            grad_signs = torch.sign(grads)
            gamma = torch.tensor(1.0).to(device)
            grads_var = grads.var(dim=0)
            grads_var[torch.isnan(grads_var)] = 1e-17
            lam = (gamma * grads_var).pow(-1)
            mask = torch.tanh(self.k * lam * (torch.abs(grad_signs.mean(dim=0)) - self.tau))
            mask = torch.max(mask, torch.zeros_like(mask))
            mask[torch.isnan(mask)] = 1e-17
            mask_t = (mask.sum() / mask.numel())
            param.grad = mask * avg_grad
            param.grad *= (1. / (1e-10 + mask_t))



class Fishr(Algorithm):
    "Invariant Gradients variances for Out-of-distribution Generalization"

    def __init__(self, input_shape, num_classes, num_domains, hparams):
        assert backpack is not None, "Install backpack with: 'pip install backpack-for-pytorch==1.3.0'"
        super(Fishr, self).__init__(input_shape, num_classes, num_domains, hparams)
        self.num_domains = num_domains

        self.featurizer = networks.Featurizer(input_shape, self.hparams)
        self.classifier = extend(
            networks.Classifier(
                self.featurizer.n_outputs,
                num_classes,
                self.hparams['nonlinear_classifier'],
            )
        )
        self.network = nn.Sequential(self.featurizer, self.classifier)

        self.register_buffer("update_count", torch.tensor([0]))
        self.bce_extended = extend(nn.CrossEntropyLoss(reduction='none'))
        self.ema_per_domain = [
            MovingAverage(ema=self.hparams["ema"], oneminusema_correction=True)
            for _ in range(self.num_domains)
        ]
        self._init_optimizer()

    def _init_optimizer(self):
        self.optimizer = torch.optim.Adam(
            list(self.featurizer.parameters()) + list(self.classifier.parameters()),
            lr=self.hparams["lr"],
            weight_decay=self.hparams["weight_decay"],
        )

    def update(self, minibatches, unlabeled=None):
        assert len(minibatches) == self.num_domains
        all_x = torch.cat([x for x, y in minibatches])
        all_y = torch.cat([y for x, y in minibatches])
        len_minibatches = [x.shape[0] for x, y in minibatches]

        all_z = self.featurizer(all_x)
        all_logits = self.classifier(all_z)

        penalty = self.compute_fishr_penalty(all_logits, all_y, len_minibatches)
        all_nll = F.l1_loss(all_logits.view(-1), all_y.view(-1))

        penalty_weight = 0
        if self.update_count >= self.hparams["penalty_anneal_iters"]:
            penalty_weight = self.hparams["lambda"]
            if self.update_count == self.hparams["penalty_anneal_iters"] != 0:
                # Reset Adam as in IRM or V-REx, because it may not like the sharp jump in
                # gradient magnitudes that happens at this step.
                self._init_optimizer()
        self.update_count += 1

        objective = all_nll + penalty_weight * penalty
        self.optimizer.zero_grad()
        objective.backward()
        self.optimizer.step()

        return {'loss': objective.item(), 'nll': all_nll.item(), 'penalty': penalty.item()}

    def compute_fishr_penalty(self, all_logits, all_y, len_minibatches):
        dict_grads = self._get_grads(all_logits, all_y)
        grads_var_per_domain = self._get_grads_var_per_domain(dict_grads, len_minibatches)
        return self._compute_distance_grads_var(grads_var_per_domain)

    def _get_grads(self, logits, y):
        self.optimizer.zero_grad()
        loss = self.bce_extended(logits, y).sum()
        with backpack(BatchGrad()):
            loss.backward(
                inputs=list(self.classifier.parameters()), retain_graph=True, create_graph=True
            )

        # compute individual grads for all samples across all domains simultaneously
        dict_grads = OrderedDict(
            [
                (name, weights.grad_batch.clone().view(weights.grad_batch.size(0), -1))
                for name, weights in self.classifier.named_parameters()
            ]
        )
        return dict_grads

    def _get_grads_var_per_domain(self, dict_grads, len_minibatches):
        # grads var per domain
        grads_var_per_domain = [{} for _ in range(self.num_domains)]
        for name, _grads in dict_grads.items():
            all_idx = 0
            for domain_id, bsize in enumerate(len_minibatches):
                env_grads = _grads[all_idx:all_idx + bsize]
                all_idx += bsize
                env_mean = env_grads.mean(dim=0, keepdim=True)
                env_grads_centered = env_grads - env_mean
                grads_var_per_domain[domain_id][name] = (env_grads_centered).pow(2).mean(dim=0)

        # moving average
        for domain_id in range(self.num_domains):
            grads_var_per_domain[domain_id] = self.ema_per_domain[domain_id].update(
                grads_var_per_domain[domain_id]
            )

        return grads_var_per_domain

    def _compute_distance_grads_var(self, grads_var_per_domain):

        # compute gradient variances averaged across domains
        grads_var = OrderedDict(
            [
                (
                    name,
                    torch.stack(
                        [
                            grads_var_per_domain[domain_id][name]
                            for domain_id in range(self.num_domains)
                        ],
                        dim=0
                    ).mean(dim=0)
                )
                for name in grads_var_per_domain[0].keys()
            ]
        )

        penalty = 0
        for domain_id in range(self.num_domains):
            penalty += l2_between_dicts(grads_var_per_domain[domain_id], grads_var)
        return penalty / self.num_domains

    def predict(self, x):
        return self.network(x)

class TRM(Algorithm):
    """
    Learning Representations that Support Robust Transfer of Predictors
    <https://arxiv.org/abs/2110.09940>
    """

    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(TRM, self).__init__(input_shape, num_classes, num_domains,hparams)
        self.register_buffer('update_count', torch.tensor([0]))
        self.num_domains = num_domains
        self.featurizer = networks.Featurizer(input_shape, self.hparams)
        self.classifier = nn.Linear(self.featurizer.n_outputs, num_classes).cuda()
        self.clist = [nn.Linear(self.featurizer.n_outputs, num_classes).cuda() for i in range(num_domains+1)]
        self.olist = [torch.optim.SGD(
            self.clist[i].parameters(),
            lr=1e-1,
        ) for i in range(num_domains+1)]

        self.optimizer_f = torch.optim.Adam(
            self.featurizer.parameters(),
            lr=self.hparams["lr"],
            weight_decay=self.hparams['weight_decay']
        )
        self.optimizer_c = torch.optim.Adam(
            self.classifier.parameters(),
            lr=self.hparams["lr"],
            weight_decay=self.hparams['weight_decay']
        )
        # initial weights
        self.alpha = torch.ones((num_domains, num_domains)).cuda() - torch.eye(num_domains).cuda()

    @staticmethod
    def neum(v, model, batch):
        def hvp(y, w, v):

            # First backprop
            first_grads = autograd.grad(y, w, retain_graph=True, create_graph=True, allow_unused=True)
            first_grads = torch.nn.utils.parameters_to_vector(first_grads)
            # Elementwise products
            elemwise_products = first_grads @ v
            # Second backprop
            return_grads = autograd.grad(elemwise_products, w, create_graph=True)
            return_grads = torch.nn.utils.parameters_to_vector(return_grads)
            return return_grads

        v = v.detach()
        h_estimate = v
        cnt = 0.
        model.eval()
        iter = 10
        for i in range(iter):
            model.weight.grad *= 0
            y = model(batch[0].detach())
            loss = F.l1_loss(y.view(-1), batch[1].detach().view(-1))
            hv = hvp(loss, model.weight, v)
            v -= hv
            v = v.detach()
            h_estimate = v + h_estimate
            h_estimate = h_estimate.detach()
            # not converge
            if torch.max(abs(h_estimate)) > 10:
                break
            cnt += 1

        model.train()
        return h_estimate.detach()

    def update(self, minibatches, unlabeled=None):

        loss_swap = 0.0
        trm = 0.0

        if self.update_count >= self.hparams['iters']:
            # TRM
            if self.hparams['class_balanced']:
                # for stability when facing unbalanced labels across environments
                for classifier in self.clist:
                    classifier.weight.data = copy.deepcopy(self.classifier.weight.data)
            self.alpha /= self.alpha.sum(1, keepdim=True)

            self.featurizer.train()
            all_x = torch.cat([x for x, y in minibatches])
            all_y = torch.cat([y for x, y in minibatches])
            all_feature = self.featurizer(all_x)
            # updating original network
            loss = F.l1_loss(self.classifier(all_feature).view(-1), all_y.view(-1))

            for i in range(30):
                all_logits_idx = 0
                loss_erm = 0.
                for j, (x, y) in enumerate(minibatches):
                    # j-th domain
                    feature = all_feature[all_logits_idx:all_logits_idx + x.shape[0]]
                    all_logits_idx += x.shape[0]
                    loss_erm += F.l1_loss(self.clist[j](feature.detach()).view(-1), y.view(-1))
                for opt in self.olist:
                    opt.zero_grad()
                loss_erm.backward()
                for opt in self.olist:
                    opt.step()

            # collect (feature, y)
            feature_split = list()
            y_split = list()
            all_logits_idx = 0
            for i, (x, y) in enumerate(minibatches):
                feature = all_feature[all_logits_idx:all_logits_idx + x.shape[0]]
                all_logits_idx += x.shape[0]
                feature_split.append(feature)
                y_split.append(y)

            # estimate transfer risk
            for Q, (x, y) in enumerate(minibatches):
                sample_list = list(range(len(minibatches)))
                sample_list.remove(Q)

                loss_Q = F.l1_loss(self.clist[Q](feature_split[Q]).view(-1), y_split[Q].view(-1))
                grad_Q = autograd.grad(loss_Q, self.clist[Q].weight, create_graph=True)
                vec_grad_Q = nn.utils.parameters_to_vector(grad_Q)

                loss_P = [F.l1_loss(self.clist[Q](feature_split[i]).view(-1), y_split[i])*(self.alpha[Q, i].data.detach().view(-1))
                          if i in sample_list else 0. for i in range(len(minibatches))]
                loss_P_sum = sum(loss_P)
                grad_P = autograd.grad(loss_P_sum, self.clist[Q].weight, create_graph=True)
                vec_grad_P = nn.utils.parameters_to_vector(grad_P).detach()
                vec_grad_P = self.neum(vec_grad_P, self.clist[Q], (feature_split[Q], y_split[Q]))

                loss_swap += loss_P_sum - self.hparams['cos_lambda'] * (vec_grad_P.detach() @ vec_grad_Q)

                for i in sample_list:
                    self.alpha[Q, i] *= (self.hparams["groupdro_eta"] * loss_P[i].data).exp()

            loss_swap /= len(minibatches)
            trm /= len(minibatches)
        else:
            # ERM
            self.featurizer.train()
            all_x = torch.cat([x for x, y in minibatches])
            all_y = torch.cat([y for x, y in minibatches])
            all_feature = self.featurizer(all_x)
            loss = F.l1_loss(self.classifier(all_feature).view(-1), all_y.view(-1))

        nll = loss.item()
        self.optimizer_c.zero_grad()
        self.optimizer_f.zero_grad()
        if self.update_count >= self.hparams['iters']:
            loss_swap = (loss + loss_swap)
        else:
            loss_swap = loss

        loss_swap.backward()
        self.optimizer_f.step()
        self.optimizer_c.step()

        loss_swap = loss_swap.item() - nll
        self.update_count += 1

        return {'nll': nll, 'trm_loss': loss_swap}

    def predict(self, x):
        return self.classifier(self.featurizer(x))

    def train(self):
        self.featurizer.train()

    def eval(self):
        self.featurizer.eval()

class IB_ERM(ERM):
    """Information Bottleneck based ERM on feature with conditionning"""

    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(IB_ERM, self).__init__(input_shape, num_classes, num_domains,
                                  hparams)
        self.optimizer = torch.optim.Adam(
            list(self.featurizer.parameters()) + list(self.classifier.parameters()),
            lr=self.hparams["lr"],
            weight_decay=self.hparams['weight_decay']
        )
        self.register_buffer('update_count', torch.tensor([0]))

    def update(self, minibatches, unlabeled=None):
        device = "cuda" if minibatches[0][0].is_cuda else "cpu"
        ib_penalty_weight = (self.hparams['ib_lambda'] if self.update_count
                          >= self.hparams['ib_penalty_anneal_iters'] else
                          0.0)

        nll = 0.
        ib_penalty = 0.

        all_x = torch.cat([x for x, y in minibatches])
        all_features = self.featurizer(all_x)
        all_logits = self.classifier(all_features)
        all_logits_idx = 0
        for i, (x, y) in enumerate(minibatches):
            features = all_features[all_logits_idx:all_logits_idx + x.shape[0]]
            logits = all_logits[all_logits_idx:all_logits_idx + x.shape[0]]
            all_logits_idx += x.shape[0]
            nll += F.l1_loss(logits.view(-1), y.view(-1))
            ib_penalty += features.var(dim=0).mean()

        nll /= len(minibatches)
        ib_penalty /= len(minibatches)

        # Compile loss
        loss = nll
        loss += ib_penalty_weight * ib_penalty

        if self.update_count == self.hparams['ib_penalty_anneal_iters']:
            # Reset Adam, because it doesn't like the sharp jump in gradient
            # magnitudes that happens at this step.
            self.optimizer = torch.optim.Adam(
                list(self.featurizer.parameters()) + list(self.classifier.parameters()),
                lr=self.hparams["lr"],
                weight_decay=self.hparams['weight_decay'])

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        self.update_count += 1
        return {'loss': loss.item(),
                'nll': nll.item(),
                'IB_penalty': ib_penalty.item()}

class IB_IRM(ERM):
    """Information Bottleneck based IRM on feature with conditionning"""

    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(IB_IRM, self).__init__(input_shape, num_classes, num_domains,
                                  hparams)
        self.optimizer = torch.optim.Adam(
            list(self.featurizer.parameters()) + list(self.classifier.parameters()),
            lr=self.hparams["lr"],
            weight_decay=self.hparams['weight_decay']
        )
        self.register_buffer('update_count', torch.tensor([0]))

    @staticmethod
    def _irm_penalty(logits, y):
        device = "cuda" if logits[0][0].is_cuda else "cpu"
        scale = torch.tensor(1.).to(device).requires_grad_()
        loss_1 = F.l1_loss(logits[::2].view(-1)* scale, y[::2].view(-1))
        loss_2 = F.l1_loss(logits[1::2].view(-1) * scale, y[1::2].view(-1))
        grad_1 = autograd.grad(loss_1, [scale], create_graph=True)[0]
        grad_2 = autograd.grad(loss_2, [scale], create_graph=True)[0]
        result = torch.sum(grad_1 * grad_2)
        return result

    def update(self, minibatches, unlabeled=None):
        device = "cuda" if minibatches[0][0].is_cuda else "cpu"
        irm_penalty_weight = (self.hparams['irm_lambda'] if self.update_count
                          >= self.hparams['irm_penalty_anneal_iters'] else
                          1.0)
        ib_penalty_weight = (self.hparams['ib_lambda'] if self.update_count
                          >= self.hparams['ib_penalty_anneal_iters'] else
                          0.0)

        nll = 0.
        irm_penalty = 0.
        ib_penalty = 0.

        all_x = torch.cat([x for x, y in minibatches])
        all_features = self.featurizer(all_x)
        all_logits = self.classifier(all_features)
        all_logits_idx = 0
        for i, (x, y) in enumerate(minibatches):
            features = all_features[all_logits_idx:all_logits_idx + x.shape[0]]
            logits = all_logits[all_logits_idx:all_logits_idx + x.shape[0]]
            all_logits_idx += x.shape[0]
            nll += F.l1_loss(logits.view(-1), y.view(-1))
            irm_penalty += self._irm_penalty(logits, y)
            ib_penalty += features.var(dim=0).mean()

        nll /= len(minibatches)
        irm_penalty /= len(minibatches)
        ib_penalty /= len(minibatches)

        # Compile loss
        loss = nll
        loss += irm_penalty_weight * irm_penalty
        loss += ib_penalty_weight * ib_penalty

        if self.update_count == self.hparams['irm_penalty_anneal_iters'] or self.update_count == self.hparams['ib_penalty_anneal_iters']:
            # Reset Adam, because it doesn't like the sharp jump in gradient
            # magnitudes that happens at this step.
            self.optimizer = torch.optim.Adam(
                list(self.featurizer.parameters()) + list(self.classifier.parameters()),
                lr=self.hparams["lr"],
                weight_decay=self.hparams['weight_decay'])

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        self.update_count += 1
        return {'loss': loss.item(),
                'nll': nll.item(),
                'IRM_penalty': irm_penalty.item(),
                'IB_penalty': ib_penalty.item()}


class AbstractCAD(Algorithm):
    """Contrastive adversarial domain bottleneck (abstract class)
    from Optimal Representations for Covariate Shift <https://arxiv.org/abs/2201.00057>
    """

    def __init__(self, input_shape, num_classes, num_domains,
                 hparams, is_conditional):
        super(AbstractCAD, self).__init__(input_shape, num_classes, num_domains, hparams)

        self.featurizer = networks.Featurizer(input_shape, self.hparams)
        self.classifier = networks.Classifier(
            self.featurizer.n_outputs,
            num_classes,
            self.hparams['nonlinear_classifier'])
        params = list(self.featurizer.parameters()) + list(self.classifier.parameters())

        # parameters for domain bottleneck loss
        self.is_conditional = is_conditional  # whether to use bottleneck conditioned on the label
        self.base_temperature = 0.07
        self.temperature = hparams['temperature']
        self.is_project = hparams['is_project']  # whether apply projection head
        self.is_normalized = hparams['is_normalized'] # whether apply normalization to representation when computing loss

        # whether flip maximize log(p) (False) to minimize -log(1-p) (True) for the bottleneck loss
        # the two versions have the same optima, but we find the latter is more stable
        self.is_flipped = hparams["is_flipped"]

        if self.is_project:
            self.project = nn.Sequential(
                nn.Linear(feature_dim, feature_dim),
                nn.ReLU(inplace=True),
                nn.Linear(feature_dim, 128),
            )
            params += list(self.project.parameters())

        # Optimizers
        self.optimizer = torch.optim.Adam(
            params,
            lr=self.hparams["lr"],
            weight_decay=self.hparams['weight_decay']
        )

    def bn_loss(self, z, y, dom_labels):
        device = z.device
        batch_size = z.shape[0]

        y = y.contiguous().view(-1, 1)
        dom_labels = dom_labels.contiguous().view(-1, 1)
        mask_y = torch.eq(y, y.T).to(device)
        mask_d = (torch.eq(dom_labels, dom_labels.T)).to(device)
        mask_drop = ~torch.eye(batch_size).bool().to(device)  # drop the "current"/"self" example
        mask_y &= mask_drop
        mask_y_n_d = mask_y & (~mask_d)  # contain the same label but from different domains
        mask_y_d = mask_y & mask_d  # contain the same label and the same domain
        mask_y, mask_drop, mask_y_n_d, mask_y_d = mask_y.float(), mask_drop.float(), mask_y_n_d.float(), mask_y_d.float()

        # compute logits
        if self.is_project:
            z = self.project(z)
        if self.is_normalized:
            z = F.normalize(z, dim=1)
        outer = z @ z.T
        logits = outer / self.temperature
        logits = logits * mask_drop
        # for numerical stability
        logits_max, _ = torch.max(logits, dim=1, keepdim=True)
        logits = logits - logits_max.detach()

        if not self.is_conditional:
            # unconditional CAD loss
            denominator = torch.logsumexp(logits + mask_drop.log(), dim=1, keepdim=True)
            log_prob = logits - denominator

            mask_valid = (mask_y.sum(1) > 0)
            log_prob = log_prob[mask_valid]
            mask_d = mask_d[mask_valid]

            if self.is_flipped:  # maximize log prob of samples from different domains
                bn_loss = - (self.temperature / self.base_temperature) * torch.logsumexp(
                    log_prob + (~mask_d).float().log(), dim=1)
            else:  # minimize log prob of samples from same domain
                bn_loss = (self.temperature / self.base_temperature) * torch.logsumexp(
                    log_prob + (mask_d).float().log(), dim=1)
        else:
            # conditional CAD loss
            if self.is_flipped:
                mask_valid = (mask_y_n_d.sum(1) > 0)
            else:
                mask_valid = (mask_y_d.sum(1) > 0)

            mask_y = mask_y[mask_valid]
            mask_y_d = mask_y_d[mask_valid]
            mask_y_n_d = mask_y_n_d[mask_valid]
            logits = logits[mask_valid]

            # compute log_prob_y with the same label
            denominator = torch.logsumexp(logits + mask_y.log(), dim=1, keepdim=True)
            log_prob_y = logits - denominator

            if self.is_flipped:  # maximize log prob of samples from different domains and with same label
                bn_loss = - (self.temperature / self.base_temperature) * torch.logsumexp(
                    log_prob_y + mask_y_n_d.log(), dim=1)
            else:  # minimize log prob of samples from same domains and with same label
                bn_loss = (self.temperature / self.base_temperature) * torch.logsumexp(
                    log_prob_y + mask_y_d.log(), dim=1)

        def finite_mean(x):
            # only 1D for now
            num_finite = (torch.isfinite(x).float()).sum()
            mean = torch.where(torch.isfinite(x), x, torch.tensor(0.0).to(x)).sum()
            if num_finite != 0:
                mean = mean / num_finite
            else:
                return torch.tensor(0.0).to(x)
            return mean

        return finite_mean(bn_loss)

    def update(self, minibatches, unlabeled=None):
        device = "cuda" if minibatches[0][0].is_cuda else "cpu"
        all_x = torch.cat([x for x, y in minibatches])
        all_y = torch.cat([y for x, y in minibatches])
        all_z = self.featurizer(all_x)
        all_d = torch.cat([
            torch.full((x.shape[0],), i, dtype=torch.int64, device=device)
            for i, (x, y) in enumerate(minibatches)
        ])

        bn_loss = self.bn_loss(all_z, all_y, all_d)
        clf_out = self.classifier(all_z)
        clf_loss = F.l1_loss(clf_out.view(-1), all_y.view(-1))
        total_loss = clf_loss + self.hparams['lmbda'] * bn_loss

        self.optimizer.zero_grad()
        total_loss.backward()
        self.optimizer.step()

        return {"clf_loss": clf_loss.item(), "bn_loss": bn_loss.item(), "total_loss": total_loss.item()}

    def predict(self, x):
        return self.classifier(self.featurizer(x))


class CAD(AbstractCAD):
    """Contrastive Adversarial Domain (CAD) bottleneck

       Properties:
       - Minimize I(D;Z)
       - Require access to domain labels but not task labels
       """

    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(CAD, self).__init__(input_shape, num_classes, num_domains, hparams, is_conditional=False)


class CondCAD(AbstractCAD):
    """Conditional Contrastive Adversarial Domain (CAD) bottleneck

    Properties:
    - Minimize I(D;Z|Y)
    - Require access to both domain labels and task labels
    """
    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(CondCAD, self).__init__(input_shape, num_classes, num_domains, hparams, is_conditional=True)


class Transfer(Algorithm):
    '''Algorithm 1 in Quantifying and Improving Transferability in Domain Generalization (https://arxiv.org/abs/2106.03632)'''
    ''' tries to ensure transferability among source domains, and thus transferabiilty between source and target'''
    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(Transfer, self).__init__(input_shape, num_classes, num_domains, hparams)
        self.register_buffer('update_count', torch.tensor([0]))
        self.d_steps_per_g = hparams['d_steps_per_g']

        # Architecture
        self.featurizer = networks.Featurizer(input_shape, self.hparams)
        self.classifier = networks.Classifier(
            self.featurizer.n_outputs,
            num_classes,
            self.hparams['nonlinear_classifier'])
        self.adv_classifier = networks.Classifier(
            self.featurizer.n_outputs,
            num_classes,
            self.hparams['nonlinear_classifier'])
        self.adv_classifier.load_state_dict(self.classifier.state_dict())

        # Optimizers
        if self.hparams['gda']:
            self.optimizer = torch.optim.SGD(self.adv_classifier.parameters(), lr=self.hparams['lr'])
        else:
            self.optimizer = torch.optim.Adam(
            (list(self.featurizer.parameters()) + list(self.classifier.parameters())),
                lr=self.hparams["lr"],
                weight_decay=self.hparams['weight_decay'])

        self.adv_opt = torch.optim.SGD(self.adv_classifier.parameters(), lr=self.hparams['lr_d'])

    def loss_gap(self, minibatches, device):
        ''' compute gap = max_i loss_i(h) - min_j loss_j(h), return i, j, and the gap for a single batch'''
        max_env_loss, min_env_loss =  torch.tensor([-float('inf')], device=device), torch.tensor([float('inf')], device=device)
        for x, y in minibatches:
            p = self.adv_classifier(self.featurizer(x))
            loss = F.l1_loss(p.view(-1), y.view(-1))
            if loss > max_env_loss:
                max_env_loss = loss
            if loss < min_env_loss:
                min_env_loss = loss
        return max_env_loss - min_env_loss

    def update(self, minibatches, unlabeled=None):
        device = "cuda" if minibatches[0][0].is_cuda else "cpu"
        # outer loop
        all_x = torch.cat([x for x, y in minibatches])
        all_y = torch.cat([y for x, y in minibatches])
        loss = F.l1_loss(self.predict(all_x).view(-1), all_y.view(-1))
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        del all_x, all_y
        gap = self.hparams['t_lambda'] * self.loss_gap(minibatches, device)
        self.optimizer.zero_grad()
        gap.backward()
        self.optimizer.step()
        self.adv_classifier.load_state_dict(self.classifier.state_dict())
        for _ in range(self.d_steps_per_g):
            self.adv_opt.zero_grad()
            gap = -self.hparams['t_lambda'] * self.loss_gap(minibatches, device)
            gap.backward()
            self.adv_opt.step()
            self.adv_classifier = proj(self.hparams['delta'], self.adv_classifier, self.classifier)
        return {'loss': loss.item(), 'gap': -gap.item()}

    def update_second(self, minibatches, unlabeled=None):
        device = "cuda" if minibatches[0][0].is_cuda else "cpu"
        self.update_count = (self.update_count + 1) % (1 + self.d_steps_per_g)
        if self.update_count.item() == 1:
            all_x = torch.cat([x for x, y in minibatches])
            all_y = torch.cat([y for x, y in minibatches])
            loss = F.l1_loss(self.predict(all_x).view(-1), all_y.view(-1))
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            del all_x, all_y
            gap = self.hparams['t_lambda'] * self.loss_gap(minibatches, device)
            self.optimizer.zero_grad()
            gap.backward()
            self.optimizer.step()
            self.adv_classifier.load_state_dict(self.classifier.state_dict())
            return {'loss': loss.item(), 'gap': gap.item()}
        else:
            self.adv_opt.zero_grad()
            gap = -self.hparams['t_lambda'] * self.loss_gap(minibatches, device)
            gap.backward()
            self.adv_opt.step()
            self.adv_classifier = proj(self.hparams['delta'], self.adv_classifier, self.classifier)
            return {'gap': -gap.item()}


    def predict(self, x):
        return self.classifier(self.featurizer(x))


class AbstractCausIRL(ERM):
    '''Abstract class for Causality based invariant representation learning algorithm from (https://arxiv.org/abs/2206.11646)'''
    def __init__(self, input_shape, num_classes, num_domains, hparams, gaussian):
        super(AbstractCausIRL, self).__init__(input_shape, num_classes, num_domains,
                                  hparams)
        if gaussian:
            self.kernel_type = "gaussian"
        else:
            self.kernel_type = "mean_cov"

    def my_cdist(self, x1, x2):
        x1_norm = x1.pow(2).sum(dim=-1, keepdim=True)
        x2_norm = x2.pow(2).sum(dim=-1, keepdim=True)
        res = torch.addmm(x2_norm.transpose(-2, -1),
                          x1,
                          x2.transpose(-2, -1), alpha=-2).add_(x1_norm)
        return res.clamp_min_(1e-30)

    def gaussian_kernel(self, x, y, gamma=[0.001, 0.01, 0.1, 1, 10, 100,
                                           1000]):
        D = self.my_cdist(x, y)
        K = torch.zeros_like(D)

        for g in gamma:
            K.add_(torch.exp(D.mul(-g)))

        return K

    def mmd(self, x, y):
        if self.kernel_type == "gaussian":
            Kxx = self.gaussian_kernel(x, x).mean()
            Kyy = self.gaussian_kernel(y, y).mean()
            Kxy = self.gaussian_kernel(x, y).mean()
            return Kxx + Kyy - 2 * Kxy
        else:
            mean_x = x.mean(0, keepdim=True)
            mean_y = y.mean(0, keepdim=True)
            cent_x = x - mean_x
            cent_y = y - mean_y
            cova_x = (cent_x.t() @ cent_x) / (len(x) - 1)
            cova_y = (cent_y.t() @ cent_y) / (len(y) - 1)

            mean_diff = (mean_x - mean_y).pow(2).mean()
            cova_diff = (cova_x - cova_y).pow(2).mean()

            return mean_diff + cova_diff

    def update(self, minibatches, unlabeled=None):
        objective = 0
        penalty = 0
        nmb = len(minibatches)

        features = [self.featurizer(xi) for xi, _ in minibatches]
        classifs = [self.classifier(fi) for fi in features]
        targets = [yi for _, yi in minibatches]

        first = None
        second = None

        for i in range(nmb):
            objective += F.l1_loss(classifs[i].view(-1) + 1e-16, targets[i].view(-1))
            slice = np.random.randint(0, len(features[i]))
            if first is None:
                first = features[i][:slice]
                second = features[i][slice:]
            else:
                first = torch.cat((first, features[i][:slice]), 0)
                second = torch.cat((second, features[i][slice:]), 0)
        if len(first) > 1 and len(second) > 1:
            penalty = torch.nan_to_num(self.mmd(first, second))
        else:
            penalty = torch.tensor(0)
        objective /= nmb

        self.optimizer.zero_grad()
        (objective + (self.hparams['mmd_gamma']*penalty)).backward()
        self.optimizer.step()

        if torch.is_tensor(penalty):
            penalty = penalty.item()

        return {'loss': objective.item(), 'penalty': penalty}


class CausIRL_MMD(AbstractCausIRL):
    '''Causality based invariant representation learning algorithm using the MMD distance from (https://arxiv.org/abs/2206.11646)'''
    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(CausIRL_MMD, self).__init__(input_shape, num_classes, num_domains,
                                  hparams, gaussian=True)


class CausIRL_CORAL(AbstractCausIRL):
    '''Causality based invariant representation learning algorithm using the CORAL distance from (https://arxiv.org/abs/2206.11646)'''
    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(CausIRL_CORAL, self).__init__(input_shape, num_classes, num_domains,
                                  hparams, gaussian=False)
