import random
from contextlib import contextmanager
import torch
import torch.nn as nn
import pdb
from sklearn.cluster import KMeans
import numpy as np
import faiss
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import DBSCAN

from sklearn.decomposition import PCA

class Faiss_Clustering:
    def __init__(self, X, n_clusters, n_init=1, max_iter=300):
        self.n_clusters = n_clusters
        self.n_init = n_init
        self.max_iter = max_iter
        self.kmeans = None
        self.cluster_centers_ = None
        self.inertia_ = None
        self.X = X
        self.kmeans = faiss.Kmeans(
            d=self.X.shape[1], k=self.n_clusters, niter=self.max_iter, nredo=self.n_init, gpu=False
        )

    def fit(self):
        self.kmeans.train(self.X)
        self.cluster_centers_ = self.kmeans.centroids
        self.inertia_ = self.kmeans.obj[-1]
        self.stats = self.kmeans.iteration_stats[-1]

    def predict(self, X):
        return self.kmeans.index.search(X, 1)[1]

def deactivate_sb(m):
    if type(m) == SB:
        m.set_activation_status(False)


def activate_sb(m):
    if type(m) == SB:
        m.set_activation_status(True)


def random_sb(m):
    if type(m) == SB:
        m.update_mix_method('random')


def crossdomain_sb(m):
    if type(m) == SB:
        m.update_mix_method('crossdomain')


@contextmanager
def run_without_sb(model):
    # Assume EFDMix was initially activated
    try:
        model.apply(deactivate_sb)
        yield
    finally:
        model.apply(activate_sb)


@contextmanager
def run_with_sb(model, mix=None):
    # Assume EFDMix was initially deactivated
    if mix == 'random':
        model.apply(random_sb)

    elif mix == 'crossdomain':
        model.apply(crossdomain_sb)

    try:
        model.apply(activate_sb)
        yield
    finally:
        model.apply(deactivate_sb)


class SB(nn.Module):
    """EFDMix.

    Reference:
      Zhang et al, Exact Feature Distribution Matching for Arbitrary Style Transfer and Domain Generalization, CVPR2022
    """

    def __init__(self, p=0.5, alpha=0.1, eps=1e-6, mix='random'):
        """
        Args:
          p (float): probability of using MixStyle.
          alpha (float): parameter of the Beta distribution.
          eps (float): scaling parameter to avoid numerical issues.
          mix (str): how to mix.
        """
        super().__init__()
        self.p = p
        self.beta = torch.distributions.Beta(alpha, alpha)
        self.eps = eps
        self.alpha = alpha
        self.mix = mix
        self._activated = True

    def __repr__(self):
        return f'EFDMix(p={self.p}, alpha={self.alpha}, eps={self.eps}, mix={self.mix})'

    def set_activation_status(self, status=True):
        self._activated = status

    def update_mix_method(self, mix='random'):
        self.mix = mix

    def forward(self, x, domain, label, cluster_idx=None, cluster=False, statistic=None, alpha=None, val2=False):
        if not self.training or not self._activated:
            # pdb.set_trace()
            return x, domain, cluster_idx, cluster, statistic

        # if not self._activated:
        #     return x, domain, cluster_idx, cluster, statistic

        if cluster == True or random.random() > self.p:
            # pdb.set_trace()
            return x, domain, cluster_idx, cluster, statistic

        # pdb.set_trace()
        B = x.size(0)
        C = x.size(1)
        W = x.size(2)
        H = x.size(3)

        # 1) My code

        mu = x.mean(dim=[2, 3], keepdim=True)
        var = x.var(dim=[2, 3], keepdim=True)
        sig = (var + self.eps).sqrt()

        n_cluster = 6
        num_class = 7
        # pdb.set_trace()

        if self.training:

            # # 2) Faiss
            #
            # n_cluster = 3
            #
            # mu_sig = torch.squeeze(torch.cat((mu, sig), dim=1))
            #
            # mu_sig_copy = mu_sig.cpu().detach().numpy().reshape(B, 2 * C)
            # #sig_copy = sig.cpu().detach().numpy().reshape(B, C)
            #
            # # mu_copy = mu.clone().detach().reshape(B, C)
            # # sig_copy = sig.clone().detach().reshape(B, C)
            # #
            # # res = faiss.StandardGpuResources()
            # # flat_config = faiss.GpuIndexFlatConfig()
            # # flat_config.device = 0
            # # index = faiss.GpuIndexFlatL2(res, 2*C, flat_config)
            # # index.add(mu_sig_copy)
            #
            #
            # mu_clustering = Faiss_Clustering(mu_sig_copy, n_cluster)
            # #sig_clustering = Faiss_Clustering(sig_copy, n_cluster)
            #
            # mu_clustering.fit()
            # #sig_clustering.fit()
            #
            #
            # _, cluster_idx = mu_clustering.kmeans.index.search(mu_sig_copy.copy(order="C"), 1)
            # #_, ids_sig = sig_clustering.kmeans.index.search(sig_copy.copy(order="C"), 1)
            #
            # domain = cluster_idx.reshape(-1)
            ############################################################################################
            # 3) with domain label

            mu_mean_list = []
            sig_mean_list = []
            mu_sig_list = []
            sig_sig_list = []
            mu_cov_list = []
            sig_cov_list = []

            if torch.is_tensor(domain):
                domain = domain.cpu().numpy()
                label = label.cpu().numpy()

            real_domain = []
            for i in range(n_cluster):
                if i in domain:
                    real_domain.append(i)

            # pdb.set_trace()

            for i in real_domain:
                # if i ==0:
                #     mu_mean = mu[np.where(domain==i)].mean(dim=[0], keepdim=True)
                #     sig_mean = sig[np.where(domain == i)].mean(dim=[0], keepdim=True)
                # else:
                #     mu_mean += mu[np.where(domain == i)].mean(dim=[0], keepdim=True)
                #     sig_mean += sig[np.where(domain == i)].mean(dim=[0], keepdim=True)
                mu_mean_list.append(mu[np.where(domain == i)].mean(dim=[0], keepdim=True))
                sig_mean_list.append(sig[np.where(domain == i)].mean(dim=[0], keepdim=True))
                mu_sig_list.append((mu[np.where(domain == i)].var(dim=[0], keepdim=True) + self.eps).sqrt())
                sig_sig_list.append((sig[np.where(domain == i)].var(dim=[0], keepdim=True) + self.eps).sqrt())

            # if cluster is None:
            #    print('mu_mean', torch.norm(sig_sig_list[0]), torch.norm(mu_sig_list[0]))

            # mu_mean = mu_mean/n_cluster
            # sig_mean = sig_mean/n_cluster

            mu_mean_list = torch.vstack(mu_mean_list)
            sig_mean_list = torch.vstack(sig_mean_list)
            mu_sig_list = torch.vstack(mu_sig_list)
            sig_sig_list = torch.vstack(sig_sig_list)

            mu_mean = mu_mean_list.mean(dim=[0], keepdim=True)
            sig_mean = sig_mean_list.mean(dim=[0], keepdim=True)
            # mu_mean = torch.vstack(mu_mean_list).mean(dim=[0], keepdim=True)
            # sig_mean = torch.vstack(sig_mean_list).mean(dim=[0], keepdim=True)

            # outlier = []
            # for i in range(n_cluster):
            #     train_statistic = torch.squeeze(torch.cat((mu_mean_list.detach(), sig_mean_list.detach()), dim=1))
            #     original = torch.squeeze(torch.cat((mu[np.where(domain == i)].detach(), sig[np.where(domain == i)].detach()), dim=1))
            #
            #     norm_t = torch.norm(train_statistic[i] - original, dim=1)
            #     max_idx = torch.argmax(norm_t)
            #     outlier.append(max_idx)

            # pdb.set_trace()

            # ##################################### 2) Naive Mean #######################################################
            # mu_mean = mu.mean(dim=[0], keepdim=True)
            # sig_mean = sig.mean(dim=[0], keepdim=True)

            # print(torch.norm(sig_mean - sig_mean1))

            # ############################################################################################
            #
            #
            #
            # sig_mu = torch.mean((mu_mean_list - mu_mean)**2, 0, True)
            # sig_sig = torch.mean((sig_mean_list - sig_mean)**2, 0, True)

            sig_mu = torch.mean((mu - mu_mean) ** 2, 0, True)
            sig_sig = torch.mean((sig - sig_mean) ** 2, 0, True)

            eps_mu = torch.randn_like(mu)  # 2) epsilon_bc_dim
            eps_sig = torch.randn_like(sig)

            # 1) Naive
            # mu_new = mu + eps_mu * torch.sqrt(sig_mu + self.eps)
            # sig_new = sig + eps_sig * torch.sqrt(sig_sig + self.eps)

            # mu_new = mu_mean + eps_mu * torch.sqrt(sig_mu + self.eps)
            # sig_new = sig_mean + eps_sig * torch.sqrt(sig_sig + self.eps)

            # GMM_p = random.random()
            # if GMM_p < 0.5:
            #     mu_new = mu + eps_mu * torch.sqrt(sig_mu + self.eps)
            #     sig_new = sig + eps_sig * torch.sqrt(sig_sig + self.eps)
            # else:
            #     mu_new = mu_mean + eps_mu * torch.sqrt(sig_mu + self.eps)
            #     sig_new = sig_mean + eps_sig * torch.sqrt(sig_sig + self.eps)

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

            GMM_p = random.random()

            # mu_batch = []
            # sig_batch = []
            # factor1 = 2
            # factor2 = 2
            # mu_a = mu_mean
            # sig_a = sig_mean
            # mu_b = torch.sqrt(sig_mu + self.eps)
            # sig_b = torch.sqrt(sig_sig + self.eps)
            #
            # # ra = np.random.choice(3,B)
            # for i in range(B):
            #
            #
            #
            #     uniform = torch.rand_like(mu_a).to(x.device)
            #     mu_c = (mu_a - mu_b * factor1) * (1 - uniform) + (mu_a + mu_b * factor1) * uniform
            #
            #
            #     #uniform = torch.rand_like(sig_a).to(x.device)
            #     sig_c = (sig_a - sig_b * factor2) * ((sig_a - sig_b * factor2) > 0) * (1 - uniform) + (
            #                 sig_a + sig_b * factor2) * uniform
            #
            #     mu_batch.append(mu_c.resize(1, C, 1, 1))
            #     sig_batch.append(sig_c.resize(1, C, 1, 1))
            #
            #
            # mu_new = torch.vstack(mu_batch) #+ eps_mu  * torch.sqrt(sig_mu + self.eps)
            # sig_new = torch.vstack(sig_batch) #+ eps_sig  * torch.sqrt(sig_sig + self.eps)
            # sig_new = sig_new * (sig_new>0)
            #
            # statistic = [mu_a.detach(), mu_b.detach() , sig_a.detach(), sig_b.detach()]

            # if self.training and random.random() > self.p:
            #     # pdb.set_trace()
            #     return x, domain, cluster_idx, cluster, statistic

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

            mu_batch = []
            sig_batch = []
            factor1 = 2
            factor2 = 2
            domain_new = []

            ra = np.random.choice(len(real_domain), B)
            #############################################################################################################
            # ra = np.zeros(64)

            balance_samples = []
            ##################### per class#########################################
            # for j in range(num_class):
            #     num_samples_domain = []
            #     balance_samples_per_class = []
            #     num_samples_class = int((label == j).sum())
            #     if num_samples_class >= len(real_domain):
            #         num_samples_per_class = num_samples_class // len(real_domain)
            #         for k in range(len(real_domain)):
            #
            #             if k == len(real_domain) - 1:
            #                 num_samples_domain.append(num_samples_class - num_samples_per_class * k)
            #             else:
            #                 num_samples_domain.append(num_samples_per_class)
            #
            #             balance_samples_per_class.append(int(
            #                 num_samples_domain[k] - np.sum(np.logical_and(label == j, domain == real_domain[k]))))
            #
            #     else:
            #         for k in range(len(real_domain)):
            #             balance_samples_per_class.append(0)
            #
            #     # pdb.set_trace()
            #     balance_samples.append(np.array(balance_samples_per_class))
            ###############################################################################
            # pdb.set_trace()

            # balance_samples = np.array(balance_samples)

            num_samples = B // len(real_domain)
            num_samples_cluster = []
            balance_samples = []

            for k in range(len(real_domain)):

                if k == len(real_domain) -1:
                    num_samples_cluster.append(B - num_samples * k)
                else:
                    num_samples_cluster.append(num_samples)

                balance_samples.append(num_samples_cluster[k] - (domain == real_domain[k]).sum())

            # num_samples_cluster = np.array(num_samples_cluster)
            balance_samples = np.array(balance_samples)

            #
            # for k in range(n_cluster):
            #     ra[num_samples*(k):num_samples*(k+1)] = k
            # np.random.shuffle(ra)
            # #pdb.set_trace()

            #############################################################################################################
            ################################ Uniform balancing  ################################
            # mu, sig = mu.detach(), sig.detach()
            # for i in range(ra.shape[0]):
            #
            #     if ra[i] == 0:
            #         mu_batch.append(mu[i].resize(1, C, 1, 1))
            #         sig_batch.append(sig[i].resize(1, C, 1, 1))
            #         domain_new.append(domain[i])
            #
            #
            #
            #     elif ra[i] == 1:
            #         mu_a = mu_mean_list[(domain[i] + 1) % 3]
            #         mu_b = mu_sig_list[(domain[i] + 1) % 3]
            #         uniform = torch.rand_like(mu_a).to(x.device)
            #         mu_c = (mu_a - mu_b * factor1) * (1 - uniform) + (mu_a + mu_b * factor1) * uniform
            #
            #         sig_a = sig_mean_list[(domain[i] + 1) % 3]  # mu_mean_list
            #         sig_b = sig_sig_list[(domain[i] + 1) % 3]  # mu_sig_list
            #         #uniform = torch.rand_like(sig_a).to(x.device)
            #         sig_c = (sig_a - sig_b * factor2) * ((sig_a - sig_b * factor2) > 0) * (1 - uniform) + (
            #                     sig_a + sig_b * factor2) * uniform
            #
            #         vector_mu = mu_c.clone() - mu[i].clone().detach()
            #         vector_sig = sig_c.clone() - sig[i].clone().detach()
            #
            #         domain_new.append(((domain[i] + 1) % 3))
            #         mu_batch.append((mu[i] + vector_mu).resize(1, C, 1, 1))
            #         sig_batch.append((sig[i] + vector_sig).resize(1, C, 1, 1))
            #
            #         # domain_new.append(((domain[i] + 1) % 3))
            #         # mu_batch.append((mu_c).resize(1, C, 1, 1))
            #         # sig_batch.append((sig_c).resize(1, C, 1, 1))
            #
            #
            #     else:
            #
            #         mu_a = mu_mean_list[(domain[i] + 2) % 3]
            #         mu_b = mu_sig_list[(domain[i] + 2) % 3]
            #         uniform = torch.rand_like(mu_a).to(x.device)
            #         mu_c = (mu_a - mu_b * factor1) * (1 - uniform) + (mu_a + mu_b * factor1) * uniform
            #
            #         sig_a = sig_mean_list[(domain[i] + 2) % 3]  # mu_mean_list
            #         sig_b = sig_sig_list[(domain[i] + 2) % 3]  # mu_sig_list
            #         #uniform = torch.rand_like(sig_a).to(x.device)
            #         sig_c = (sig_a - sig_b * factor2) * ((sig_a - sig_b * factor2) > 0) * (1 - uniform) + (
            #                     sig_a + sig_b * factor2) * uniform
            #
            #         vector_mu = mu_c.clone() - mu[i].clone().detach()
            #         vector_sig = sig_c.clone() - sig[i].clone().detach()
            #
            #         domain_new.append(((domain[i] + 2) % 3))
            #         mu_batch.append((mu[i] + vector_mu).resize(1, C, 1, 1))
            #         sig_batch.append((sig[i] + vector_sig).resize(1, C, 1, 1))
            #
            #         # domain_new.append(((domain[i] + 2) % 3))
            #         # mu_batch.append((mu_c).resize(1, C, 1, 1))
            #         # sig_batch.append((sig_c).resize(1, C, 1, 1))
            #
            #
            # mu_new = torch.vstack(mu_batch)  # + eps_mu  * torch.sqrt(sig_mu + self.eps)
            # sig_new = torch.vstack(sig_batch)  # + eps_sig  * torch.sqrt(sig_sig + self.eps)
            # sig_new = sig_new * (sig_new > 0)
            # cluster = True
            # x_normed = (x - mu) / sig
            # x_new = x_normed * sig_new + mu_new

            ########################Gaussian balancing#####################################################################################

            # mu, sig = mu.detach(), sig.detach()

            # for i in range(ra.shape[0]):
            #
            #     if ra[i] == 0:
            #         mu_batch.append(mu[i].resize(1, C, 1, 1))
            #         sig_batch.append(sig[i].resize(1, C, 1, 1))
            #         domain_new.append(domain[i])
            #
            #
            #
            #     elif ra[i] == 1:
            #
            #         sampling = np.random.choice(np.where(domain == ((domain[i] + 1) % 3))[0], 1, False)
            #         epsilon1 = torch.randn_like(mu[sampling]).to(x.device)
            #         epsilon2 = torch.randn_like(mu[sampling]).to(x.device)
            #
            #         mu_c = mu_mean_list[(domain[i] + 1) % 3] + mu_sig_list[(domain[i] + 1) % 3] * epsilon1
            #         sig_c = sig_mean_list[(domain[i] + 1) % 3] + sig_sig_list[(domain[i] + 1) % 3] * epsilon2
            #
            #         # vector_mu = mu_c.clone() - mu[i].clone().detach()
            #         # vector_sig = sig_c.clone() - sig[i].clone().detach()
            #         #
            #         # domain_new.append(((domain[i] + 1) % 3))
            #         # mu_batch.append((mu[i] + vector_mu).resize(1, C, 1, 1))
            #         # sig_batch.append((sig[i] + vector_sig).resize(1, C, 1, 1))
            #
            #         domain_new.append(((domain[i] + 1) % 3))
            #         mu_batch.append((mu_c).resize(1, C, 1, 1))
            #         sig_batch.append((sig_c).resize(1, C, 1, 1))
            #
            #
            #     else:
            #
            #         sampling = np.random.choice(np.where(domain == ((domain[i] + 2) % 3))[0], 1, False)
            #         epsilon1 = torch.randn_like(mu[sampling]).to(x.device)
            #         epsilon2 = torch.randn_like(mu[sampling]).to(x.device)
            #
            #         mu_c = mu_mean_list[(domain[i] + 2) % 3] + mu_sig_list[(domain[i] + 2) % 3] * epsilon1
            #         sig_c = sig_mean_list[(domain[i] + 2) % 3] + sig_sig_list[(domain[i] + 2) % 3] * epsilon2
            #
            #         # vector_mu = mu_c.clone() - mu[i].clone().detach()
            #         # vector_sig = sig_c.clone() - sig[i].clone().detach()
            #         #
            #         # domain_new.append(((domain[i] + 2) % 3))
            #         # mu_batch.append((mu[i] + vector_mu).resize(1, C, 1, 1))
            #         # sig_batch.append((sig[i] + vector_sig).resize(1, C, 1, 1))
            #
            #         domain_new.append(((domain[i] + 2) % 3))
            #         mu_batch.append((mu_c).resize(1, C, 1, 1))
            #         sig_batch.append((sig_c).resize(1, C, 1, 1))
            #
            # mu_new = torch.vstack(mu_batch)  # + eps_mu  * torch.sqrt(sig_mu + self.eps)
            # sig_new = torch.vstack(sig_batch)  # + eps_sig  * torch.sqrt(sig_sig + self.eps)
            # sig_new = sig_new * (sig_new > 0)
            # cluster = True
            # x_normed = (x - mu) / sig
            # x_new = x_normed * sig_new + mu_new

            ######################################################################################################
            ### Remove redudant data

            # mu, sig = mu.detach(), sig.detach()
            concat_stat = torch.squeeze(torch.cat((mu, sig), dim=1))
            concat_stat_m = torch.squeeze(torch.cat((mu_mean_list, sig_mean_list), dim=1))

            remove_ind = []
            domain_replace = []
            domain_new = domain.copy()

            for k in range(len(real_domain)):
                if balance_samples[k] > 0:
                    for j in range(balance_samples[k]):
                        domain_replace.append(real_domain[k])

                # domain_replace.append(domain_replace_perclass)
            # pdb.set_trace()

            # for k in range(n_cluster):
            #     if balance_samples[k] < 0:
            #         dis_vec = torch.norm(concat_stat_m[k] - concat_stat[domain == k], dim=1)
            #         _, min_ind = torch.topk(dis_vec, np.abs(balance_samples[k]))
            #         if np.abs(balance_samples[k]) == 1:
            #             remove_ind.append(np.where(domain == k)[0][min_ind.cpu()].tolist())
            #             domain_new[np.where(domain == k)[0][min_ind.cpu()]] = domain_replace[
            #                                                                   index]
            #         else:
            #             remove_ind += np.where(domain == k)[0][min_ind.cpu()].tolist()
            #             domain_new[np.where(domain == k)[0][min_ind.cpu()]] = domain_replace[
            #                                                                   index:index + min_ind.size(0)]
            #
            #         index += min_ind.size(0)



            index = 0


            for k in range(len(real_domain)):
                if balance_samples[k] < 0:
                    cor_ind = (domain == real_domain[k])
                    dis_matrix = torch.cdist(concat_stat[cor_ind], concat_stat[cor_ind], p=2)
                    dis_matrix.fill_diagonal_(np.inf)

                    for n in range(np.abs(balance_samples[k])):
                        dis_matrix_1d = dis_matrix.view(-1)
                        min_ind = torch.argmin(dis_matrix_1d)

                        point1, point2 = int(min_ind) // (domain == real_domain[k]).sum(), int(
                            min_ind) % (domain == real_domain[k]).sum()

                        values1, _ = torch.topk(dis_matrix[point1, :], 2, largest=False)

                        values2, _ = torch.topk(dis_matrix[point2, :], 2, largest=False)

                        if values1[1] > values2[1]:
                            point = point2
                        else:
                            point = point1

                        remove_ind.append(
                            np.where(domain == real_domain[k])[0][point])

                        domain_new[np.where(domain == real_domain[k])[0][point]] = \
                        domain_replace[index]

                        dis_matrix[point, :] = np.inf
                        dis_matrix[:, point] = np.inf
                        index += 1

            #pdb.set_trace()

            ###
            # EFDM
            x_new_batch = []
            x_view = x.view(B, C, -1)
            value_x, index_x = torch.sort(x_view)  # sort inputs
            inverse_index = index_x.argsort(-1)

            for i in range(ra.shape[0]):

                if i in remove_ind:
                    #pdb.set_trace()
                    ################################################################################################
                    # gaussian

                    # mu, sig = mu.detach(), sig.detach()
                    # epsilon1 = torch.randn_like(mu[i]).to(x.device)
                    # epsilon2 = torch.randn_like(mu[i]).to(x.device)
                    #
                    # mu_c = mu_mean_list[domain_new[i]] + mu_sig_list[domain_new[i]] * epsilon1
                    # sig_c = sig_mean_list[domain_new[i]] + sig_sig_list[domain_new[i]] * epsilon2
                    #
                    # # vector_mu = mu_c.clone().detach() - mu[i].clone().detach()
                    # # vector_sig = sig_c.clone().detach() - sig[i].clone().detach()
                    # #
                    # # mu_batch.append((mu[i] + vector_mu).resize(1, C, 1, 1))
                    # # sig_batch.append((sig[i] + vector_sig).resize(1, C, 1, 1))
                    #
                    # mu_batch.append((mu_c).resize(1, C, 1, 1))
                    # sig_batch.append((sig_c).resize(1, C, 1, 1))

                    ################################################################################################
                    # mix
                    # mu, sig = mu.detach(), sig.detach()
                    # lmda = self.beta.sample((1, 1, 1, 1))
                    # #lmda = self.beta.sample((1, C, 1, 1))
                    # lmda = lmda.to(x.device)
                    # try:
                    #     mix_a, mix_b = np.random.choice(np.where(domain == domain_new[i])[0], 2, False)
                    #     mu1, sig1 = mu[mix_a], sig[mix_a]
                    #     mu2, sig2 = mu[mix_b], sig[mix_b]
                    # except:
                    #     mix_a = np.random.choice(np.where(domain == domain_new[i])[0], 1, False)
                    #     mu1, sig1 = mu[mix_a], sig[mix_a]
                    #     mu2, sig2 = mu[mix_a] + self.eps, sig[mix_a] + self.eps
                    #     #pdb.set_trace()
                    #
                    # mu_c = mu1 * lmda + mu2 * (1 - lmda)
                    # sig_c = sig1 * lmda + sig2 * (1 - lmda)
                    #
                    #
                    # mu_batch.append((mu_c).resize(1, C, 1, 1))
                    # sig_batch.append((sig_c).resize(1, C, 1, 1))
                    ###############################################################################################
                    # uniform
                    # mu_a = mu_mean_list[domain_new[i]]
                    # mu_b = mu_sig_list[domain_new[i]]
                    # uniform = torch.rand_like(mu_a).to(x.device)
                    # mu_c = (mu_a - mu_b * factor1) * (1 - uniform) + (mu_a + mu_b * factor1) * uniform

                    # sig_a = sig_mean_list[domain_new[i]]  # mu_mean_list
                    # sig_b = sig_sig_list[domain_new[i]]  # mu_sig_list
                    # uniform = torch.rand_like(sig_a).to(x.device)
                    # sig_c = (sig_a - sig_b * factor2) * ((sig_a - sig_b * factor2) > 0) * (1 - uniform) + (
                    #         sig_a + sig_b * factor2) * uniform
                    #
                    # # vector_mu = mu_c.clone() - mu[i].clone().detach()
                    # # vector_sig = sig_c.clone() - sig[i].clone().detach()
                    # #
                    # # domain_new.append(((domain[i] + 1) % 3))
                    # # mu_batch.append((mu[i] + vector_mu).resize(1, C, 1, 1))
                    # # sig_batch.append((sig[i] + vector_sig).resize(1, C, 1, 1))
                    #
                    # mu_batch.append((mu_c).resize(1, C, 1, 1))
                    # sig_batch.append((sig_c).resize(1, C, 1, 1))
                    ################################################################################################
                    # EFDM
                    lmda = self.beta.sample((1, 1, 1))
                    # lmda = self.beta.sample((1, C, 1))
                    lmda = lmda.to(x.device)

                    # try:
                    #     mix_a, mix_b = np.random.choice(np.where(domain == domain_new[i])[0], 2, False)
                    #
                    #     x_view_copy = value_x[mix_a[0]].gather(-1, inverse_index[i])
                    #     x_view_copy_2 = value_x[mix_b[0]].gather(-1, inverse_index[i])
                    #     new_x = x_view_copy_2 + (x_view_copy - x_view_copy_2.detach()) * (1 - lmda)
                    #     #new_x = x_view_copy_2 + (x_view_copy - x_view_copy_2) * (1 - lmda)
                    #
                    #     x_new_batch.append(new_x.view(1, C, W, H))
                    # except:
                    #     mix_a = np.random.choice(np.where(domain == domain_new[i])[0], 1, False)
                    #     # pdb.set_trace()
                    #     x_view_copy = value_x[mix_a[0]].gather(-1, inverse_index[i])
                    #     x_new_batch.append(x_view_copy.view(1, C, W, H))

                    try:
                        mix_a, mix_b = np.random.choice(np.where(domain == domain_new[i])[0], 2, False)

                        x_view_copy = value_x[mix_a].gather(-1, inverse_index[i])
                        x_view_copy_2 = value_x[mix_b].gather(-1, inverse_index[i])
                        new_x_temp = x_view_copy_2 + (x_view_copy - x_view_copy_2) * (1 - lmda)
                        new_x = x_view[i] * 0.1 + (new_x_temp - x_view[i].detach() * 0.1)  # 0.1

                        # new_x_temp = x_view[i] + (x_view_copy - x_view[i].detach())
                        # new_x_temp_2 = x_view[i] + (x_view_copy_2 - x_view[i].detach())
                        # new_x = new_x_temp + (x_view_copy_2 - new_x_temp) * (1 - lmda)

                        # new_x = x_view[i] + (new_x_temp - x_view[i].detach())
                        x_new_batch.append(new_x.view(1, C, W, H))
                    except:
                        mix_a = np.random.choice(np.where(domain == domain_new[i])[0], 1, False)
                        #new_list = [i for i in range(B) if i != mix_a[0]]
                        #mix_b = np.random.choice(np.array(new_list), 1, False)

                        #pdb.set_trace()

                        x_view_copy = value_x[mix_a[0]].gather(-1, inverse_index[i])
                        #x_view_copy_2 = value_x[i].gather(-1, inverse_index[i])
                        #new_x_temp = x_view_copy_2 + (x_view_copy - x_view_copy_2) * (1 - lmda)
                        new_x = x_view[i]  + (x_view_copy - x_view[i].detach() * lmda)  # 0.1
                        x_new_batch.append(new_x.view(1, C, W, H))

                        # x_view_copy = value_x[mix_a[0]].gather(-1, inverse_index[i])
                        # new_x = x_view[i] * 0.1 + (x_view_copy - x_view[i].detach() * 0.1)
                        # x_new_batch.append(new_x.view(1, C, W, H))



                else:
                    # mu_batch.append(mu[i].resize(1, C, 1, 1))
                    # sig_batch.append(sig[i].resize(1, C, 1, 1))
                    #############################################################
                    x_new_batch.append(x[i].resize(1, C, W, H))

            cluster = True
            #pdb.set_trace()
            x_new = torch.vstack(x_new_batch)  # + eps_mu  * torch.sqrt(sig_mu + self.eps)

            #######################################################################################
            # mu_new = torch.vstack(mu_batch)  # + eps_mu  * torch.sqrt(sig_mu + self.eps)
            # sig_new = torch.vstack(sig_batch)  # + eps_sig  * torch.sqrt(sig_sig + self.eps)
            # sig_new = sig_new * (sig_new > 0)
            # cluster = True
            # x_normed = (x - mu) / sig
            # #pdb.set_trace()
            # x_new = x_normed * sig_new + mu_new

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

            ########################MIX balancing#####################################################################################
            # mu, sig = mu.detach(), sig.detach()
            #  for i in range(ra.shape[0]):
            #
            #      #lmda = self.beta.sample((1, 1, 1, 1))
            #      lmda = self.beta.sample((1, C, 1, 1))
            #      lmda = lmda.to(x.device)
            #
            #      if ra[i] == 0:
            #          mu_batch.append(mu[i].resize(1, C, 1, 1))
            #          sig_batch.append(sig[i].resize(1, C, 1, 1))
            #          #num_samples_cluster[domain[i]] -= 1
            #          domain_new.append(domain[i])
            #
            #
            #      elif ra[i] == 1:
            #
            #          try:
            #              mix_a, mix_b = np.random.choice(np.where(domain == ((domain[i] + 1) % 3))[0], 2, False)
            #              mu1, sig1 = mu[mix_a], sig[mix_a]
            #              mu2, sig2 = mu[mix_b], sig[mix_b]
            #          except:
            #              mix_a = np.random.choice(np.where(domain == ((domain[i] + 1) % 3))[0], 1, False)
            #              mu1, sig1 = mu[mix_a], sig[mix_a]
            #              mu2, sig2 = mu[mix_a] + self.eps, sig[mix_a] + self.eps
            #              #pdb.set_trace()
            #
            #          mu_c = mu1 * lmda + mu2 * (1 - lmda)
            #          sig_c = sig1 * lmda + sig2 * (1 - lmda)
            #
            #          # vector_mu = mu_c.clone() - mu[i].clone().detach()
            #          # vector_sig = sig_c.clone()- sig[i].clone().detach()
            #          #
            #          # domain_new.append(((domain[i]+1) %3))
            #          # mu_batch.append((mu[i]+vector_mu).resize(1, C, 1, 1))
            #          # sig_batch.append((sig[i]+vector_sig).resize(1, C, 1, 1))
            #
            #          domain_new.append(((domain[i] + 1) % 3))
            #          mu_batch.append((mu_c).resize(1, C, 1, 1))
            #          sig_batch.append((sig_c).resize(1, C, 1, 1))
            #
            #
            #      else:
            #          try:
            #              mix_a, mix_b = np.random.choice(np.where(domain == ((domain[i]+2) %3))[0], 2, False)
            #              mu1, sig1 = mu[mix_a], sig[mix_a]
            #              mu2, sig2 = mu[mix_b], sig[mix_b]
            #          except:
            #              mix_a = np.random.choice(np.where(domain == ((domain[i] + 2) % 3))[0], 1, False)
            #              mu1, sig1 = mu[mix_a], sig[mix_a]
            #              mu2, sig2 = mu[mix_a] + self.eps, sig[mix_a] + self.eps
            #              #mu2, sig2 = mu[mix_a] + 1e-3, sig[mix_a] + 1e-3
            #              #pdb.set_trace()
            #
            #
            #          mu_c = mu1 * lmda + mu2 * (1 - lmda)
            #          sig_c = sig1 * lmda + sig2 * (1 - lmda)
            #
            #          # vector_mu = mu_c.clone() - mu[i].clone().detach()
            #          # vector_sig = sig_c.clone() - sig[i].clone().detach()
            #          #
            #          # domain_new.append(((domain[i] + 2) % 3))
            #          # mu_batch.append((mu[i] + vector_mu).resize(1, C, 1, 1))
            #          # sig_batch.append((sig[i] + vector_sig).resize(1, C, 1, 1))
            #
            #          domain_new.append(((domain[i] + 2) % 3))
            #          mu_batch.append((mu_c).resize(1, C, 1, 1))
            #          sig_batch.append((sig_c).resize(1, C, 1, 1))
            #
            #
            #  mu_new = torch.vstack(mu_batch)  # + eps_mu  * torch.sqrt(sig_mu + self.eps)
            #  sig_new = torch.vstack(sig_batch)  # + eps_sig  * torch.sqrt(sig_sig + self.eps)
            #  sig_new = sig_new * (sig_new > 0)
            #  cluster = True
            #  x_normed = (x - mu) / sig
            # # pdb.set_trace()
            #  x_new = x_normed * sig_new + mu_new

            ########################MIX mix balancing#####################################################################################
            # ra = np.random.choice(3, 64)
            # for i in range(ra.shape[0]):
            #
            #     #lmda = self.beta.sample((1, 1, 1, 1))
            #     lmda = self.beta.sample((1, C, 1, 1))
            #     lmda = lmda.to(x.device)
            #
            #     if ra[i] == 0:
            #         mu_batch.append(mu[i].resize(1, C, 1, 1))
            #         sig_batch.append(sig[i].resize(1, C, 1, 1))
            #         num_samples_cluster[domain[i]] -= 1
            #         domain_new.append(domain[i])
            #
            #
            #     elif ra[i] == 1 or ra[i] == 2:
            #         cluster_1, cluster_2 = np.random.choice(range(n_cluster), 2, False)
            #
            #
            #         mix_a = np.random.choice(np.where(domain == cluster_1)[0], 1, False)
            #         mix_b = np.random.choice(np.where(domain == cluster_2)[0], 1, False)
            #         mu1, sig1 = mu[mix_a], sig[mix_a]
            #         mu2, sig2 = mu[mix_b], sig[mix_b]
            #
            #
            #         mu_c = mu1 * lmda + mu2 * (1 - lmda)
            #         sig_c = sig1 * lmda + sig2 * (1 - lmda)
            #
            #         domain_new.append(cluster_1)
            #         mu_batch.append(mu_c.resize(1, C, 1, 1))
            #         sig_batch.append(sig_c.resize(1, C, 1, 1))
            #
            #
            #
            #
            # mu_new = torch.vstack(mu_batch)  # + eps_mu  * torch.sqrt(sig_mu + self.eps)
            # sig_new = torch.vstack(sig_batch)  # + eps_sig  * torch.sqrt(sig_sig + self.eps)
            # sig_new = sig_new * (sig_new > 0)
            # cluster = True
            # x_normed = (x - mu) / sig
            # x_new = x_normed * sig_new + mu_new
            ##############################EFDM balancing###############################################################################
            # x_new_batch = []
            # x_view = x.view(B, C, -1)
            # value_x, index_x = torch.sort(x_view)  # sort inputs
            # inverse_index = index_x.argsort(-1)
            # #pdb.set_trace()
            # for i in range(ra.shape[0]):
            #
            #     lmda = self.beta.sample((1, 1, 1))
            #     # lmda = self.beta.sample((1, C, 1, 1))
            #     lmda = lmda.to(x.device)
            #
            #     if ra[i] == 0:
            #         x_new_batch.append(x[i].resize(1, C, H, W))
            #         domain_new.append(domain[i])
            #
            #
            #     elif ra[i] == 1:
            #
            #         try:
            #             mix_a, mix_b = np.random.choice(np.where(domain == ((domain[i] + 1) % 3))[0], 2, False)
            #
            #             x_view_copy = value_x[mix_a[0]].gather(-1, inverse_index[i])
            #             x_view_copy_2 = value_x[mix_b[0]].gather(-1, inverse_index[i])
            #             new_x = x_view_copy_2 + (x_view_copy - x_view_copy_2.detach()) * (1 - lmda)
            #             x_new_batch.append(new_x.view(1, C, W, H))
            #         except:
            #             mix_a = np.random.choice(np.where(domain == ((domain[i] + 1) % 3))[0], 1, False)
            #             #pdb.set_trace()
            #             x_view_copy = value_x[mix_a[0]].gather(-1, inverse_index[i])
            #             x_new_batch.append(x_view_copy.view(1, C, W, H))
            #
            #
            #
            #         domain_new.append(((domain[i] + 1) % 3))
            #
            #
            #
            #     else:
            #         try:
            #             mix_a, mix_b = np.random.choice(np.where(domain == ((domain[i] + 2) % 3))[0], 2, False)
            #
            #             x_view_copy = value_x[mix_a[0]].gather(-1, inverse_index[i])
            #             x_view_copy_2 = value_x[mix_b[0]].gather(-1, inverse_index[i])
            #             new_x = x_view_copy_2 + (x_view_copy - x_view_copy_2.detach()) * (1 - lmda)
            #             x_new_batch.append(new_x.view(1, C, W, H))
            #         except:
            #             mix_a = np.random.choice(np.where(domain == ((domain[i] + 2) % 3))[0], 1, False)
            #             #pdb.set_trace()
            #             x_view_copy = value_x[mix_a[0]].gather(-1, inverse_index[i])
            #             x_new_batch.append(x_view_copy.view(1, C, W, H))
            #
            #         domain_new.append(((domain[i] + 2) % 3))
            #
            # cluster = True
            # x_new = torch.vstack(x_new_batch)  # + eps_mu  * torch.sqrt(sig_mu + self.eps)

            ######################################Perfect balance#######################################################################
            # mu, sig = mu.detach(), sig.detach()
            #  for i in range(ra.shape[0]):
            #
            #      lmda = self.beta.sample((1, 1, 1, 1))
            #      #lmda = self.beta.sample((1, C, 1, 1))
            #      lmda = lmda.to(x.device)
            #
            #      # if i in max_idx:
            #      #     mu_batch.append(mu[i].resize(1, C, 1, 1))
            #      #     sig_batch.append(sig[i].resize(1, C, 1, 1))
            #      #     num_samples_cluster[domain[i]] -= 1
            #      #     domain_new.append(domain[i])
            #      #     continue
            #
            #      if ra[i] == 0:
            #          if num_samples_cluster[domain[i]] <= 0:
            #              domain_aux = np.random.choice(np.where(num_samples_cluster != 0)[0], 1, False)
            #
            #              ######################################################################################
            #              try:
            #                  mix_a, mix_b = np.random.choice(np.where(domain == (domain_aux))[0], 2, False)
            #                  mu1, sig1 = mu[mix_a], sig[mix_a]
            #                  mu2, sig2 = mu[mix_b], sig[mix_b]
            #              except:
            #                  mix_a = np.random.choice(np.where(domain == (domain_aux))[0], 1, False)
            #                  mu1, sig1 = mu[mix_a], sig[mix_a]
            #                  mu2, sig2 = mu[mix_a] + self.eps, sig[mix_a] + self.eps
            #              mu_c = mu1 * lmda + mu2 * (1 - lmda)
            #              sig_c = sig1 * lmda + sig2 * (1 - lmda)
            #              ######################################################################################
            #              # epsilon1 = torch.randn_like(mu[i]).to(x.device)
            #              # epsilon2 = torch.randn_like(mu[i]).to(x.device)
            #              #
            #              # mu_c = mu_mean_list[domain_aux] + mu_sig_list[domain_aux] * epsilon1
            #              # sig_c = sig_mean_list[domain_aux] + sig_sig_list[domain_aux] * epsilon2
            #
            #              # mu_c = mu_mean + sig_mu * epsilon1
            #              # sig_c = sig_mean + sig_sig * epsilon2
            #
            #              ######################################################################################
            #
            #
            #
            #              domain_new.append(int(domain_aux))
            #              mu_batch.append(mu_c.resize(1, C, 1, 1))
            #              sig_batch.append(sig_c.resize(1, C, 1, 1))
            #              num_samples_cluster[domain_aux] -= 1
            #
            #
            #          else:
            #              mu_batch.append(mu[i].resize(1, C, 1, 1))
            #              sig_batch.append(sig[i].resize(1, C, 1, 1))
            #              num_samples_cluster[domain[i]] -= 1
            #              domain_new.append(domain[i])
            #
            #
            #      elif ra[i] == 1:
            #          if num_samples_cluster[(domain[i] + 1) % 3] <= 0:
            #              domain_aux = np.random.choice(np.where(num_samples_cluster != 0)[0], 1, False)
            #          else:
            #              domain_aux = (domain[i] + 1) % 3
            #
            #          num_samples_cluster[domain_aux] -= 1
            #
            #          ######################################################################################
            #          try:
            #              mix_a, mix_b = np.random.choice(np.where(domain == domain_aux)[0], 2, False)
            #              mu1, sig1 = mu[mix_a], sig[mix_a]
            #              mu2, sig2 = mu[mix_b], sig[mix_b]
            #          except:
            #              mix_a = np.random.choice(np.where(domain == domain_aux)[0], 1, False)
            #              mu1, sig1 = mu[mix_a], sig[mix_a]
            #              mu2, sig2 = mu[mix_a] + self.eps, sig[mix_a] + self.eps
            #              # pdb.set_trace()
            #
            #          mu_c = mu1 * lmda + mu2 * (1 - lmda)
            #          sig_c = sig1 * lmda + sig2 * (1 - lmda)
            #          ######################################################################################
            #
            #          # epsilon1 = torch.randn_like(mu[i]).to(x.device)
            #          # epsilon2 = torch.randn_like(mu[i]).to(x.device)
            #          #
            #          # mu_c = mu_mean_list[domain_aux] + mu_sig_list[domain_aux] * epsilon1
            #          # sig_c = sig_mean_list[domain_aux] + sig_sig_list[domain_aux] * epsilon2
            #
            #          # mu_c = mu_mean + sig_mu * epsilon1
            #          # sig_c = sig_mean + sig_sig * epsilon2
            #          ######################################################################################
            #
            #
            #
            #          domain_new.append(int(domain_aux))
            #          mu_batch.append(mu_c.resize(1, C, 1, 1))
            #          sig_batch.append(sig_c.resize(1, C, 1, 1))
            #
            #
            #      else:
            #          if num_samples_cluster[(domain[i] + 2) % 3] <= 0:
            #              domain_aux = np.random.choice(np.where(num_samples_cluster != 0)[0], 1, False)
            #          else:
            #              domain_aux = (domain[i] + 2) % 3
            #          num_samples_cluster[domain_aux] -= 1
            #          ######################################################################################
            #          try:
            #              mix_a, mix_b = np.random.choice(np.where(domain == domain_aux)[0], 2, False)
            #              mu1, sig1 = mu[mix_a], sig[mix_a]
            #              mu2, sig2 = mu[mix_b], sig[mix_b]
            #          except:
            #              mix_a = np.random.choice(np.where(domain == domain_aux)[0], 1, False)
            #              mu1, sig1 = mu[mix_a], sig[mix_a]
            #              mu2, sig2 = mu[mix_a] + self.eps, sig[mix_a] + self.eps
            #              # pdb.set_trace()
            #
            #          mu_c = mu1 * lmda + mu2 * (1 - lmda)
            #          sig_c = sig1 * lmda + sig2 * (1 - lmda)
            #          ######################################################################################
            #
            #          # epsilon1 = torch.randn_like(mu[i]).to(x.device)
            #          # epsilon2 = torch.randn_like(mu[i]).to(x.device)
            #          #
            #          # mu_c = mu_mean_list[domain_aux] + mu_sig_list[domain_aux] * epsilon1
            #          # sig_c = sig_mean_list[domain_aux] + sig_sig_list[domain_aux] * epsilon2
            #
            #          # mu_c = mu_mean + sig_mu * epsilon1
            #          # sig_c = sig_mean + sig_sig * epsilon2
            #          ######################################################################################
            #
            #          domain_new.append(int(domain_aux))
            #          mu_batch.append(mu_c.resize(1, C, 1, 1))
            #          sig_batch.append(sig_c.resize(1, C, 1, 1))
            #
            #
            #
            #  mu_new = torch.vstack(mu_batch)  # + eps_mu  * torch.sqrt(sig_mu + self.eps)
            #  sig_new = torch.vstack(sig_batch)  # + eps_sig  * torch.sqrt(sig_sig + self.eps)
            #  sig_new = sig_new * (sig_new > 0)
            #  cluster = True
            #  x_normed = (x - mu) / sig
            # # pdb.set_trace()
            #  x_new = x_normed * sig_new + mu_new

            ########################################################################################################################
            # a = np.array(domain_new)
            # b = (a==0).sum()
            # c = (a == 1).sum()
            # d = (a == 2).sum()
            # pdb.set_trace()

            # mu_1_idx = np.random.choice(np.where(domain == 0)[0], 20, True)
            # mu_2_idx = np.random.choice(np.where(domain == 1)[0], 20, True)
            # mu_3_idx = np.random.choice(np.where(domain == 2)[0], 20, True)
            #
            # mu_out = torch.vstack([mu[mu_1_idx], mu[mu_2_idx], mu[mu_3_idx]])
            # sig_out = torch.vstack([sig[mu_1_idx], sig[mu_2_idx], sig[mu_3_idx]])

            # statistic = [mu_mean_list.detach(), sig_mean_list.detach()]  # , sig_a.detach(), sig_b.detach()]

            statistic = [mu, sig]

            return x_new, np.array(domain_new), cluster_idx, cluster, statistic

            # lmda2 = self.beta.sample((B, 1, 1, 1))
            # lmda2 = lmda2.to(x.device)
            # perm = torch.randperm(B)
            #
            # mu2, sig2 = mu_new[perm], sig_new[perm]
            # mu_mix = mu_new * lmda2 + mu2 * (1 - lmda2)
            # sig_mix = sig_new * lmda2 + sig2 * (1 - lmda2)
            #
            # statistic = [mu_mean_list.detach(), sig_mean_list.detach()] #, sig_a.detach(), sig_b.detach()]
            #
            # x_normed = (x - mu) / sig
            # x_new = x_normed * sig_mix + mu_mix
            #
            # return x_new, domain, cluster_idx, cluster, statistic


        else:
            if statistic is None:

                #############################################################################################################
                # 2) Faiss

                # n_cluster = 6
                #
                # mu_sig = torch.squeeze(torch.cat((mu, sig), dim=1))
                #
                # mu_sig_copy = mu_sig.cpu().detach().numpy().reshape(B, 2 * C)
                # # sig_copy = sig.cpu().detach().numpy().reshape(B, C)
                #
                # # mu_copy = mu.clone().detach().reshape(B, C)
                # # sig_copy = sig.clone().detach().reshape(B, C)
                # #
                # # res = faiss.StandardGpuResources()
                # # flat_config = faiss.GpuIndexFlatConfig()
                # # flat_config.device = 0
                # # index = faiss.GpuIndexFlatL2(res, 2*C, flat_config)
                # # index.add(mu_sig_copy)
                #
                # mu_clustering = Faiss_Clustering(mu_sig_copy, n_cluster)
                # # sig_clustering = Faiss_Clustering(sig_copy, n_cluster)
                #
                # mu_clustering.fit()
                # # sig_clustering.fit()
                #
                # _, cluster_idx = mu_clustering.kmeans.index.search(mu_sig_copy.copy(order="C"), 1)
                # # _, ids_sig = sig_clustering.kmeans.index.search(sig_copy.copy(order="C"), 1)
                #
                # # mu_cluster_labels = np.array([int(n[0]) for n in ids_mu])
                # # sig_cluster_labels = np.array([int(n[0]) for n in ids_sig])
                #
                # # mu_list = []
                # # sig_list = []
                # # for i in range(n_cluster):
                # #     if i ==0:
                # #         mu_mean = mu[np.where(ids_mu==i)].mean(dim=[0], keepdim=True)
                # #         sig_mean = sig[np.where(ids_mu == i)].mean(dim=[0], keepdim=True)
                # #     else:
                # #         mu_mean += mu[np.where(ids_mu == i)].mean(dim=[0], keepdim=True)
                # #         sig_mean += sig[np.where(ids_mu == i)].mean(dim=[0], keepdim=True)
                # #     mu_list.append(mu[np.where(ids_mu == i)].mean(dim=[0], keepdim=True))
                # #     sig_list.append(sig[np.where(ids_mu == i)].mean(dim=[0], keepdim=True))
                # #
                # # mu_mean = mu_mean/n_cluster
                # # sig_mean = sig_mean/n_cluster
                # # pdb.set_trace()
                # domain = cluster_idx.reshape(-1)
                #############################################################################################################

                mu_mean_list = []
                sig_mean_list = []

                if torch.is_tensor(domain):
                    domain = domain.cpu().numpy()

                # real_domain = []
                # for i in range(n_cluster):
                #     if i in domain:
                #         real_domain.append(i)

                # for i in range(n_cluster):
                #     if i in real_domain:
                #         mu_mean_list.append(mu[np.where(domain == i)].mean(dim=[0], keepdim=True))
                #         sig_mean_list.append(sig[np.where(domain == i)].mean(dim=[0], keepdim=True))
                #     else:
                #         mu_mean_list.append(mu[np.where(domain == i)].mean(dim=[0], keepdim=True))
                #         sig_mean_list.append(sig[np.where(domain == i)].mean(dim=[0], keepdim=True))
                #
                # mu_mean_list = torch.vstack(mu_mean_list)
                # sig_mean_list = torch.vstack(sig_mean_list)
                #
                # statistic = [mu_mean_list.detach(), sig_mean_list.detach()]  # , sig_a.detach(), sig_b.detach()]

                #############################################################################################################
                statistic = [mu.detach(), sig.detach(),
                             x.detach()]  # mu_sig_list.detach(), sig_sig_list.detach()]  # , sig_a.detach(), sig_b.detach()]
                #############################################################################################################
                # statistic = [mu, sig]

                return x, domain, cluster_idx, cluster, statistic

            #############################################################################################################
            # mu_a, mu_b, sig_a, sig_b = statistic
            # mu_batch = []
            # sig_batch = []
            # factor1 = 1
            # factor2 = 1
            # eps_mu = torch.randn_like(mu)  # 2) epsilon_bc_dim
            # eps_sig = torch.randn_like(sig)
            #
            #
            # for i in range(B):
            #
            #
            #     uniform = torch.rand_like(mu_a).to(x.device)
            #     mu_c = (mu_a - mu_b * factor1) * (1 - uniform) + (mu_a + mu_b * factor1) * uniform
            #
            #
            #     #uniform = torch.rand_like(sig_a).to(x.device)
            #     sig_c = (sig_a - sig_b * factor2) * ((sig_a - sig_b * factor2) > 0) * (1 - uniform) + (
            #                 sig_a + sig_b * factor2) * uniform
            #
            #     mu_batch.append(mu_c.resize(1, C, 1, 1))
            #     sig_batch.append(sig_c.resize(1, C, 1, 1))
            #
            #
            # mu_new = torch.vstack(mu_batch)  #+ eps_mu  * mu_b
            # sig_new = torch.vstack(sig_batch) # + eps_sig  * sig_b
            # sig_new = sig_new * (sig_new>0)
            #############################################################################################################
            # Just mean
            else:

                mu_statistic, sig_statistic = statistic[0].to(mu.device), statistic[1].to(mu.device)
                #threshold = statistic[2]
                mu_batch = []
                sig_batch = []

                train_statistic = torch.squeeze(torch.cat((mu_statistic, sig_statistic), dim=1))

                # train_statistic_mu = torch.squeeze(mu_statistic)
                # train_statistic_sig = torch.squeeze(sig_statistic)

                mean_train_stat = train_statistic.mean(dim=[0], keepdim=True)
                # sig_train_stat = train_statistic_sig.mean(dim=[0], keepdim=True)

                norm_s = torch.norm(train_statistic - mean_train_stat, dim=(1))
                # real_mean = (train_statistic[0] + train_statistic[1] + train_statistic[2])/3
                max_distance = torch.mean(norm_s)

                max_distance = torch.sum(norm_s)

                # pdb.set_trace()
                #############################################################################################################
                # arg min
                for i in range(B):

                    test_statistic = torch.squeeze(torch.cat((mu[i], sig[i]), dim=0))

                    ra = np.random.choice(B, 1)
                    ra = ra[0]

                    vector = (torch.squeeze(torch.cat((mu[i], sig[i]), dim=0)) - mean_train_stat) * 10

                    # test_statistic = torch.squeeze(mu[i])

                    if val2 == True:
                        test_statistic = mean_train_stat + vector * torch.rand(1).to(mu.device)
                        x_normed = (x[i] - mu[i]) / sig[i]
                        mu_hat = test_statistic[0, :C]
                        sig_hat = test_statistic[0, C:]
                        x[i] = x_normed * sig_hat.resize(1, C, 1, 1) + mu_hat.resize(1, C, 1, 1)
                        mu[i] = mu_hat.resize(1, C, 1, 1)
                        sig[i] = sig_hat.resize(1, C, 1, 1)

                    distance_test = torch.norm(test_statistic - mean_train_stat)

                    distance_test_c = torch.norm(test_statistic - train_statistic, dim=(1))
                    distance_test_c = torch.sum(distance_test_c)

                    print(distance_test_c)
                    #pdb.set_trace()
                    if distance_test_c < max_distance * alpha:
                        # if distance_test_c[0] < norm_s[0] or distance_test_c[1] < norm_s[1] or distance_test_c[2] < norm_s[2]:

                        mu_batch.append(mu[i].resize(1, C, 1, 1))
                        sig_batch.append(sig[i].resize(1, C, 1, 1))

                    else:
                        # pdb.set_trace()
                        norm_t = torch.norm(train_statistic - test_statistic, dim=1)
                        min_idx = torch.argmin(norm_t)

                        # print(min_idx)
                        mu_c = mu_statistic[min_idx]
                        # mu_c = mean_train_stat

                        sig_c = sig_statistic[min_idx]
                        # sig_c = sig_train_stat

                        mu_batch.append(mu_c.resize(1, C, 1, 1))
                        sig_batch.append(sig_c.resize(1, C, 1, 1))

                mu_new = torch.vstack(mu_batch)  # + eps_mu  * mu_b
                sig_new = torch.vstack(sig_batch)  # + eps_sig  * sig_b
                sig_new = sig_new * (sig_new > 0)

                x_normed = (x - mu) / sig
                x_new = x_normed * sig_new + mu_new
                cluster = True

                return x_new, domain, cluster_idx, cluster, statistic

            #############################################################################################################
            # Mahalanobis distance based

            # else:
            #     mu_batch = []
            #     sig_batch = []
            #
            #     n_cluster = 3
            #     mean_samples, cov_per_domain = statistic[0], statistic[1]
            #     #mean_per_domain = statistic[2]
            #     #threshold =statistic[3]
            #     identity = torch.eye(cov_per_domain[0].size(0)).to(cov_per_domain[0].device)
            #
            #     #train_statistic = torch.vstack(mean_per_domain)
            #
            #
            #     #mean_train_stat = train_statistic.mean(dim=[0], keepdim=True)
            #
            #     #score_each_domain = []
            #     # for k in range(n_cluster):
            #     #     score_each_domain.append(torch.sqrt(self.mahalanobis_score(mean_per_domain[k], cov_per_domain[k], torch.squeeze(mean_train_stat))))
            #
            #
            #     #score = torch.sqrt(self.mahalanobis_score(mean_per_domain, cov_per_domain, torch.squeeze(mean_train_stat))
            #
            #
            #     #max_distance = torch.sum(torch.tensor(score_each_domain))
            #
            #     #
            #     #############################################################################################################
            #     # arg min
            #     for i in range(B):
            #
            #         score_each_domain_test = []
            #         test_statistic = torch.squeeze(torch.cat((mu[i],sig[i]),dim=0))
            #         #test_statistic = torch.squeeze(mu[i])
            #
            #         for k in range(n_cluster):
            #             score_each_domain_test.append(torch.sqrt(self.mahalanobis_score(mean_samples[k], cov_per_domain[k],
            #                                                             torch.squeeze(test_statistic))))
            #
            #         # score = torch.sqrt(self.mahalanobis_score(mean_samples, identity,
            #         #                                      torch.squeeze(test_statistic)))
            #
            #
            #
            #         distance_test_c = torch.min(torch.tensor(score_each_domain_test))
            #         min_idx = torch.argmin(torch.tensor(score_each_domain_test))
            #
            #
            #
            #         #print(min_idx)
            #
            #         if distance_test_c < 10:
            #         #if distance_test_c[0] < norm_s[0] or distance_test_c[1] < norm_s[1] or distance_test_c[2] < norm_s[2]:
            #
            #             mu_batch.append(mu[i].resize(1, C, 1, 1))
            #             sig_batch.append(sig[i].resize(1, C, 1, 1))
            #
            #         else:
            #             # norm_t = torch.norm(mean_samples - test_statistic, dim=1)
            #             # min_idx = torch.argmin(norm_t)
            #             #pdb.set_trace()
            #
            #             #print(min_idx)
            #             mu_c = mean_samples[min_idx][:C]
            #             #mu_c = mean_train_stat
            #
            #
            #             sig_c = mean_samples[min_idx][C:]
            #             #sig_c = sig_train_stat
            #
            #             mu_batch.append(mu_c.resize(1, C, 1, 1))
            #             sig_batch.append(sig_c.resize(1, C, 1, 1))
            #
            #     mu_new = torch.vstack(mu_batch)  # + eps_mu  * mu_b
            #     sig_new = torch.vstack(sig_batch)  # + eps_sig  * sig_b
            #     sig_new = sig_new * (sig_new > 0)
            #
            #     x_normed = (x - mu) / sig
            #     x_new = x_normed * sig_new + mu_new
            #     cluster = True
            #
            #     return x_new, domain, cluster_idx, cluster, statistic

            #     #############################################################################################################
            # DBSCAN
            # else:
            #
            #     mu_statistic, sig_statistic, clustering = statistic[0].to(mu.device), statistic[1].to(mu.device), statistic[2]
            #     mu_batch = []
            #     sig_batch = []
            #     #pdb.set_trace()
            #     train_statistic = torch.squeeze(torch.cat((mu_statistic, sig_statistic), dim=1))
            #
            #     #train_statistic_mu = torch.squeeze(mu_statistic)
            #     #train_statistic_sig = torch.squeeze(sig_statistic)
            #
            #     mean_train_stat = train_statistic.mean(dim=[0], keepdim=True)
            #     # sig_train_stat = train_statistic_sig.mean(dim=[0], keepdim=True)
            #
            #     norm_s = torch.norm( train_statistic - mean_train_stat, dim=(1))
            #     #real_mean = (train_statistic[0] + train_statistic[1] + train_statistic[2])/3
            #     max_distance = torch.mean(norm_s)
            #
            #     max_distance = torch.sum(norm_s)
            #
            #     #pdb.set_trace()
            #     #############################################################################################################
            #     # arg min
            #     for i in range(B):
            #
            #
            #         test_statistic = torch.squeeze(torch.cat((mu[i],sig[i]),dim=0))
            #         #test_statistic = torch.squeeze(mu[i])
            #
            #         test_statistic_numpy = (test_statistic.clone().cpu().numpy()).reshape(1,-1)
            #         label_test = clustering.fit_predict(test_statistic_numpy)
            #
            #
            #         print(label_test)
            #
            #
            #         if label_test == -1:
            #             norm_t = torch.norm(train_statistic - test_statistic, dim=1)
            #             min_idx = torch.argmin(norm_t)
            #
            #             #print(min_idx)
            #             mu_c = mu_statistic[min_idx]
            #             # mu_c = mean_train_stat
            #
            #             sig_c = sig_statistic[min_idx]
            #             # sig_c = sig_train_stat
            #
            #             mu_batch.append(mu_c.resize(1, C, 1, 1))
            #             sig_batch.append(sig_c.resize(1, C, 1, 1))
            #
            #
            #
            #         else:
            #             mu_batch.append(mu[i].resize(1, C, 1, 1))
            #             sig_batch.append(sig[i].resize(1, C, 1, 1))
            #
            #
            #
            #     mu_new = torch.vstack(mu_batch)  # + eps_mu  * mu_b
            #     sig_new = torch.vstack(sig_batch)  # + eps_sig  * sig_b
            #     sig_new = sig_new * (sig_new > 0)
            #
            #     x_normed = (x - mu) / sig
            #     x_new = x_normed * sig_new + mu_new
            #     cluster = True
            #
            #     return x_new, domain, cluster_idx, cluster, statistic
            #     #############################################################################################################
            # efdm based
            # else:
            #
            #     mu_statistic, sig_statistic = statistic[0].to(mu.device), statistic[1].to(mu.device)
            #     x_samples = statistic[2].to(mu.device)
            #
            #     N = x_samples.size(0)
            #
            #     x_new_batch = []
            #     x_view = x_samples.view(N, C, -1)
            #     value_x, index_x = torch.sort(x_view)  # sort inputs
            #     inverse_index = index_x.argsort(-1)
            #
            #     mu_batch = []
            #     sig_batch = []
            #
            #     train_statistic = torch.squeeze(torch.cat((mu_statistic, sig_statistic), dim=1))
            #
            #     #train_statistic_mu = torch.squeeze(mu_statistic)
            #     #train_statistic_sig = torch.squeeze(sig_statistic)
            #
            #     mean_train_stat = train_statistic.mean(dim=[0], keepdim=True)
            #     # sig_train_stat = train_statistic_sig.mean(dim=[0], keepdim=True)
            #
            #     norm_s = torch.norm( train_statistic - mean_train_stat, dim=(1))
            #     #real_mean = (train_statistic[0] + train_statistic[1] + train_statistic[2])/3
            #     max_distance = torch.mean(norm_s)
            #
            #     max_distance = torch.sum(norm_s)
            #
            #     #pdb.set_trace()
            #     #############################################################################################################
            #     # arg min
            #     for i in range(B):
            #
            #
            #         test_statistic = torch.squeeze(torch.cat((mu[i],sig[i]),dim=0))
            #         #test_statistic = torch.squeeze(mu[i])
            #         distance_test = torch.norm(test_statistic - mean_train_stat)
            #
            #         distance_test_c = torch.norm(test_statistic - train_statistic, dim=(1))
            #         distance_test_c = torch.sum(distance_test_c)
            #
            #
            #         #pdb.set_trace()
            #         if distance_test_c < 0 * 3:
            #         #if distance_test_c[0] < norm_s[0] or distance_test_c[1] < norm_s[1] or distance_test_c[2] < norm_s[2]:
            #
            #             mu_batch.append(mu[i].resize(1, C, 1, 1))
            #             sig_batch.append(sig[i].resize(1, C, 1, 1))
            #
            #         else:
            #            # pdb.set_trace()
            #             norm_t = torch.norm(train_statistic - test_statistic, dim=1)
            #             min_idx = torch.argmin(norm_t)
            #
            #
            #             print(min_idx)
            #             mu_c = mu_statistic[min_idx]
            #             #mu_c = mean_train_stat
            #
            #
            #             sig_c = sig_statistic[min_idx]
            #             #sig_c = sig_train_stat
            #
            #             mu_batch.append(mu_c.resize(1, C, 1, 1))
            #             sig_batch.append(sig_c.resize(1, C, 1, 1))
            #
            #     mu_new = torch.vstack(mu_batch)  # + eps_mu  * mu_b
            #     sig_new = torch.vstack(sig_batch)  # + eps_sig  * sig_b
            #     sig_new = sig_new * (sig_new > 0)
            #
            #     x_normed = (x - mu) / sig
            #     x_new = x_normed * sig_new + mu_new
            #     cluster = True
            #
            #     return x_new, domain, cluster_idx, cluster, statistic
            ########################################################################################################
            # ConvexHull
            # else:
            #     from scipy.spatial import ConvexHull
            #
            #     mu_statistic, sig_statistic = statistic[0].to(mu.device), statistic[1].to(mu.device)
            #     samples = statistic[2]
            #     pca = statistic[3]
            #
            #
            #
            #     mu_batch = []
            #     sig_batch = []
            #
            #     train_statistic = torch.squeeze(torch.cat((mu_statistic, sig_statistic), dim=1))
            #
            #     #train_statistic_mu = torch.squeeze(mu_statistic)
            #     #train_statistic_sig = torch.squeeze(sig_statistic)
            #
            #     mean_train_stat = train_statistic.mean(dim=[0], keepdim=True)
            #     # sig_train_stat = train_statistic_sig.mean(dim=[0], keepdim=True)
            #
            #     norm_s = torch.norm( train_statistic - mean_train_stat, dim=(1))
            #     #real_mean = (train_statistic[0] + train_statistic[1] + train_statistic[2])/3
            #     max_distance = torch.mean(norm_s)
            #
            #     max_distance = torch.sum(norm_s)
            #
            #     #pdb.set_trace()
            #     #############################################################################################################
            #     # arg min
            #     for i in range(B):
            #
            #
            #         test_statistic = torch.squeeze(torch.cat((mu[i],sig[i]),dim=0))
            #
            #         test_sample = test_statistic.cpu().numpy()
            #         #test_sample = test_sample.reshape(1, -1)
            #         #test_sample = pca.transform(test_sample)
            #         #pdb.set_trace()
            #         ishull = in_hull(samples, test_sample ) #samples 2D, test_sample 1D
            #
            #
            #
            #         if ishull:
            #         #if distance_test_c[0] < norm_s[0] or distance_test_c[1] < norm_s[1] or distance_test_c[2] < norm_s[2]:
            #
            #             mu_batch.append(mu[i].resize(1, C, 1, 1))
            #             sig_batch.append(sig[i].resize(1, C, 1, 1))
            #
            #         else:
            #            # pdb.set_trace()
            #             norm_t = torch.norm(train_statistic - test_statistic, dim=1)
            #             min_idx = torch.argmin(norm_t)
            #
            #
            #             print(min_idx)
            #             mu_c = mu_statistic[min_idx]
            #             #mu_c = mean_train_stat
            #
            #
            #             sig_c = sig_statistic[min_idx]
            #             #sig_c = sig_train_stat
            #
            #             mu_batch.append(mu_c.resize(1, C, 1, 1))
            #             sig_batch.append(sig_c.resize(1, C, 1, 1))
            #
            #     mu_new = torch.vstack(mu_batch)  # + eps_mu  * mu_b
            #     sig_new = torch.vstack(sig_batch)  # + eps_sig  * sig_b
            #     sig_new = sig_new * (sig_new > 0)
            #
            #     x_normed = (x - mu) / sig
            #     x_new = x_normed * sig_new + mu_new
            #     cluster = True
            #
            #     return x_new, domain, cluster_idx, cluster, statistic






