import random
from contextlib import contextmanager
import torch
import torch.nn as nn
import numpy as np
import pdb
def deactivate_efdmix(m):
    if type(m) == EFDMix:
        m.set_activation_status(False)


def activate_efdmix(m):
    if type(m) == EFDMix:
        m.set_activation_status(True)


def random_efdmix(m):
    if type(m) == EFDMix:
        m.update_mix_method('random')


def crossdomain_efdmix(m):
    if type(m) == EFDMix:
        m.update_mix_method('crossdomain')


@contextmanager
def run_without_efdmix(model):
    # Assume EFDMix was initially activated
    try:
        model.apply(deactivate_efdmix)
        yield
    finally:
        model.apply(activate_efdmix)


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

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

    try:
        model.apply(activate_efdmix)
        yield
    finally:
        model.apply(deactivate_efdmix)


class EFDMix(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=None, p=0.5):
        if not self.training or not self._activated:
            return x

        if random.random() > p:
            return x

        B,C,W,H = x.size(0), x.size(1), x.size(2),  x.size(3)
        x_view = x.view(B,C, -1)
        ## sort input vectors.
        value_x, index_x = torch.sort(x_view)
        lmda = self.beta.sample((B, 1, 1))
        lmda = lmda.to(x.device)

        if torch.is_tensor(domain):
            domain = domain.cpu().numpy()
        if self.mix == 'random':
            # random shuffle
            # perm = torch.randperm(B)

            original_index = []
            new_index = []
            num_domain = 6 #pacs 3

            ########################################################################################################
            # label matching for our algorithm

            for i in range(num_domain):
                original_index.append(set(np.where(domain == i)[0]))

            sample_union = set()

            for idx in range(domain.size):
                union = set()
                for domain_label in range(num_domain):
                    if domain[idx] != domain_label:
                        union = union | (set(original_index[domain_label]) - sample_union)
                try:
                    sample = random.sample(union, 1)[0]
                except:
                    # union = set(original_index[domain_label]) - sample_union

                    union = set(original_index[domain[idx]]) - sample_union

                    sample = random.sample(union, 1)[0]


                new_index.append(sample)
                sample_union.add(sample)
            perm = new_index
            # pdb.set_trace()
            ########################################################################################################

        elif self.mix == 'crossdomain':
            # split into two halves and swap the order
            perm = torch.arange(B - 1, -1, -1) # inverse index
            perm_b, perm_a = perm.chunk(2)
            perm_b = perm_b[torch.randperm(B // 2)]
            perm_a = perm_a[torch.randperm(B // 2)]
            perm = torch.cat([perm_b, perm_a], 0)

        else:
            raise NotImplementedError

        inverse_index = index_x.argsort(-1)
        x_view_copy = value_x[perm].gather(-1, inverse_index) * (1-lmda)
        new_x = x_view + (x_view_copy - x_view.detach() * (1-lmda))
        return new_x.view(B, C, W, H)


